Potential promise context lost fix #4583#4584
Conversation
|
Lightbeam link |
There was a problem hiding this comment.
Pull request overview
This PR aims to address #4583 by ensuring ExecutionContext (and therefore AsyncLocal-backed logging scopes/attributes) flows correctly across await points when awaiting Beamable.Common.Promise.
Changes:
- Capture and restore
ExecutionContextinPromise<T>.INotifyCompletion.OnCompleted. - Capture and restore
ExecutionContextwhen resuming the async state machine inPromiseAsyncMethodBuilder/PromiseAsyncMethodBuilder<T>.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Mirror TaskAwaiter.OnCompleted: flow ExecutionContext so AsyncLocal | ||
| // values set by the caller survive the await and are restored on the | ||
| // thread that resolves the promise. | ||
| var capturedContext = System.Threading.ExecutionContext.Capture(); | ||
| if (capturedContext == null) | ||
| { | ||
| ((ICriticalNotifyCompletion)this).UnsafeOnCompleted(continuation); | ||
| return; | ||
| } | ||
| ((ICriticalNotifyCompletion)this).UnsafeOnCompleted(() => | ||
| { | ||
| System.Threading.ExecutionContext.Run(capturedContext, s => ((Action)s)(), continuation); | ||
| }); |
There was a problem hiding this comment.
System.Threading.ExecutionContext is referenced here, but this file defines DISABLE_THREADING for UNITY_WEBGL and elsewhere conditionally compiles out System.Threading usage. As-is, this new code path can break WebGL builds (or any build with DISABLE_THREADING) because ExecutionContext lives in System.Threading. Consider wrapping the ExecutionContext capture/run logic in #if !DISABLE_THREADING and falling back to the previous UnsafeOnCompleted(continuation) behavior when threading is disabled.
| private void MoveNextWithCapturedContext(System.Threading.ExecutionContext capturedContext) | ||
| { | ||
| if (capturedContext == null) | ||
| { | ||
| _stateMachine.MoveNext(); | ||
| } | ||
| else | ||
| { | ||
| System.Threading.ExecutionContext.Run(capturedContext, s => ((IAsyncStateMachine)s).MoveNext(), _stateMachine); | ||
| } |
There was a problem hiding this comment.
This helper uses System.Threading.ExecutionContext, which conflicts with the existing DISABLE_THREADING compilation mode (defined for UNITY_WEBGL at the top of this file) that is meant to avoid System.Threading APIs. Please gate ExecutionContext usage behind #if !DISABLE_THREADING (or equivalent) and keep a non-threading fallback that just calls _stateMachine.MoveNext() so builds with DISABLE_THREADING still compile.
| if (_stateMachine == null) | ||
| { | ||
| _stateMachine = stateMachine; | ||
| _stateMachine.SetStateMachine(stateMachine); | ||
| } | ||
|
|
||
| var capturedContext = System.Threading.ExecutionContext.Capture(); | ||
| awaiter.UnsafeOnCompleted(() => MoveNextWithCapturedContext(capturedContext)); |
There was a problem hiding this comment.
System.Threading.ExecutionContext is captured/used in the async method builder here, but this file supports a DISABLE_THREADING build mode (notably for WebGL) where System.Threading APIs are avoided. Please add conditional compilation so the builder still compiles when DISABLE_THREADING is defined (eg, skip capture and resume the state machine directly).
| var capturedContext = System.Threading.ExecutionContext.Capture(); | ||
| awaiter.OnCompleted(() => MoveNextWithCapturedContext(capturedContext)); | ||
| } | ||
|
|
||
| public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>( | ||
| ref TAwaiter awaiter, ref TStateMachine stateMachine) | ||
| where TAwaiter : ICriticalNotifyCompletion | ||
| where TStateMachine : IAsyncStateMachine | ||
| { | ||
| AwaitOnCompleted(ref awaiter, ref stateMachine); | ||
| if (_stateMachine == null) | ||
| { | ||
| _stateMachine = stateMachine; | ||
| _stateMachine.SetStateMachine(stateMachine); | ||
| } | ||
|
|
||
| var capturedContext = System.Threading.ExecutionContext.Capture(); | ||
| awaiter.UnsafeOnCompleted(() => MoveNextWithCapturedContext(capturedContext)); |
There was a problem hiding this comment.
Same ExecutionContext dependency exists in the non-generic async method builder; without conditional compilation this can break DISABLE_THREADING targets (eg UNITY_WEBGL). Please apply the same #if !DISABLE_THREADING guard/fallback here as well so this code remains portable across supported Unity platforms.
| // Mirror AsyncTaskMethodBuilder: capture ExecutionContext at the await | ||
| // point and restore it before resuming the state machine so AsyncLocal | ||
| // values flow across the await. | ||
| var capturedContext = System.Threading.ExecutionContext.Capture(); | ||
| awaiter.OnCompleted(() => MoveNextWithCapturedContext(capturedContext)); |
There was a problem hiding this comment.
This change alters core await/continuation behavior (ExecutionContext/AsyncLocal flow) but there doesn’t appear to be a unit test that asserts AsyncLocal (or logging scope) values survive await Promise boundaries. Given there are existing PromiseTests, please add a focused test that sets an AsyncLocal value, awaits a Promise that completes on a different callback, and verifies the value is present after resumption (and ideally that suppression/clearing behaves as expected).
Fixes #4583 (potentially)
Needs in depth review since it was made with AI agent.