Reduce stack consumption if unable to switch to apartment_context (#1276)

Author: oldnewthing

strings/base_coroutine_foundation.h

@@ -169,7 +169,7 @@ namespace winrt::impl
         }
 
         template <typename T>
-        auto await_suspend(coroutine_handle<T> handle)
+        bool await_suspend(coroutine_handle<T> handle)
         {
             this->set_cancellable_promise_from_handle(handle);
             return register_completed_callback(handle);
@@ -183,17 +183,10 @@ namespace winrt::impl
         }
 
     private:
-        auto register_completed_callback(coroutine_handle<> handle)
+        bool register_completed_callback(coroutine_handle<> handle)
         {
             async.Completed(disconnect_aware_handler(this, handle));
-#ifdef _RESUMABLE_FUNCTIONS_SUPPORTED
-            if (!suspending.exchange(false, std::memory_order_acquire))
-            {
-                handle.resume();
-            }
-#else
             return suspending.exchange(false, std::memory_order_acquire);
-#endif
         }
 
         static fire_and_forget cancel_asynchronously(Async async)

strings/base_coroutine_threadpool.h

@@ -77,7 +77,7 @@ namespace winrt::impl
         return 0;
     };
 
-    inline void resume_apartment_sync(com_ptr<IContextCallback> const& context, coroutine_handle<> handle, int32_t* failure)
+    [[nodiscard]] inline bool resume_apartment_sync(com_ptr<IContextCallback> const& context, coroutine_handle<> handle, int32_t* failure)
     {
         com_callback_args args{};
         args.data = handle.address();
@@ -87,8 +87,9 @@ namespace winrt::impl
         {
             // Resume the coroutine on the wrong apartment, but tell it why.
             *failure = result;
-            handle();
+            return false;
         }
+        return true;
     }
 
     struct threadpool_resume
@@ -103,7 +104,10 @@ namespace winrt::impl
     inline void __stdcall fallback_submit_threadpool_callback(void*, void* p) noexcept
     {
         std::unique_ptr<threadpool_resume> state{ static_cast<threadpool_resume*>(p) };
-        resume_apartment_sync(state->m_context, state->m_handle, state->m_failure);
+        if (!resume_apartment_sync(state->m_context, state->m_handle, state->m_failure))
+        {
+            state->m_handle.resume();
+        }
     }
 
     inline void resume_apartment_on_threadpool(com_ptr<IContextCallback> const& context, coroutine_handle<> handle, int32_t* failure)
@@ -113,7 +117,7 @@ namespace winrt::impl
         state.release();
     }
 
-    inline auto resume_apartment(resume_apartment_context const& context, coroutine_handle<> handle, int32_t* failure)
+    [[nodiscard]] inline auto resume_apartment(resume_apartment_context const& context, coroutine_handle<> handle, int32_t* failure)
     {
         WINRT_ASSERT(context.valid());
         if ((context.m_context == nullptr) || (context.m_context == try_capture<IContextCallback>(WINRT_IMPL_CoGetObjectContext)))
@@ -132,8 +136,7 @@ namespace winrt::impl
         }
         else
         {
-            resume_apartment_sync(context.m_context, handle, failure);
-            return true;
+            return resume_apartment_sync(context.m_context, handle, failure);
         }
     }
 }
@@ -331,17 +334,10 @@ namespace winrt::impl
             check_hresult(failure);
         }
 
-        auto await_suspend(impl::coroutine_handle<> handle)
+        bool await_suspend(impl::coroutine_handle<> handle)
         {
             auto context_copy = context;
-#ifdef _RESUMABLE_FUNCTIONS_SUPPORTED
-            if (!impl::resume_apartment(context_copy.context, handle, &failure))
-            {
-                handle.resume();
-            }
-#else
             return impl::resume_apartment(context_copy.context, handle, &failure);
-#endif
         }
     };
 

test/test/await_completed.cpp

@@ -27,46 +27,12 @@ namespace
     }
 #endif
 
-    // Simple awaiter that (inefficiently) resumes from inside a function nested in
-    // await_suspend, for the purpose of measuring how much stack it consumes.
-    // This is the best we can do with MSVC prerelease coroutines prior to 16.11.
-    // This simulates the behavior of await_adapter.
-    struct resume_sync_from_await_suspend
-    {
-        bool await_ready() { return false; }
-        template <typename T>
-        void await_suspend(winrt::impl::coroutine_handle<T> h) { resume_inner(h); }
-        void await_resume() { }
-
-    private:
-        void resume_inner(winrt::impl::coroutine_handle<> h) { h(); }
-    };
-
     IAsyncAction SyncCompletion()
     {
         uintptr_t initial = approximate_stack_pointer();
-        co_await resume_sync_from_await_suspend();
-        uintptr_t sync_usage = initial - approximate_stack_pointer();
-
-        initial = approximate_stack_pointer();
         co_await AlreadyCompleted();
         uintptr_t consumed = initial - approximate_stack_pointer();
-#ifdef _RESUMABLE_FUNCTIONS_SUPPORTED
-        // This branch is taken only for MSVC prerelease coroutines.
-        //
-        // MSVC prerelease coroutines prior to 16.11 do not implement "bool await_suspend" reliably,
-        // so we can't use it impl::await_adapter. We must resume inline inside await_suspend,
-        // so there is a small amount of stack usage. (Pre-16.11 and post-16.11 prerelease coroutines
-        // are interoperable, so we cannot change behavior based on which compiler we are using,
-        // because that would introduce ODR violations. Our first opportunity to change behavior
-        // is the ABI breaking change with MSVC standard-conforming coroutines.)
-        REQUIRE(consumed <= sync_usage);
-#else
-        (void)sync_usage;
-        // MSVC standard-conforming coroutines (as well as gcc and clang coroutines)
-        // support "bool await_suspend" just fine.
         REQUIRE(consumed == 0);
-#endif
     }
 
     // co_await the same apartment context and confirm that stack does not grow.
@@ -77,28 +43,9 @@ namespace
         winrt::apartment_context same_context;
 
         uintptr_t initial = approximate_stack_pointer();
-        co_await resume_sync_from_await_suspend();
-        uintptr_t sync_usage = initial - approximate_stack_pointer();
-
-        initial = approximate_stack_pointer();
         co_await same_context;
         uintptr_t consumed = initial - approximate_stack_pointer();
-
-#ifdef _RESUMABLE_FUNCTIONS_SUPPORTED
-        // This branch is taken only for MSVC prerelease coroutines.
-        //
-        // MSVC prerelease coroutines prior to 16.11 do not implement "bool await_suspend" reliably,
-        // so we can't use it impl::apartment_awaiter. We must resume inline inside await_suspend,
-        // so there is a small amount of stack usage. (Pre-16.11 and post-16.11 prerelease coroutines
-        // are interoperable, so we cannot change behavior based on which compiler we are using,
-        // because that would introduce ODR violations. Our first opportunity to change behavior
-        // is the ABI breaking change with MSVC standard-conforming coroutines.)
-        REQUIRE(consumed <= sync_usage);
-#else
-        // MSVC standard-conforming coroutines (as well as gcc and clang coroutines)
-        // support "bool await_suspend" just fine.
         REQUIRE(consumed == 0);
-#endif
     }
 }
 TEST_CASE("await_completed_await")