diff --git a/README.md b/README.md index ea10efb9..7c5dd681 100644 --- a/README.md +++ b/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)). diff --git a/src/UniTask.NetCoreTests/ExecutionContextTest.cs b/src/UniTask.NetCoreTests/ExecutionContextTest.cs new file mode 100644 index 00000000..4df6e6dd --- /dev/null +++ b/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 ambient = new AsyncLocal(); + + public ExecutionContextTest() + { + UniTaskScheduler.FlowExecutionContext = true; + } + + public void Dispose() + { + UniTaskScheduler.FlowExecutionContext = false; + } + + [Fact] + public async Task ResumedByAFlowWithoutTheContext_KeepsAsyncLocal() + { + ambient.Value = Before; + TaskCompletionSource read = new TaskCompletionSource(); + + async UniTask Body() + { + await read.Task; + return ambient.Value; + } + + // Suspends at its await before anything completes the read, so the resume is real. + UniTask body = Body(); + await CompleteFromAFlowWithoutTheContext(read); + + (await body).Should().Be(Before); + } + + [Fact] + public async Task WithoutResult_ResumedByAFlowWithoutTheContext_KeepsAsyncLocal() + { + ambient.Value = Before; + TaskCompletionSource read = new TaskCompletionSource(); + 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 read = new TaskCompletionSource(); + TaskCompletionSource seen = new TaskCompletionSource(); + + 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 read = new TaskCompletionSource(); + + 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 read = new TaskCompletionSource(); + + async UniTask Callee() + { + string previous = ambient.Value; + ambient.Value = "callee scope"; + try + { + await read.Task; + } + finally + { + ambient.Value = previous; + } + } + + async UniTask Caller() + { + await Callee(); + return ambient.Value; + } + + UniTask 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 read) + { + using (ExecutionContext.SuppressFlow()) + { + return Task.Run(() => read.SetResult(true)); + } + } + } +} diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskMethodBuilder.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskMethodBuilder.cs index 1aa990d9..6ee45082 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskMethodBuilder.cs +++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskMethodBuilder.cs @@ -83,6 +83,12 @@ public void AwaitOnCompleted(ref TAwaiter awaiter, ref AsyncUniTask.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(ref TAwaiter awaiter AsyncUniTask.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(ref TAwaiter awaiter public void Start(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(ref TAwaiter awaiter, ref AsyncUniTask.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(ref TAwaiter awaiter AsyncUniTask.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(ref TAwaiter awaiter public void Start(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(); } diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskVoidMethodBuilder.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskVoidMethodBuilder.cs index 82e91f38..04519f99 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskVoidMethodBuilder.cs +++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/AsyncUniTaskVoidMethodBuilder.cs @@ -83,6 +83,12 @@ public void AwaitOnCompleted(ref TAwaiter awaiter, ref AsyncUniTaskVoid.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(ref TAwaiter awaiter AsyncUniTaskVoid.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(ref TAwaiter awaiter public void Start(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(); } diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/StateMachineRunner.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/StateMachineRunner.cs index 1cffeced..a50e91d9 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/StateMachineRunner.cs +++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/CompilerServices/StateMachineRunner.cs @@ -5,6 +5,7 @@ using System.Linq; using System.Diagnostics; using System.Runtime.CompilerServices; +using System.Threading; namespace Cysharp.Threading.Tasks.CompilerServices { @@ -15,6 +16,9 @@ namespace Cysharp.Threading.Tasks.CompilerServices internal interface IStateMachineRunner { Action MoveNext { get; } +#if UNITASK_NETCORE + void CaptureExecutionContext(); +#endif void Return(); #if ENABLE_IL2CPP @@ -25,6 +29,9 @@ internal interface IStateMachineRunner internal interface IStateMachineRunnerPromise : IUniTaskSource { Action MoveNext { get; } +#if UNITASK_NETCORE + void CaptureExecutionContext(); +#endif UniTask Task { get; } void SetResult(); void SetException(Exception exception); @@ -33,6 +40,9 @@ internal interface IStateMachineRunnerPromise : IUniTaskSource internal interface IStateMachineRunnerPromise : IUniTaskSource { Action MoveNext { get; } +#if UNITASK_NETCORE + void CaptureExecutionContext(); +#endif UniTask Task { get; } void SetResult(T result); void SetException(Exception exception); @@ -61,6 +71,17 @@ internal sealed class AsyncUniTaskVoid : IStateMachineRunner, ITa TStateMachine stateMachine; public Action MoveNext { get; } +#if UNITASK_NETCORE + // Restored around the resume, as an async Task method does: without it the code after an + // await runs in whatever ExecutionContext the completing thread carries, losing AsyncLocals. + ExecutionContext executionContext; + static readonly ContextCallback RunInExecutionContext = state => ((AsyncUniTaskVoid)state).stateMachine.MoveNext(); + + public void CaptureExecutionContext() + { + executionContext = ExecutionContext.Capture(); + } +#endif public AsyncUniTaskVoid() { @@ -94,6 +115,9 @@ public void Return() { TaskTracker.RemoveTracking(this); stateMachine = default; +#if UNITASK_NETCORE + executionContext = null; +#endif pool.TryPush(this); } @@ -101,6 +125,15 @@ public void Return() [MethodImpl(MethodImplOptions.AggressiveInlining)] void Run() { +#if UNITASK_NETCORE + ExecutionContext context = executionContext; + if (context != null) + { + executionContext = null; + ExecutionContext.Run(context, RunInExecutionContext, this); + return; + } +#endif stateMachine.MoveNext(); } @@ -134,6 +167,17 @@ internal sealed class AsyncUniTask : IStateMachineRunnerPromise, readonly Action returnDelegate; #endif public Action MoveNext { get; } +#if UNITASK_NETCORE + // Restored around the resume, as an async Task method does: without it the code after an + // await runs in whatever ExecutionContext the completing thread carries, losing AsyncLocals. + ExecutionContext executionContext; + static readonly ContextCallback RunInExecutionContext = state => ((AsyncUniTask)state).stateMachine.MoveNext(); + + public void CaptureExecutionContext() + { + executionContext = ExecutionContext.Capture(); + } +#endif TStateMachine stateMachine; UniTaskCompletionSourceCore core; @@ -171,6 +215,9 @@ void Return() TaskTracker.RemoveTracking(this); core.Reset(); stateMachine = default; +#if UNITASK_NETCORE + executionContext = null; +#endif pool.TryPush(this); } @@ -179,6 +226,9 @@ bool TryReturn() TaskTracker.RemoveTracking(this); core.Reset(); stateMachine = default; +#if UNITASK_NETCORE + executionContext = null; +#endif return pool.TryPush(this); } @@ -186,6 +236,15 @@ bool TryReturn() [MethodImpl(MethodImplOptions.AggressiveInlining)] void Run() { +#if UNITASK_NETCORE + ExecutionContext context = executionContext; + if (context != null) + { + executionContext = null; + ExecutionContext.Run(context, RunInExecutionContext, this); + return; + } +#endif stateMachine.MoveNext(); } @@ -257,6 +316,17 @@ internal sealed class AsyncUniTask : IStateMachineRunnerPromis #endif public Action MoveNext { get; } +#if UNITASK_NETCORE + // Restored around the resume, as an async Task method does: without it the code after an + // await runs in whatever ExecutionContext the completing thread carries, losing AsyncLocals. + ExecutionContext executionContext; + static readonly ContextCallback RunInExecutionContext = state => ((AsyncUniTask)state).stateMachine.MoveNext(); + + public void CaptureExecutionContext() + { + executionContext = ExecutionContext.Capture(); + } +#endif TStateMachine stateMachine; UniTaskCompletionSourceCore core; @@ -294,6 +364,9 @@ void Return() TaskTracker.RemoveTracking(this); core.Reset(); stateMachine = default; +#if UNITASK_NETCORE + executionContext = null; +#endif pool.TryPush(this); } @@ -302,6 +375,9 @@ bool TryReturn() TaskTracker.RemoveTracking(this); core.Reset(); stateMachine = default; +#if UNITASK_NETCORE + executionContext = null; +#endif return pool.TryPush(this); } @@ -309,6 +385,15 @@ bool TryReturn() [MethodImpl(MethodImplOptions.AggressiveInlining)] void Run() { +#if UNITASK_NETCORE + ExecutionContext context = executionContext; + if (context != null) + { + executionContext = null; + ExecutionContext.Run(context, RunInExecutionContext, this); + return; + } +#endif // UnityEngine.Debug.Log($"MoveNext State:" + StateMachineUtility.GetState(stateMachine)); stateMachine.MoveNext(); } diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskScheduler.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskScheduler.cs index 2f91f2ad..7b1457f6 100644 --- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskScheduler.cs +++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTaskScheduler.cs @@ -17,6 +17,14 @@ public static class UniTaskScheduler /// public static bool PropagateOperationCanceledException = false; +#if UNITASK_NETCORE + + /// + /// Restore the ExecutionContext when an async UniTask method resumes, as async Task does, so AsyncLocal values survive an await. Adds a few nanoseconds per await and no allocation; set it at startup. Default is false. + /// + public static bool FlowExecutionContext = false; +#endif + #if UNITY_2018_3_OR_NEWER ///