Resume async coroutines co_awaiting another coroutine immediately ins…

…tead of marshaling back to the original thread

Docs/Async.md

@@ -192,9 +192,8 @@ caller when the callee coroutine finishes for any reason, **including**
 The return type of co_awaiting TCoroutine<T> is T.
 If the coroutine completed without co_returning a value, the result will be T().
 
-Async coroutines try to resume on a similar thread as they were on when co_await
-was issued (game thread to game thread, render thread to render thread, etc.),
-latent coroutines resume on the next tick after the callee ended.
+Async coroutines resume on the thread where the awaited coroutine finished.
+Latent coroutines resume on the next tick after the callee ended.
 co_awaiting a coroutine that's already complete will not release the current
 thread and will continue running with the result obtained synchronously.
 

Plugins/UE5Coro/Source/UE5Coro/Private/AsyncAwaiter.cpp

@@ -64,27 +64,15 @@ class FResumeTask
 
 bool FAsyncAwaiter::await_ready()
 {
-	// This needs to be scheduled after the coroutine's completion regardless of
-	// the target thread
-	if (ResumeAfter.has_value() && !ResumeAfter->IsDone())
-		return false;
-
 	// Don't move threads if we're already on the target thread
 	auto ThisThread = FTaskGraphInterface::Get().GetCurrentThreadIfKnown();
 	return (ThisThread & ThreadTypeMask) == (Thread & ThreadTypeMask);
 }
 
 void FAsyncAwaiter::Suspend(FPromise& Promise)
 {
-	auto* Task = TGraphTask<FResumeTask>::CreateTask()
-	                                     .ConstructAndHold(Thread, Promise);
-
-	// await_ready returning false and the coroutine having finished since is OK,
-	// ContinueWith will run this synchronously
-	if (ResumeAfter.has_value())
-		ResumeAfter->ContinueWith([Task] { Task->Unlock(); });
-	else
-		Task->Unlock();
+	TGraphTask<FResumeTask>::CreateTask().ConstructAndDispatchWhenReady(Thread,
+	                                                                    Promise);
 }
 
 FAsyncTimeAwaiter::FAsyncTimeAwaiter(const FAsyncTimeAwaiter& Other)

Plugins/UE5Coro/Source/UE5Coro/Private/AsyncAwaiters.cpp

@@ -76,12 +76,12 @@ struct FAutoStartResumeRunnable final : FRunnable
 
 FAsyncAwaiter Async::MoveToThread(ENamedThreads::Type Thread)
 {
-	return FAsyncAwaiter(Thread, {});
+	return FAsyncAwaiter(Thread);
 }
 
 FAsyncAwaiter Async::MoveToGameThread()
 {
-	return FAsyncAwaiter(ENamedThreads::GameThread, {});
+	return FAsyncAwaiter(ENamedThreads::GameThread);
 }
 
 FAsyncYieldAwaiter Async::Yield()

Plugins/UE5Coro/Source/UE5Coro/Public/UE5Coro/AsyncAwaiters.h

@@ -116,13 +116,9 @@ class [[nodiscard]] UE5CORO_API FAsyncAwaiter : public TAwaiter<FAsyncAwaiter>
 {
 	ENamedThreads::Type Thread;
 
-protected:
-	std::optional<TCoroutine<>> ResumeAfter;
-
 public:
-	explicit FAsyncAwaiter(ENamedThreads::Type Thread,
-	                       std::optional<TCoroutine<>> ResumeAfter)
-		: Thread(Thread), ResumeAfter(std::move(ResumeAfter)) { }
+	explicit FAsyncAwaiter(ENamedThreads::Type Thread)
+		: Thread(Thread) { }
 
 	bool await_ready();
 	void Suspend(FPromise&);

Plugins/UE5Coro/Source/UE5Coro/Public/UE5Coro/CoroutineAwaiters.h

@@ -41,30 +41,25 @@
 namespace UE5Coro::Private
 {
 template<typename T>
-class TAsyncCoroutineAwaiter : public FAsyncAwaiter
+class TAsyncCoroutineAwaiter : public TAwaiter<TAsyncCoroutineAwaiter<T>>
 {
+	TCoroutine<T> Antecedent;
+
 public:
 	explicit TAsyncCoroutineAwaiter(TCoroutine<T> Antecedent)
-		: FAsyncAwaiter(FTaskGraphInterface::Get().GetCurrentThreadIfKnown(),
-		                std::move(Antecedent)) { }
+		: Antecedent(std::move(Antecedent)) { }
 
-	// Prevent surprises with `co_await SomeCoroutine();` by making a copy.
-	// This cannot be moved as there could be another TCoroutine still owning it
-	T await_resume()
+	void Suspend(FPromise& Promise)
 	{
-		auto& Coro = static_cast<TCoroutine<T>&>(*ResumeAfter);
-		checkf(Coro.IsDone(), TEXT("Internal error: resuming too early"));
-		return Coro.GetResult();
+		Antecedent.ContinueWith([&Promise] { Promise.Resume(); });
 	}
-};
 
-template<>
-class TAsyncCoroutineAwaiter<void> : public FAsyncAwaiter
-{
-public:
-	explicit TAsyncCoroutineAwaiter(TCoroutine<> Antecedent)
-		: FAsyncAwaiter(FTaskGraphInterface::Get().GetCurrentThreadIfKnown(),
-		                std::move(Antecedent)) { }
+	T await_resume()
+	{
+		checkf(Antecedent.IsDone(), TEXT("Internal error: resuming too early"));
+		if constexpr (!std::is_void_v<T>)
+			return Antecedent.GetResult();
+	}
 };
 
 template<typename T>