Add opt-in UniTaskScheduler.FlowExecutionContext for .NET

README.md

@@ -1129,7 +1129,7 @@ For .NET Core, use NuGet.
 
 UniTask of .NET Core version is a subset of Unity UniTask with PlayerLoop dependent methods removed.
 
-It runs at higher performance than the standard Task/ValueTask, but you should be careful to ignore the ExecutionContext/SynchronizationContext when using it. `AsyncLocal` also does not work because it ignores ExecutionContext.
+It runs at higher performance than the standard Task/ValueTask, but you should be careful to ignore the ExecutionContext/SynchronizationContext when using it. `AsyncLocal` also does not work because it ignores ExecutionContext. To make it work, set `UniTaskScheduler.FlowExecutionContext = true` at startup: an `async UniTask` method then restores the ExecutionContext when it resumes, as `async Task` does, for a few nanoseconds per await and no allocation.
 
 If you use UniTask internally, but provide ValueTask as an external API, you can write it like the following(Inspired by [PooledAwait](https://github.com/mgravell/PooledAwait)).
 

src/UniTask.NetCoreTests/ExecutionContextTest.cs

@@ -0,0 +1,156 @@
+using Cysharp.Threading.Tasks;
+using FluentAssertions;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace NetCoreTests
+{
+    // UniTaskScheduler.FlowExecutionContext is global, so these run apart from the other tests.
+    [CollectionDefinition(nameof(ExecutionContextTest), DisableParallelization = true)]
+    public class ExecutionContextTestCollection
+    {
+    }
+
+    // With FlowExecutionContext set, an async UniTask method treats the ExecutionContext as an
+    // async Task method does: it is restored when the method resumes, and what the method sets
+    // before its first await does not reach its caller.
+    [Collection(nameof(ExecutionContextTest))]
+    public class ExecutionContextTest : IDisposable
+    {
+        const string Before = "set before the call";
+
+        static readonly AsyncLocal<string> ambient = new AsyncLocal<string>();
+
+        public ExecutionContextTest()
+        {
+            UniTaskScheduler.FlowExecutionContext = true;
+        }
+
+        public void Dispose()
+        {
+            UniTaskScheduler.FlowExecutionContext = false;
+        }
+
+        [Fact]
+        public async Task ResumedByAFlowWithoutTheContext_KeepsAsyncLocal()
+        {
+            ambient.Value = Before;
+            TaskCompletionSource<bool> read = new TaskCompletionSource<bool>();
+
+            async UniTask<string> Body()
+            {
+                await read.Task;
+                return ambient.Value;
+            }
+
+            // Suspends at its await before anything completes the read, so the resume is real.
+            UniTask<string> body = Body();
+            await CompleteFromAFlowWithoutTheContext(read);
+
+            (await body).Should().Be(Before);
+        }
+
+        [Fact]
+        public async Task WithoutResult_ResumedByAFlowWithoutTheContext_KeepsAsyncLocal()
+        {
+            ambient.Value = Before;
+            TaskCompletionSource<bool> read = new TaskCompletionSource<bool>();
+            string seen = null;
+
+            async UniTask Body()
+            {
+                await read.Task;
+                seen = ambient.Value;
+            }
+
+            UniTask body = Body();
+            await CompleteFromAFlowWithoutTheContext(read);
+            await body;
+
+            seen.Should().Be(Before);
+        }
+
+        [Fact]
+        public async Task UniTaskVoid_ResumedByAFlowWithoutTheContext_KeepsAsyncLocal()
+        {
+            ambient.Value = Before;
+            TaskCompletionSource<bool> read = new TaskCompletionSource<bool>();
+            TaskCompletionSource<string> seen = new TaskCompletionSource<string>();
+
+            async UniTaskVoid Body()
+            {
+                await read.Task;
+                seen.SetResult(ambient.Value);
+            }
+
+            Body().Forget();
+            await CompleteFromAFlowWithoutTheContext(read);
+
+            (await seen.Task).Should().Be(Before);
+        }
+
+        [Fact]
+        public async Task WriteBeforeTheFirstAwait_DoesNotReachTheCaller()
+        {
+            ambient.Value = Before;
+            TaskCompletionSource<bool> read = new TaskCompletionSource<bool>();
+
+            async UniTask Callee()
+            {
+                ambient.Value = "set by the callee";
+                await read.Task;
+            }
+
+            UniTask callee = Callee();
+
+            ambient.Value.Should().Be(Before);
+            read.SetResult(true);
+            await callee;
+        }
+
+        // Restoring the context on resume without also isolating Start would hand a callee's
+        // scope back to its caller.
+        [Fact]
+        public async Task ScopeOpenedAcrossAnAwait_DoesNotLeakIntoTheCallerOnResume()
+        {
+            ambient.Value = Before;
+            TaskCompletionSource<bool> read = new TaskCompletionSource<bool>();
+
+            async UniTask Callee()
+            {
+                string previous = ambient.Value;
+                ambient.Value = "callee scope";
+                try
+                {
+                    await read.Task;
+                }
+                finally
+                {
+                    ambient.Value = previous;
+                }
+            }
+
+            async UniTask<string> Caller()
+            {
+                await Callee();
+                return ambient.Value;
+            }
+
+            UniTask<string> caller = Caller();
+            await CompleteFromAFlowWithoutTheContext(read);
+
+            (await caller).Should().Be(Before);
+        }
+
+        // Completes the read from a flow that never saw the value, as I/O completion does.
+        static Task CompleteFromAFlowWithoutTheContext(TaskCompletionSource<bool> read)
+        {
+            using (ExecutionContext.SuppressFlow())
+            {
+                return Task.Run(() => read.SetResult(true));
+            }
+        }
+    }
+}

src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskMethodBuilder.cs

@@ -83,6 +83,12 @@ public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref
                 AsyncUniTask<TStateMachine>.SetStateMachine(ref stateMachine, ref runnerPromise);
             }
 
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                runnerPromise.CaptureExecutionContext();
+            }
+#endif
             awaiter.OnCompleted(runnerPromise.MoveNext);
         }
 
@@ -99,6 +105,12 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
                 AsyncUniTask<TStateMachine>.SetStateMachine(ref stateMachine, ref runnerPromise);
             }
 
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                runnerPromise.CaptureExecutionContext();
+            }
+#endif
             awaiter.UnsafeOnCompleted(runnerPromise.MoveNext);
         }
 
@@ -108,6 +120,14 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
         public void Start<TStateMachine>(ref TStateMachine stateMachine)
             where TStateMachine : IAsyncStateMachine
         {
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                // Undo what the code before the first await sets, as an async Task method does.
+                default(AsyncTaskMethodBuilder).Start(ref stateMachine);
+                return;
+            }
+#endif
             stateMachine.MoveNext();
         }
 
@@ -214,6 +234,12 @@ public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref
                 AsyncUniTask<TStateMachine, T>.SetStateMachine(ref stateMachine, ref runnerPromise);
             }
 
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                runnerPromise.CaptureExecutionContext();
+            }
+#endif
             awaiter.OnCompleted(runnerPromise.MoveNext);
         }
 
@@ -230,6 +256,12 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
                 AsyncUniTask<TStateMachine, T>.SetStateMachine(ref stateMachine, ref runnerPromise);
             }
 
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                runnerPromise.CaptureExecutionContext();
+            }
+#endif
             awaiter.UnsafeOnCompleted(runnerPromise.MoveNext);
         }
 
@@ -239,6 +271,14 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
         public void Start<TStateMachine>(ref TStateMachine stateMachine)
             where TStateMachine : IAsyncStateMachine
         {
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                // Undo what the code before the first await sets, as an async Task method does.
+                default(AsyncTaskMethodBuilder).Start(ref stateMachine);
+                return;
+            }
+#endif
             stateMachine.MoveNext();
         }
 

src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskVoidMethodBuilder.cs

@@ -83,6 +83,12 @@ public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref
                 AsyncUniTaskVoid<TStateMachine>.SetStateMachine(ref stateMachine, ref runner);
             }
 
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                runner.CaptureExecutionContext();
+            }
+#endif
             awaiter.OnCompleted(runner.MoveNext);
         }
 
@@ -99,6 +105,12 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
                 AsyncUniTaskVoid<TStateMachine>.SetStateMachine(ref stateMachine, ref runner);
             }
 
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                runner.CaptureExecutionContext();
+            }
+#endif
             awaiter.UnsafeOnCompleted(runner.MoveNext);
         }
 
@@ -107,6 +119,14 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
         public void Start<TStateMachine>(ref TStateMachine stateMachine)
             where TStateMachine : IAsyncStateMachine
         {
+#if UNITASK_NETCORE
+            if (UniTaskScheduler.FlowExecutionContext)
+            {
+                // Undo what the code before the first await sets, as an async Task method does.
+                default(AsyncTaskMethodBuilder).Start(ref stateMachine);
+                return;
+            }
+#endif
             stateMachine.MoveNext();
         }