Fix UIScheduler deadlock on shutdown (#12738)

vmoroz · web-flow · commit e2bf3e27f1ad · 2024-02-15T14:35:26.000-08:00
* Fix UIScheduler deadlock on shutdown

* Change files
diff --git a/change/react-native-windows-3f32eace-b18b-47b9-bc5d-0a51b9daf1c8.json b/change/react-native-windows-3f32eace-b18b-47b9-bc5d-0a51b9daf1c8.json
@@ -0,0 +1,7 @@
+{
+  "type": "prerelease",
+  "comment": "Fix UIScheduler deadlock on shutdown",
+  "packageName": "react-native-windows",
+  "email": "vmorozov@microsoft.com",
+  "dependentChangeType": "patch"
+}
diff --git a/vnext/Mso.UnitTests/dispatchQueue/dispatchQueueTest.cpp b/vnext/Mso.UnitTests/dispatchQueue/dispatchQueueTest.cpp
@@ -5,7 +5,19 @@
 #include <atomic>
 #include <condition_variable>
 #include <mutex>
+#include "eventWaitHandle/eventWaitHandle.h"
 #include "motifCpp/testCheck.h"
+#include "winrt/Windows.Foundation.h"
+
+using namespace winrt;
+using namespace Windows::Foundation;
+#ifndef USE_WINUI3
+#include "winrt/Windows.System.h"
+using namespace Windows::System;
+#else
+#include "winrt/Microsoft.UI.Dispatching.h"
+using namespace Microsoft::UI::Dispatching;
+#endif
 
 namespace Mso {
 extern void Test_ThreadPoolSchedulerWin_EnableThreadPoolWorkTracking(bool enable) noexcept;
@@ -15,55 +27,111 @@ extern void Test_ThreadPoolSchedulerWin_WaitForThreadPoolWorkCompletion() noexce
 namespace DispatchQueueTests {
 
 TEST_CLASS (DispatchQueueTest) {
-  TEST_METHOD(DispatchQueue_Shutdown) {
+  TEST_METHOD(DispatchQueue_SerialQueue_Shutdown) {
+    WithSerialQueue(TestShutdown);
+  }
+
+  TEST_METHOD(DispatchQueue_UIQueue_Shutdown) {
+    WithUIQueue(TestShutdown);
+  }
+
+  TEST_METHOD(DispatchQueue_UIQueue_ShutdownFromTask) {
+    WithUIQueue(TestShutdownFromTask);
+  }
+
+  TEST_METHOD(DispatchQueue_SerialQueue_ShutdownFromTask) {
+    WithSerialQueue(TestShutdownFromTask);
+  }
+
+  TEST_METHOD(DispatchQueue_UIQueue_ShutdownFromPreviousTask) {
+    WithUIQueue(TestShutdownFromPreviousTask);
+  }
+
+  TEST_METHOD(DispatchQueue_SerialQueue_ShutdownFromPreviousTask) {
+    WithSerialQueue(TestShutdownFromPreviousTask);
+  }
+
+ private:
+  static void TestShutdown(Mso::DispatchQueue queue, Mso::VoidFunctor drainQueue) {
     // Check that there is no dead lock if queue is released outside of task.
-    std::mutex mutex;
-    std::atomic<bool> taskStarted{false};
-    std::condition_variable whenTaskStarted;
     int callCount = 0;
-    {
-      Mso::DispatchQueue queue = Mso::DispatchQueue::MakeSerialQueue();
-      queue.Post([&]() {
-        ++callCount;
-        taskStarted = true;
-        whenTaskStarted.notify_all();
-      });
-      std::unique_lock lock(mutex);
-      whenTaskStarted.wait(lock, [&]() { return taskStarted.load(); });
-    }
+    queue.Post([&]() { ++callCount; });
+    drainQueue();
+
+    queue = nullptr;
+    TestCheck(callCount == 1);
+  };
+
+  static void TestShutdownFromTask(Mso::DispatchQueue queue, Mso::VoidFunctor drainQueue) {
+    // Check that there is no dead lock if queue is released form a task.
+    int callCount = 0;
+    Mso::ManualResetEvent queueReleased;
+    queue.Post([&]() {
+      ++callCount;
+      // Wait until the queue is released outside of task.
+      // So, the last release of the queue will be done from the task.
+      queueReleased.Wait();
+    });
+
+    queue = nullptr;
+    queueReleased.Set();
+
+    drainQueue();
     TestCheck(callCount == 1);
   }
 
-  TEST_METHOD(DispatchQueue_ShutdownFromTask) {
+  static void TestShutdownFromPreviousTask(Mso::DispatchQueue queue, Mso::VoidFunctor drainQueue) {
     // Check that there is no dead lock if queue is released form a task.
-    // Use shared_ptr for mutex because the test completes before the DispatchQueue task.
-    std::shared_ptr<std::mutex> mutex = std::make_shared<std::mutex>();
-    std::atomic<bool> taskStarted{false};
-    std::condition_variable whenTaskStarted;
-    std::atomic<bool> queueReleased{false};
-    std::condition_variable whenQueueReleased;
     int callCount = 0;
+    Mso::ManualResetEvent taskStarted;
+    Mso::ManualResetEvent queueReleased;
+    queue.Post([&]() {
+      taskStarted.Set();
+      ++callCount;
+      // Wait until the queue is released outside of task.
+      // So, the last release of the queue will be done from the task.
+      queueReleased.Wait();
+    });
+
+    queue.Post([&]() { ++callCount; });
+
+    taskStarted.Wait();
+    queue = nullptr;
+    queueReleased.Set();
+
+    drainQueue();
+    TestCheck(callCount == 2);
+  }
+
+  static void WithUIQueue(Mso::Functor<void(Mso::DispatchQueue queue, Mso::VoidFunctor drainQueue)> test) {
+    DispatcherQueueController queueController = DispatcherQueueController::CreateOnDedicatedThread();
+
+    Mso::VoidFunctor drainQueue = [queueController]() noexcept {
+      Mso::ManualResetEvent queueDrained;
+      queueController.DispatcherQueue().TryEnqueue([queueDrained]() noexcept { queueDrained.Set(); });
+      queueDrained.Wait();
+    };
+
+    Mso::DispatchQueue queue;
+    queueController.DispatcherQueue().TryEnqueue(
+        [&queue]() noexcept { queue = Mso::DispatchQueue::GetCurrentUIThreadQueue(); });
+    drainQueue();
+
+    test(std::move(queue), std::move(drainQueue));
+
+    queueController.ShutdownQueueAsync().get();
+  }
+
+  static void WithSerialQueue(Mso::Functor<void(Mso::DispatchQueue queue, Mso::VoidFunctor drainQueue)> test) {
     Mso::Test_ThreadPoolSchedulerWin_EnableThreadPoolWorkTracking(true);
-    {
-      Mso::DispatchQueue queue = Mso::DispatchQueue::MakeSerialQueue();
-      queue.Post([&, mutex]() {
-        ++callCount;
-        std::unique_lock lock(*mutex);
-        taskStarted = true;
-        whenTaskStarted.notify_all();
-        // Wait until the queue is released outside of task.
-        whenQueueReleased.wait(lock, [&]() { return queueReleased.load(); });
-      });
-      std::unique_lock lock(*mutex);
-      whenTaskStarted.wait(lock, [&]() { return taskStarted.load(); });
-    }
-    {
-      std::unique_lock lock(*mutex);
-      queueReleased = true;
-      whenQueueReleased.notify_all();
-    }
-    TestCheck(callCount == 1);
-    Mso::Test_ThreadPoolSchedulerWin_WaitForThreadPoolWorkCompletion();
+
+    Mso::VoidFunctor drainQueue = []() noexcept { Mso::Test_ThreadPoolSchedulerWin_WaitForThreadPoolWorkCompletion(); };
+
+    Mso::DispatchQueue queue = Mso::DispatchQueue::MakeSerialQueue();
+
+    test(std::move(queue), drainQueue);
+
+    drainQueue();
     Mso::Test_ThreadPoolSchedulerWin_EnableThreadPoolWorkTracking(false);
   }
 };
diff --git a/vnext/Mso/src/dispatchQueue/uiScheduler_winrt.cpp b/vnext/Mso/src/dispatchQueue/uiScheduler_winrt.cpp
@@ -337,7 +337,9 @@ void UISchedulerWinRT<TDispatcherTraits>::Shutdown() noexcept {
 template <typename TDispatcherTraits>
 void UISchedulerWinRT<TDispatcherTraits>::AwaitTermination() noexcept {
   Shutdown();
-  m_terminationEvent.Wait();
+  if (m_threadId != std::this_thread::get_id()) {
+    m_terminationEvent.Wait();
+  }
 }
 
 template <typename TDispatcherTraits>
@@ -411,7 +413,7 @@ void UISchedulerWinRT<TDispatcherTraits>::CleanupContext::CheckTermination() noe
 }
 
 //=============================================================================
-// DispatchQueueStatic::MakeCurrentThreadUIScheduler implementation
+// DispatchQueueStatic::GetCurrentUIThreadQueue implementation
 //=============================================================================
 
 DispatchQueue DispatchQueueStatic::GetCurrentUIThreadQueue() noexcept {
