[FancyZones] Send message from VirtualDesktopUpdates thread to FZ thread when update happens (#2568)

* Move part of the virtual desktops related logic from FancyZones to VirtualDesktopUtils.

* Post WM message from vritual desktop tracker thread to FZ thread.

* Minor improvements in RegisterVirtualDesktopUpdates method.

* Close registry key after HandleVirtualDesktopUpdates thread finishes execution.

* Remove comment explaining workaround to VirtualDesktopUtils namespace.

* Move HandleVirtualDesktopUpdates to VirtualDesktopUtils namespace. Resolve PR comments.

* Fix typos in window messages description.

* Remove lock from OnKeyDown method to avoid deadlock.

src/modules/fancyzones/lib/FancyZones.cpp

@@ -190,8 +190,7 @@ private:
     void CycleActiveZoneSet(DWORD vkCode) noexcept;
     bool OnSnapHotkey(DWORD vkCode) noexcept;
     
-    void HandleVirtualDesktopUpdates(HANDLE fancyZonesDestroyedEvent) noexcept;
+    void RegisterVirtualDesktopUpdates(std::vector<GUID>& ids) noexcept;
-    void RegisterVirtualDesktopUpdates(std::unordered_set<GUID>& currentVirtualDesktopIds) noexcept;
     void RegisterNewWorkArea(GUID virtualDesktopId, HMONITOR monitor) noexcept;
     bool IsNewWorkArea(GUID virtualDesktopId, HMONITOR monitor) noexcept;
 
@@ -202,8 +201,6 @@ private:
 
     const HINSTANCE m_hinstance{};
 
-    HKEY m_virtualDesktopsRegKey{ nullptr };
-
     mutable std::shared_mutex m_lock;
     HWND m_window{};
     WindowMoveHandler m_windowMoveHandler;
@@ -218,9 +215,10 @@ private:
     OnThreadExecutor m_dpiUnawareThread;
     OnThreadExecutor m_virtualDesktopTrackerThread;
 
-    static UINT WM_PRIV_VDCHANGED; // Message to get back on to the UI thread when virtual desktop changes
+    static UINT WM_PRIV_VD_INIT; // Message to get back to the UI thread when FancyZones is initialized
-    static UINT WM_PRIV_VDINIT; // Message to get back to the UI thread when FancyZones are initialized
+    static UINT WM_PRIV_VD_SWITCH; // Message to get back to the UI thread when virtual desktop switch occurs
-    static UINT WM_PRIV_EDITOR; // Message to get back on to the UI thread when the editor exits
+    static UINT WM_PRIV_VD_UPDATE; // Message to get back to the UI thread on virtual desktops update (creation/deletion)
+    static UINT WM_PRIV_EDITOR; // Message to get back to the UI thread when the editor exits
 
     // Did we terminate the editor or was it closed cleanly?
     enum class EditorExitKind : byte
@@ -230,8 +228,9 @@ private:
     };
 };
 
-UINT FancyZones::WM_PRIV_VDCHANGED = RegisterWindowMessage(L"{128c2cb0-6bdf-493e-abbe-f8705e04aa95}");
+UINT FancyZones::WM_PRIV_VD_INIT = RegisterWindowMessage(L"{469818a8-00fa-4069-b867-a1da484fcd9a}");
-UINT FancyZones::WM_PRIV_VDINIT = RegisterWindowMessage(L"{469818a8-00fa-4069-b867-a1da484fcd9a}");
+UINT FancyZones::WM_PRIV_VD_SWITCH = RegisterWindowMessage(L"{128c2cb0-6bdf-493e-abbe-f8705e04aa95}");
+UINT FancyZones::WM_PRIV_VD_UPDATE = RegisterWindowMessage(L"{b8b72b46-f42f-4c26-9e20-29336cf2f22e}");
 UINT FancyZones::WM_PRIV_EDITOR = RegisterWindowMessage(L"{87543824-7080-4e91-9d9c-0404642fc7b6}");
 
 // IFancyZones
@@ -263,12 +262,9 @@ FancyZones::Run() noexcept
                       } })
         .wait();
 
-    if (RegOpenKeyEx(HKEY_CURRENT_USER, L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\VirtualDesktops", 0, KEY_ALL_ACCESS, &m_virtualDesktopsRegKey) == ERROR_SUCCESS)
-    {
     m_terminateVirtualDesktopTrackerEvent.reset(CreateEvent(nullptr, FALSE, FALSE, nullptr));
-        m_virtualDesktopTrackerThread.submit(
+    m_virtualDesktopTrackerThread.submit(OnThreadExecutor::task_t{ [&] {
-            OnThreadExecutor::task_t{ std::bind(&FancyZones::HandleVirtualDesktopUpdates, this, m_terminateVirtualDesktopTrackerEvent.get()) });
+        VirtualDesktopUtils::HandleVirtualDesktopUpdates(m_window, WM_PRIV_VD_UPDATE, m_terminateVirtualDesktopTrackerEvent.get()); } });
-    }
 }
 
 // IFancyZones
@@ -287,11 +283,6 @@ FancyZones::Destroy() noexcept
     {
         SetEvent(m_terminateVirtualDesktopTrackerEvent.get());
     }
-    if (m_virtualDesktopsRegKey)
-    {
-        RegCloseKey(m_virtualDesktopsRegKey);
-        m_virtualDesktopsRegKey = nullptr;
-    }
 }
 
 // IFancyZonesCallback
@@ -301,14 +292,14 @@ FancyZones::VirtualDesktopChanged() noexcept
     // VirtualDesktopChanged is called from another thread but results in new windows being created.
     // Jump over to the UI thread to handle it.
     std::shared_lock readLock(m_lock);
-    PostMessage(m_window, WM_PRIV_VDCHANGED, 0, 0);
+    PostMessage(m_window, WM_PRIV_VD_SWITCH, 0, 0);
 }
 
 // IFancyZonesCallback
 IFACEMETHODIMP_(void)
 FancyZones::VirtualDesktopInitialize() noexcept
 {
-    PostMessage(m_window, WM_PRIV_VDINIT, 0, 0);
+    PostMessage(m_window, WM_PRIV_VD_INIT, 0, 0);
 }
 
 // IFancyZonesCallback
@@ -390,7 +381,6 @@ FancyZones::OnKeyDown(PKBDLLHOOKSTRUCT info) noexcept
     //    return false;
     //}
 
-    std::shared_lock readLock(m_lock);
     if (m_windowMoveHandler.IsDragEnabled() && shift)
     {
         return true;
@@ -567,13 +557,21 @@ LRESULT FancyZones::WndProc(HWND window, UINT message, WPARAM wparam, LPARAM lpa
 
     default:
     {
-        if (message == WM_PRIV_VDCHANGED)
+        if (message == WM_PRIV_VD_INIT)
+        {
+            OnDisplayChange(DisplayChangeType::Initialization);
+        }
+        else if (message == WM_PRIV_VD_SWITCH)
         {
             OnDisplayChange(DisplayChangeType::VirtualDesktop);
         }
-        else if (message == WM_PRIV_VDINIT)
+        else if (message == WM_PRIV_VD_UPDATE)
         {
-            OnDisplayChange(DisplayChangeType::Initialization);
+            std::vector<GUID> ids{};
+            if (VirtualDesktopUtils::GetVirtualDekstopIds(ids))
+            {
+                RegisterVirtualDesktopUpdates(ids);
+            }
         }
         else if (message == WM_PRIV_EDITOR)
         {
@@ -604,23 +602,11 @@ void FancyZones::OnDisplayChange(DisplayChangeType changeType) noexcept
     if (changeType == DisplayChangeType::VirtualDesktop ||
         changeType == DisplayChangeType::Initialization)
     {
-        // Explorer persists this value to the registry on a per session basis but only after
-        // the first virtual desktop switch happens. If the user hasn't switched virtual desktops in this session
-        // then this value will be empty. This means loading the first virtual desktop's configuration can be
-        // funky the first time we load up at boot since the user will not have switched virtual desktops yet.
         GUID currentVirtualDesktopId{};
         if (VirtualDesktopUtils::GetCurrentVirtualDesktopId(&currentVirtualDesktopId))
         {
             std::unique_lock writeLock(m_lock);
             m_currentVirtualDesktopId = currentVirtualDesktopId;
-        }
-        else
-        {
-            std::vector<GUID> ids{};
-            if (VirtualDesktopUtils::GetVirtualDekstopIds(m_virtualDesktopsRegKey, ids) && !ids.empty())
-            {
-                std::unique_lock writeLock(m_lock);
-                m_currentVirtualDesktopId = ids[0];
             wil::unique_cotaskmem_string id;
             if (changeType == DisplayChangeType::Initialization &&
                 SUCCEEDED_LOG(StringFromCLSID(m_currentVirtualDesktopId, &id)))
@@ -629,7 +615,6 @@ void FancyZones::OnDisplayChange(DisplayChangeType changeType) noexcept
             }
         }
     }
-    }
 
     UpdateZoneWindows();
 
@@ -827,38 +812,15 @@ bool FancyZones::OnSnapHotkey(DWORD vkCode) noexcept
     return false;
 }
 
-void FancyZones::HandleVirtualDesktopUpdates(HANDLE fancyZonesDestroyedEvent) noexcept
+void FancyZones::RegisterVirtualDesktopUpdates(std::vector<GUID>& ids) noexcept
-{
-    HANDLE regKeyEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
-    HANDLE events[2] = { regKeyEvent, fancyZonesDestroyedEvent };
-    while (1)
-    {
-        if (RegNotifyChangeKeyValue(m_virtualDesktopsRegKey, TRUE, REG_NOTIFY_CHANGE_LAST_SET, regKeyEvent, TRUE) != ERROR_SUCCESS)
-        {
-            return;
-        }
-        if (WaitForMultipleObjects(2, events, FALSE, INFINITE) != (WAIT_OBJECT_0 + 0))
-        {
-            // if fancyZonesDestroyedEvent is signalized or WaitForMultipleObjects failed, terminate thread execution
-            return;
-        }
-        std::vector<GUID> ids{};
-        if (VirtualDesktopUtils::GetVirtualDekstopIds(m_virtualDesktopsRegKey, ids))
-        {
-            std::unordered_set<GUID> idSet(std::begin(ids), std::end(ids));
-            RegisterVirtualDesktopUpdates(idSet);
-        }
-    }
-}
-
-void FancyZones::RegisterVirtualDesktopUpdates(std::unordered_set<GUID>& currentVirtualDesktopIds) noexcept
 {
+    std::unordered_set<GUID> activeVirtualDesktops(std::begin(ids), std::end(ids));
     std::unique_lock writeLock(m_lock);
     bool modified{ false };
-    for (auto it = begin(m_processedWorkAreas); it != end(m_processedWorkAreas);)
+    for (auto it = std::begin(m_processedWorkAreas); it != std::end(m_processedWorkAreas);)
     {
-        auto iter = currentVirtualDesktopIds.find(it->first);
+        auto iter = activeVirtualDesktops.find(it->first);
-        if (iter == currentVirtualDesktopIds.end())
+        if (iter == activeVirtualDesktops.end())
         {
             // if we couldn't find the GUID in currentVirtualDesktopIds, we must remove it from both m_processedWorkAreas and deviceInfoMap
             wil::unique_cotaskmem_string virtualDesktopId;
@@ -870,7 +832,7 @@ void FancyZones::RegisterVirtualDesktopUpdates(std::unordered_set<GUID>& current
         }
         else
         {
-            currentVirtualDesktopIds.erase(it->first); // virtual desktop already in map, skip it
+            activeVirtualDesktops.erase(it->first); // virtual desktop already in map, skip it
             ++it;
         }
     }
@@ -879,7 +841,7 @@ void FancyZones::RegisterVirtualDesktopUpdates(std::unordered_set<GUID>& current
         JSONHelpers::FancyZonesDataInstance().SaveFancyZonesData();
     }
     // register new virtual desktops, if any
-    for (const auto& id : currentVirtualDesktopIds)
+    for (const auto& id : activeVirtualDesktops)
     {
         m_processedWorkAreas[id] = std::vector<HMONITOR>();
     }

src/modules/fancyzones/lib/VirtualDesktopUtils.cpp

@@ -9,6 +9,7 @@ namespace VirtualDesktopUtils
 
     const wchar_t RegCurrentVirtualDesktop[] = L"CurrentVirtualDesktop";
     const wchar_t RegVirtualDesktopIds[] = L"VirtualDesktopIDs";
+    const wchar_t RegKeyVirtualDesktops[] = L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\VirtualDesktops";
 
     IServiceProvider* GetServiceProvider()
     {
@@ -50,7 +51,7 @@ namespace VirtualDesktopUtils
         return SUCCEEDED(CLSIDFromString(virtualDesktopId.c_str(), desktopId));
     }
 
-    bool GetCurrentVirtualDesktopId(GUID* desktopId)
+    bool GetDesktopIdFromCurrentSession(GUID* desktopId)
     {
         DWORD sessionId;
         ProcessIdToSessionId(GetCurrentProcessId(), &sessionId);
@@ -77,8 +78,30 @@ namespace VirtualDesktopUtils
         return false;
     }
 
+    bool GetCurrentVirtualDesktopId(GUID* desktopId)
+    {
+        if (!GetDesktopIdFromCurrentSession(desktopId))
+        {
+            // Explorer persists current virtual desktop identifier to registry on a per session basis,
+            // but only after first virtual desktop switch happens. If the user hasn't switched virtual
+            // 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())
+            {
+                return false;
+            }
+            *desktopId = ids[0];
+        }
+        return true;
+    }
+
     bool GetVirtualDekstopIds(HKEY hKey, std::vector<GUID>& ids)
     {
+        if (!hKey)
+        {
+            return false;
+        }
         DWORD bufferCapacity;
         // request regkey binary buffer capacity only
         if (RegQueryValueExW(hKey, RegVirtualDesktopIds, 0, nullptr, nullptr, &bufferCapacity) != ERROR_SUCCESS)
@@ -102,4 +125,49 @@ namespace VirtualDesktopUtils
         ids = std::move(temp);
         return true;
     }
+
+    bool GetVirtualDekstopIds(std::vector<GUID>& ids)
+    {
+        return GetVirtualDekstopIds(GetVirtualDesktopsRegKey(), ids);
+    }
+
+    HKEY OpenVirtualDesktopsRegKey()
+    {
+        HKEY hKey{ nullptr };
+        if (RegOpenKeyEx(HKEY_CURRENT_USER, RegKeyVirtualDesktops, 0, KEY_ALL_ACCESS, &hKey) == ERROR_SUCCESS)
+        {
+            return hKey;
+        }
+        return nullptr;
+    }
+
+    HKEY GetVirtualDesktopsRegKey()
+    {
+        static wil::unique_hkey virtualDesktopsKey{ OpenVirtualDesktopsRegKey() };
+        return virtualDesktopsKey.get();
+    }
+
+    void HandleVirtualDesktopUpdates(HWND window, UINT message, HANDLE terminateEvent)
+    {
+        HKEY virtualDesktopsRegKey = GetVirtualDesktopsRegKey();
+        if (!virtualDesktopsRegKey)
+        {
+            return;
+        }
+        HANDLE regKeyEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
+        HANDLE events[2] = { regKeyEvent, terminateEvent };
+        while (1)
+        {
+            if (RegNotifyChangeKeyValue(virtualDesktopsRegKey, TRUE, REG_NOTIFY_CHANGE_LAST_SET, regKeyEvent, TRUE) != ERROR_SUCCESS)
+            {
+                return;
+            }
+            if (WaitForMultipleObjects(2, events, FALSE, INFINITE) != (WAIT_OBJECT_0 + 0))
+            {
+                // if terminateEvent is signalized or WaitForMultipleObjects failed, terminate thread execution
+                return;
+            }
+            PostMessage(window, message, 0, 0);
+        }
+    }
 }

src/modules/fancyzones/lib/VirtualDesktopUtils.h

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