[FancyZones] Clean up resources taken by closed virtual desktops on module startup (#2616)

* Clean up resources taken by closed virtual desktops on FancyZones module startup.

* Overload GetVirtualDesktopIds to return wide string.

* Refactor OnDisplayChange method.

src/modules/fancyzones/lib/FancyZones.cpp

@@ -568,7 +568,7 @@ LRESULT FancyZones::WndProc(HWND window, UINT message, WPARAM wparam, LPARAM lpa
         else if (message == WM_PRIV_VD_UPDATE)
         {
             std::vector<GUID> ids{};
-            if (VirtualDesktopUtils::GetVirtualDekstopIds(ids))
+            if (VirtualDesktopUtils::GetVirtualDesktopIds(ids))
             {
                 RegisterVirtualDesktopUpdates(ids);
             }
@@ -597,6 +597,7 @@ LRESULT FancyZones::WndProc(HWND window, UINT message, WPARAM wparam, LPARAM lpa
     return 0;
 }
 
+
 void FancyZones::OnDisplayChange(DisplayChangeType changeType) noexcept
 {
     if (changeType == DisplayChangeType::VirtualDesktop ||
@@ -605,13 +606,15 @@ void FancyZones::OnDisplayChange(DisplayChangeType changeType) noexcept
         GUID currentVirtualDesktopId{};
         if (VirtualDesktopUtils::GetCurrentVirtualDesktopId(&currentVirtualDesktopId))
         {
-            std::unique_lock writeLock(m_lock);
             m_currentVirtualDesktopId = currentVirtualDesktopId;
-            wil::unique_cotaskmem_string id;
+        }
-            if (changeType == DisplayChangeType::Initialization &&
+        if (changeType == DisplayChangeType::Initialization)
-                SUCCEEDED_LOG(StringFromCLSID(m_currentVirtualDesktopId, &id)))
+        {
+            std::vector<std::wstring> ids{};
+            if (VirtualDesktopUtils::GetVirtualDesktopIds(ids) && !ids.empty())
             {
-                JSONHelpers::FancyZonesDataInstance().UpdatePrimaryDesktopData(id.get());
+                JSONHelpers::FancyZonesDataInstance().UpdatePrimaryDesktopData(ids[0]);
+                JSONHelpers::FancyZonesDataInstance().RemoveDeletedDesktops(ids);
             }
         }
     }

src/modules/fancyzones/lib/JsonHelpers.cpp

@@ -10,6 +10,7 @@
 #include <fstream>
 #include <regex>
 #include <sstream>
+#include <unordered_set>
 
 namespace
 {
@@ -348,6 +349,25 @@ namespace JSONHelpers
         SaveFancyZonesData();
     }
 
+    void FancyZonesData::RemoveDeletedDesktops(const std::vector<std::wstring>& activeDesktops)
+    {
+        std::unordered_set<std::wstring> active(std::begin(activeDesktops), std::end(activeDesktops));
+        std::scoped_lock lock{ dataLock };
+        for (auto it = std::begin(deviceInfoMap); it != std::end(deviceInfoMap);)
+        {
+            auto foundId = active.find(ExtractVirtualDesktopId(it->first));
+            if (foundId == std::end(active))
+            {
+                it = deviceInfoMap.erase(it);
+            }
+            else
+            {
+                ++it;
+            }
+        }
+        SaveFancyZonesData();
+    }
+
     int FancyZonesData::GetAppLastZoneIndex(HWND window, const std::wstring_view& deviceId, const std::wstring_view& zoneSetId) const
     {
         std::scoped_lock lock{ dataLock };

src/modules/fancyzones/lib/JsonHelpers.h

@@ -234,6 +234,7 @@ namespace JSONHelpers
         bool RemoveDevicesByVirtualDesktopId(const std::wstring& virtualDesktopId);
         void CloneDeviceInfo(const std::wstring& source, const std::wstring& destination);
         void UpdatePrimaryDesktopData(const std::wstring& desktopId);
+        void RemoveDeletedDesktops(const std::vector<std::wstring>& activeDesktops);
 
         int GetAppLastZoneIndex(HWND window, const std::wstring_view& deviceId, const std::wstring_view& zoneSetId) const;
         bool RemoveAppLastZone(HWND window, const std::wstring_view& deviceId, const std::wstring_view& zoneSetId);

src/modules/fancyzones/lib/VirtualDesktopUtils.cpp

@@ -87,7 +87,7 @@ namespace VirtualDesktopUtils
             // desktops (only primary desktop) in this session value in registry will be empty.
             // If this value is empty take first element from array of virtual desktops (not kept per session).
             std::vector<GUID> ids{};
-            if (!GetVirtualDekstopIds(ids) || ids.empty())
+            if (!GetVirtualDesktopIds(ids) || ids.empty())
             {
                 return false;
             }
@@ -96,7 +96,7 @@ namespace VirtualDesktopUtils
         return true;
     }
 
-    bool GetVirtualDekstopIds(HKEY hKey, std::vector<GUID>& ids)
+    bool GetVirtualDesktopIds(HKEY hKey, std::vector<GUID>& ids)
     {
         if (!hKey)
         {
@@ -126,9 +126,27 @@ namespace VirtualDesktopUtils
         return true;
     }
 
-    bool GetVirtualDekstopIds(std::vector<GUID>& ids)
+    bool GetVirtualDesktopIds(std::vector<GUID>& ids)
     {
-        return GetVirtualDekstopIds(GetVirtualDesktopsRegKey(), ids);
+        return GetVirtualDesktopIds(GetVirtualDesktopsRegKey(), ids);
+    }
+
+    bool GetVirtualDesktopIds(std::vector<std::wstring>& ids)
+    {
+        std::vector<GUID> guids{};
+        if (GetVirtualDesktopIds(guids))
+        {
+            for (auto& guid : guids)
+            {
+                wil::unique_cotaskmem_string guidString;
+                if (SUCCEEDED(StringFromCLSID(guid, &guidString)))
+                {
+                    ids.push_back(guidString.get());
+                }
+            }
+            return true;
+        }
+        return false;
     }
 
     HKEY OpenVirtualDesktopsRegKey()

src/modules/fancyzones/lib/VirtualDesktopUtils.h

@@ -7,7 +7,8 @@ namespace VirtualDesktopUtils
     bool GetWindowDesktopId(HWND topLevelWindow, GUID* desktopId);
     bool GetZoneWindowDesktopId(IZoneWindow* zoneWindow, GUID* desktopId);
     bool GetCurrentVirtualDesktopId(GUID* desktopId);
-    bool GetVirtualDekstopIds(std::vector<GUID>& ids);
+    bool GetVirtualDesktopIds(std::vector<GUID>& ids);
+    bool GetVirtualDesktopIds(std::vector<std::wstring>& ids);
     HKEY GetVirtualDesktopsRegKey();
     void HandleVirtualDesktopUpdates(HWND window, UINT message, HANDLE terminateEvent);
 }