Fix weak references to coroutines (#1097)

Author: oldnewthing

strings/base_implements.h

@@ -831,7 +831,7 @@ namespace winrt::impl
         virtual ~root_implements() noexcept
         {
             // If a weak reference is created during destruction, this ensures that it is also destroyed.
-            subtract_reference();
+            subtract_final_reference();
         }
 
         int32_t __stdcall GetIids(uint32_t* count, guid** array) noexcept
@@ -897,10 +897,6 @@ namespace winrt::impl
 
             if (target == 0)
             {
-                // If a weak reference was previously created, the m_references value will not be stable value (won't be zero).
-                // This ensures destruction has a stable value during destruction.
-                m_references = 1;
-
                 if constexpr (has_final_release::value)
                 {
                     D::final_release(std::unique_ptr<D>(static_cast<D*>(this)));
@@ -992,7 +988,7 @@ namespace winrt::impl
         }
         catch (...) { return to_hresult(); }
 
-        uint32_t subtract_reference() noexcept
+        uint32_t subtract_final_reference() noexcept
         {
             if constexpr (is_weak_ref_source::value)
             {
@@ -1019,6 +1015,19 @@ namespace winrt::impl
             }
         }
 
+        uint32_t subtract_reference() noexcept
+        {
+            uint32_t result = subtract_final_reference();
+
+            if (result == 0)
+            {
+                // Ensure destruction happens with a stable reference count that isn't a weak reference.
+                m_references.store(1, std::memory_order_relaxed);
+            }
+
+            return result;
+        }
+
         template <typename T>
         winrt::weak_ref<T> get_weak()
         {

test/old_tests/UnitTests/weak.cpp

@@ -68,6 +68,49 @@ namespace
             REQUIRE(weak_self.get() == nullptr);
         }
     };
+
+    struct WeakCreateWeakInDestructor : implements<WeakCreateWeakInDestructor, IStringable>
+    {
+        winrt::weak_ref<WeakCreateWeakInDestructor>& weak_self;
+
+        WeakCreateWeakInDestructor(winrt::weak_ref<WeakCreateWeakInDestructor>& magic) : weak_self(magic) {}
+
+        ~WeakCreateWeakInDestructor()
+        {
+            // Creates a weak reference to itself in the destructor.
+            weak_self = get_weak();
+        }
+
+        hstring ToString()
+        {
+            return L"WeakCreateWeakInDestructor";
+        }
+    };
+
+#ifdef WINRT_IMPL_COROUTINES
+    // Returns an IAsyncAction that has already completed.
+    winrt::Windows::Foundation::IAsyncAction Action()
+    {
+        co_return;
+    }
+
+    // Returns an IAsyncAction that has not completed.
+    // Call the resume() handle to complete it.
+    winrt::Windows::Foundation::IAsyncAction SuspendAction(impl::coroutine_handle<>& resume)
+    {
+        struct awaiter
+        {
+            impl::coroutine_handle<>& resume;
+            bool await_ready() { return false; }
+            void await_suspend(impl::coroutine_handle<> handle) { resume = handle; }
+            void await_resume() {}
+        };
+
+        co_await awaiter{ resume };
+        co_return;
+    }
+
+#endif
 }
 
 TEST_CASE("weak,source")
@@ -413,3 +456,37 @@ TEST_CASE("weak,self")
     a.ToString();
     a = nullptr;
 }
+
+TEST_CASE("weak,create_weak_in_destructor")
+{
+    weak_ref<WeakCreateWeakInDestructor> magic;
+    IStringable a = make<WeakCreateWeakInDestructor>(magic);
+    a.ToString();
+    a = nullptr;
+    REQUIRE(magic.get() == nullptr);
+}
+
+#ifdef WINRT_IMPL_COROUTINES
+TEST_CASE("weak,coroutine")
+{
+    // Run a coroutine to completion. Confirm that weak references fail to resolve.
+    auto weak = winrt::weak_ref(Action());
+    REQUIRE(weak.get() == nullptr);
+
+    // Start a coroutine but don't complete it yet.
+    // Confirm that weak references resolve.
+    impl::coroutine_handle<> resume;
+    weak = winrt::weak_ref(SuspendAction(resume));
+    REQUIRE(weak.get() != nullptr);
+    // Now complete the coroutine. Confirm that weak references no longer resolve.
+    resume();
+    REQUIRE(weak.get() == nullptr);
+
+    // Verify that weak reference resolves as long as strong reference exists.
+    auto action = Action();
+    weak = winrt::weak_ref(action);
+    REQUIRE(weak.get() == action);
+    action = nullptr;
+    REQUIRE(weak.get() == nullptr);
+}
+#endif