From 359595116a43dcc41cded734394f9c7864f12ec3 Mon Sep 17 00:00:00 2001
From: Raymond Chen <oldnewthing@users.noreply.github.com>
Date: Wed, 13 Sep 2023 13:12:52 -0700
Subject: [PATCH] Allow resume_agile to be stored in a variable

`resume_agile` exposes the ability to save the `await_adapter`
in a variable. This was not possible without `resume_agile` because
the `await_adapter` had previously been available only via
`operator co_await`, which means that it is created only in
response to an immediate attempt to `co_await` it, so we knew
that it would be consumed before its argument (possibly a temporary)
was destructed.

`resume_agile` returns the `await_adapter`, and we expect people
to await it immediately, but it's possible that they decide to
save it in a variable and await it later. In that case, we have
to record the `Async` as a value instead of a reference. We forward
the `resume_agile` argument into the `Async` so that it moves
if given an rvalue reference, or copies if given an lvalue reference.
This ensure that the common case where somebody does
`co_await resume_agile(DoSomething())`, we do not incur any additional
AddRefs or Releases.

Now that it's possible to `co_await` the `await_adapter` twice,
we have to worry about `await_suspend` being called twice. It had
previously assumed that `suspending` was true (since that's how it
was constructed), but that is no longer valid in the `resume_agile`
case if somebody tries to await the `resume_agile` twice. So we have
to force it to `true`. (Now, the second await will fail with
"illegal delegate assignment", but our failure to set `suspending`
to `true` led to double-resumption, which is super-bad.)
---
 strings/base_coroutine_foundation.h | 16 +++++++++-----
 test/test/await_adapter.cpp         | 33 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/strings/base_coroutine_foundation.h b/strings/base_coroutine_foundation.h
index ba61cd49a..670a65403 100644
--- a/strings/base_coroutine_foundation.h
+++ b/strings/base_coroutine_foundation.h
@@ -151,11 +151,12 @@ namespace winrt::impl
 
 #ifdef WINRT_IMPL_COROUTINES
     template <typename Async, bool preserve_context = true>
-    struct await_adapter : cancellable_awaiter<await_adapter<Async>>
+    struct await_adapter : cancellable_awaiter<await_adapter<Async, preserve_context>>
     {
-        await_adapter(Async const& async) : async(async) { }
+        template<typename T>
+        await_adapter(T&& async) : async(std::forward<T>(async)) { }
 
-        Async const& async;
+        std::conditional_t<preserve_context, Async const&, Async> async;
         Windows::Foundation::AsyncStatus status = Windows::Foundation::AsyncStatus::Started;
         int32_t failure = 0;
         std::atomic<bool> suspending = true;
@@ -190,6 +191,11 @@ namespace winrt::impl
     private:
         bool register_completed_callback(coroutine_handle<> handle)
         {
+            if constexpr (!preserve_context)
+            {
+                // Ensure that the illegal delegate assignment propagates properly.
+                suspending.store(true, std::memory_order_relaxed);
+            }
             async.Completed(disconnect_aware_handler<preserve_context, await_adapter>(this, handle));
             return suspending.exchange(false, std::memory_order_acquire);
         }
@@ -257,9 +263,9 @@ namespace winrt::impl
 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)
+    inline impl::await_adapter<std::decay_t<Async>, false> resume_agile(Async&& async)
     {
-        return { async };
+        return { std::forward<Async>(async) };
     };
 }
 
diff --git a/test/test/await_adapter.cpp b/test/test/await_adapter.cpp
index 701bc6b15..dccb9ce3d 100644
--- a/test/test/await_adapter.cpp
+++ b/test/test/await_adapter.cpp
@@ -136,3 +136,36 @@ TEST_CASE("await_adapter_agile")
     AgileAsync(dispatcher).get();
     controller.ShutdownQueueAsync().get();
 }
+
+namespace
+{
+    IAsyncAction AgileAsyncVariable(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. Store the resume_agile in a variable
+        // and await the variable.
+        auto op = resume_agile(OtherBackgroundDelayAsync());
+        co_await op;
+        // We should be on the background thread now.
+        REQUIRE(!is_sta());
+
+        // Second attempt to await the op should fail cleanly.
+        REQUIRE_THROWS_AS(co_await op, hresult_illegal_delegate_assignment);
+        // We should still be on the background thread.
+        REQUIRE(!is_sta());
+    }
+}
+
+
+TEST_CASE("await_adapter_agile_variable")
+{
+    auto controller = DispatcherQueueController::CreateOnDedicatedThread();
+    auto dispatcher = controller.DispatcherQueue();
+
+    AgileAsyncVariable(dispatcher).get();
+    controller.ShutdownQueueAsync().get();
+}