Add `resume_agile` to allow coroutine to resume in any apartment (#1356)

strings/base_coroutine_foundation.h

@@ -99,44 +99,49 @@ namespace winrt::impl
         return async.GetResults();
     }
 
-    template<typename Awaiter>
-    struct disconnect_aware_handler
+    struct ignore_apartment_context {};
+
+    template<bool preserve_context, typename Awaiter>
+    struct disconnect_aware_handler : private std::conditional_t<preserve_context, resume_apartment_context, ignore_apartment_context>
     {
         disconnect_aware_handler(Awaiter* awaiter, coroutine_handle<> handle) noexcept
             : m_awaiter(awaiter), m_handle(handle) { }
 
-        disconnect_aware_handler(disconnect_aware_handler&& other) noexcept
-            : m_context(std::move(other.m_context))
-            , m_awaiter(std::exchange(other.m_awaiter, {}))
-            , m_handle(std::exchange(other.m_handle, {})) { }
+        disconnect_aware_handler(disconnect_aware_handler&& other) = default;
 
         ~disconnect_aware_handler()
         {
-            if (m_handle) Complete();
+            if (m_handle.value) Complete();
         }
 
         template<typename Async>
         void operator()(Async&&, Windows::Foundation::AsyncStatus status)
         {
-            m_awaiter->status = status;
+            m_awaiter.value->status = status;
             Complete();
         }
 
     private:
-        resume_apartment_context m_context;
-        Awaiter* m_awaiter;
-        coroutine_handle<> m_handle;
+        movable_primitive<Awaiter*> m_awaiter;
+        movable_primitive<coroutine_handle<>, nullptr> m_handle;
 
         void Complete()
         {
-            if (m_awaiter->suspending.exchange(false, std::memory_order_release))
+            if (m_awaiter.value->suspending.exchange(false, std::memory_order_release))
             {
-                m_handle = nullptr; // resumption deferred to await_suspend
+                m_handle.value = nullptr; // resumption deferred to await_suspend
             }
             else
             {
-                auto handle = std::exchange(m_handle, {});
-                if (!resume_apartment(m_context, handle, &m_awaiter->failure))
+                auto handle = m_handle.detach();
+                if constexpr (preserve_context)
+                {
+                    if (!resume_apartment(*this, handle, &m_awaiter.value->failure))
+                    {
+                        handle.resume();
+                    }
+                }
+                else
                 {
                     handle.resume();
                 }
@@ -145,7 +150,7 @@ namespace winrt::impl
     };
 
 #ifdef WINRT_IMPL_COROUTINES
-    template <typename Async>
+    template <typename Async, bool preserve_context = true>
     struct await_adapter : cancellable_awaiter<await_adapter<Async>>
     {
         await_adapter(Async const& async) : async(async) { }
@@ -185,7 +190,7 @@ namespace winrt::impl
     private:
         bool register_completed_callback(coroutine_handle<> handle)
         {
-            async.Completed(disconnect_aware_handler(this, handle));
+            async.Completed(disconnect_aware_handler<preserve_context, await_adapter>(this, handle));
             return suspending.exchange(false, std::memory_order_acquire);
         }
 
@@ -249,6 +254,15 @@ namespace winrt::impl
 }
 
 #ifdef WINRT_IMPL_COROUTINES
+WINRT_EXPORT namespace winrt
+{
+    template<typename Async, typename = std::enable_if_t<std::is_convertible_v<Async, winrt::Windows::Foundation::IAsyncInfo>>>
+    inline impl::await_adapter<Async, false> resume_agile(Async const& async)
+    {
+        return { async };
+    };
+}
+
 WINRT_EXPORT namespace winrt::Windows::Foundation
 {
     inline impl::await_adapter<IAsyncAction> operator co_await(IAsyncAction const& async)

strings/base_coroutine_threadpool.h

@@ -52,23 +52,14 @@ namespace winrt::impl
     {
         resume_apartment_context() = default;
         resume_apartment_context(std::nullptr_t) : m_context(nullptr), m_context_type(-1) {}
-        resume_apartment_context(resume_apartment_context const&) = default;
-        resume_apartment_context(resume_apartment_context&& other) noexcept :
-            m_context(std::move(other.m_context)), m_context_type(std::exchange(other.m_context_type, -1)) {}
-        resume_apartment_context& operator=(resume_apartment_context const&) = default;
-        resume_apartment_context& operator=(resume_apartment_context&& other) noexcept
-        {
-            m_context = std::move(other.m_context);
-            m_context_type = std::exchange(other.m_context_type, -1);
-            return *this;
-        }
+
         bool valid() const noexcept
         {
-            return m_context_type >= 0;
+            return m_context_type.value >= 0;
         }
 
         com_ptr<IContextCallback> m_context = try_capture<IContextCallback>(WINRT_IMPL_CoGetObjectContext);
-        int32_t m_context_type = get_apartment_type().first;
+        movable_primitive<int32_t, -1> m_context_type = get_apartment_type().first;
     };
 
     inline int32_t __stdcall resume_apartment_callback(com_callback_args* args) noexcept
@@ -124,7 +115,7 @@ namespace winrt::impl
         {
             return false;
         }
-        else if (context.m_context_type == 1 /* APTTYPE_MTA */)
+        else if (context.m_context_type.value == 1 /* APTTYPE_MTA */)
         {
             resume_background(handle);
             return true;

strings/base_meta.h

@@ -193,6 +193,25 @@ namespace winrt::impl
         }
     }
 
+    template<typename T, auto empty_value = T{}>
+    struct movable_primitive
+    {
+        T value = empty_value;
+        movable_primitive() = default;
+        movable_primitive(T const& init) : value(init) {}
+        movable_primitive(movable_primitive const&) = default;
+        movable_primitive(movable_primitive&& other) :
+            value(other.detach()) {}
+        movable_primitive& operator=(movable_primitive const&) = default;
+        movable_primitive& operator=(movable_primitive&& other)
+        {
+            value = other.detach();
+            return *this;
+        }
+
+        T detach() { return std::exchange(value, empty_value); }
+    };
+
     template <typename T, typename Enable = void>
     struct arg
     {

test/test/await_adapter.cpp

@@ -18,13 +18,9 @@ namespace
 
     static handle signal{ CreateEventW(nullptr, false, false, nullptr) };
 
-    IAsyncAction OtherForegroundAsync()
+    IAsyncAction OtherForegroundAsync(DispatcherQueue dispatcher)
     {
-        // Simple coroutine that completes on a unique STA thread.
-
-        auto controller = DispatcherQueueController::CreateOnDedicatedThread();
-        auto dispatcher = controller.DispatcherQueue();
-
+        // Simple coroutine that completes on the specified STA thread.
         co_await resume_foreground(dispatcher);
     }
 
@@ -35,37 +31,37 @@ namespace
         co_await resume_background();
     }
 
-    IAsyncAction ForegroundAsync(DispatcherQueue dispatcher)
+    // Coroutine that completes on dispatcher1, while potentially blocking dispatcher2.
+    IAsyncAction ForegroundAsync(DispatcherQueue dispatcher1, DispatcherQueue dispatcher2)
     {
         REQUIRE(!is_sta());
-        co_await resume_foreground(dispatcher);
+        co_await resume_foreground(dispatcher1);
         REQUIRE(is_sta());
 
         // This exercises one STA thread waiting on another thus one context callback
         // completing on another.
         uint32_t id = GetCurrentThreadId();
-        co_await OtherForegroundAsync();
+        co_await OtherForegroundAsync(dispatcher2);
         REQUIRE(id == GetCurrentThreadId());
 
-        // This just avoids the ForegroundAsync coroutine completing before
-        // BackgroundAsync waits on the result, forcing the Completed handler
-        // to be called on the foreground thread. This just makes the test
-        // success/failure more predictable.
+        // This Sleep() makes it more likely that the caller will actually suspend in await_suspend,
+        // so that the Completed handler triggers a resumption from the dispatcher1 thread.
         Sleep(100);
     }
 
-    fire_and_forget SignalFromForeground(DispatcherQueue dispatcher)
+    fire_and_forget SignalFromForeground(DispatcherQueue dispatcher1)
     {
         REQUIRE(!is_sta());
-        co_await resume_foreground(dispatcher);
+        co_await resume_foreground(dispatcher1);
         REQUIRE(is_sta());
 
-        // Previously, this signal was never raised because the foreground thread
-        // was always blocked waiting for ContextCallback to return.
+        // Previously, we never got here because of a deadlock:
+        // The dispatcher1 thread was blocked waiting for ContextCallback to return,
+        // but the ContextCallback is waiting for this event to get signaled.
         REQUIRE(SetEvent(signal.get()));
     }
 
-    IAsyncAction BackgroundAsync(DispatcherQueue dispatcher)
+    IAsyncAction BackgroundAsync(DispatcherQueue dispatcher1, DispatcherQueue dispatcher2)
     {
         // Switch to a background (MTA) thread.
         co_await resume_background();
@@ -76,19 +72,19 @@ namespace
         co_await OtherBackgroundAsync();
         REQUIRE(!is_sta());
 
-        // Wait for a coroutine that completes on a foreground (STA) thread.
-        co_await ForegroundAsync(dispatcher);
+        // Wait for a coroutine that completes on a the dispatcher1 thread (STA).
+        co_await ForegroundAsync(dispatcher1, dispatcher2);
 
         // Resumption should automatically switch to a background (MTA) thread
-        // without blocking the Completed handler (which would in turn block the foreground thread).
+        // without blocking the Completed handler (which would in turn block the dispatcher1 thread).
         REQUIRE(!is_sta());
 
-        // Attempt to signal from the foreground thread under the assumption
-        // that the foreground thread is not blocked.
-        SignalFromForeground(dispatcher);
+        // Attempt to signal from the dispatcher1 thread under the assumption
+        // that the dispatcher1 thread is not blocked.
+        SignalFromForeground(dispatcher1);
 
-        // Block the background (MTA) thread indefinitely until the signal is raied.
-        // Previously this would deadlock.
+        // Block the background (MTA) thread indefinitely until the signal is raised.
+        // Previously this would hang because the signal never got raised.
         REQUIRE(WAIT_OBJECT_0 == WaitForSingleObject(signal.get(), INFINITE));
     }
 }
@@ -99,9 +95,44 @@ TEST_CASE("await_adapter", "[.clang-crash]")
 #else
 TEST_CASE("await_adapter")
 #endif
+{
+    auto controller1 = DispatcherQueueController::CreateOnDedicatedThread();
+    auto controller2 = DispatcherQueueController::CreateOnDedicatedThread();
+
+    BackgroundAsync(controller1.DispatcherQueue(), controller2.DispatcherQueue()).get();
+    controller1.ShutdownQueueAsync().get();
+    controller2.ShutdownQueueAsync().get();
+}
+
+namespace
+{
+    IAsyncAction OtherBackgroundDelayAsync()
+    {
+        // Simple coroutine that completes on some MTA thread after a brief delay
+        // to ensure that the caller has suspended.
+
+        co_await resume_after(100ms);
+    }
+
+    IAsyncAction AgileAsync(DispatcherQueue dispatcher)
+    {
+        // Switch to the STA.
+        co_await resume_foreground(dispatcher);
+        REQUIRE(is_sta());
+
+        // Ask for agile resumption of a coroutine that finishes on a background thread.
+        // Add a 100ms delay to ensure we suspend.
+        co_await resume_agile(OtherBackgroundDelayAsync());
+        // We should be on the background thread now.
+        REQUIRE(!is_sta());
+    }
+}
+
+TEST_CASE("await_adapter_agile")
 {
     auto controller = DispatcherQueueController::CreateOnDedicatedThread();
     auto dispatcher = controller.DispatcherQueue();
 
-    BackgroundAsync(dispatcher).get();
+    AgileAsync(dispatcher).get();
+    controller.ShutdownQueueAsync().get();
 }