diff --git a/Plugins/UE5Coro/Source/UE5Coro/Private/AsyncAwaiters.cpp b/Plugins/UE5Coro/Source/UE5Coro/Private/AsyncAwaiters.cpp index 81c7a79..94c651d 100644 --- a/Plugins/UE5Coro/Source/UE5Coro/Private/AsyncAwaiters.cpp +++ b/Plugins/UE5Coro/Source/UE5Coro/Private/AsyncAwaiters.cpp @@ -84,6 +84,11 @@ FAsyncAwaiter Async::MoveToGameThread() return FAsyncAwaiter(ENamedThreads::GameThread); } +FAsyncAwaiter Async::MoveToSimilarThread() +{ + return FAsyncAwaiter(FTaskGraphInterface::Get().GetCurrentThreadIfKnown()); +} + FAsyncYieldAwaiter Async::Yield() { return {}; diff --git a/Plugins/UE5Coro/Source/UE5Coro/Public/UE5Coro/AsyncAwaiters.h b/Plugins/UE5Coro/Source/UE5Coro/Public/UE5Coro/AsyncAwaiters.h index f0003f8..888b6be 100644 --- a/Plugins/UE5Coro/Source/UE5Coro/Public/UE5Coro/AsyncAwaiters.h +++ b/Plugins/UE5Coro/Source/UE5Coro/Public/UE5Coro/AsyncAwaiters.h @@ -65,6 +65,12 @@ UE5CORO_API Private::FAsyncAwaiter MoveToThread(ENamedThreads::Type); * the game thread. */ UE5CORO_API Private::FAsyncAwaiter MoveToGameThread(); +/** Convenience function to resume on the same kind of named thread that this + * function was called on.
+ * co_await MoveToSimilarThread() is not useful. The return value should be + * stored to "remember" the original thread, then co_awaited later. */ +UE5CORO_API Private::FAsyncAwaiter MoveToSimilarThread(); + /** Always suspends the coroutine and resumes it on the same kind of named * thread that it's currently running on, or AnyThread otherwise.
* The return value of this function is reusable and always refers to the diff --git a/Plugins/UE5Coro/Source/UE5CoroTests/Private/AsyncAwaiterTest.cpp b/Plugins/UE5Coro/Source/UE5CoroTests/Private/AsyncAwaiterTest.cpp index 9cddc9d..0966c5a 100644 --- a/Plugins/UE5Coro/Source/UE5CoroTests/Private/AsyncAwaiterTest.cpp +++ b/Plugins/UE5Coro/Source/UE5CoroTests/Private/AsyncAwaiterTest.cpp @@ -139,6 +139,24 @@ void DoTest(FAutomationTestBase& Test) Test.TestTrue(TEXT("Triggered"), CoroToTest->Wait()); } + { + FEventRef CoroToTest; + std::atomic bMovedOut = false; + std::atomic bMovedIn = false; + World.Run(CORO + { + auto Return = Async::MoveToSimilarThread(); + co_await Async::MoveToNewThread(); + bMovedOut = !IsInGameThread(); + co_await Return; + bMovedIn = IsInGameThread(); + CoroToTest->Trigger(); + }); + FTestHelper::PumpGameThread(World, [&] { return CoroToTest->Wait(0); }); + Test.TestTrue(TEXT("Moved out"), bMovedOut); + Test.TestTrue(TEXT("Moved back in"), bMovedIn); + } + { std::atomic State = 0; World.Run(CORO