Support moving continuations to SynchronizationContext/TaskScheduler
#117314
Conversation
|
Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch |
There was a problem hiding this comment.
Pull Request Overview
Adds support for configuring where async continuations run—either on the captured SynchronizationContext/TaskScheduler or on the thread pool—by extending the JIT, VM, and BCL, and includes a new end-to-end test for SynchronizationContext behavior.
- JIT & VM: Introduce a new JIT flag and enum to track “continue on captured context” vs. thread-pool, update inlining logic to fatal on unsupported sites, and inject calls to
AsyncHelpers.CaptureContinuationContext. - BCL (
AsyncHelpers): AddCaptureContinuationContext, expand the continuation flags enum, and implementThunkTask/ThunkTaskCoreto run continuations respecting the captured context. - Tests: New
synchronization-context.csverifies that continuations post to the custom context when appropriate.
Reviewed Changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| src/tests/async/synchronization-context/synchronization-context.csproj | Add IL test project for sync-context tests |
| src/tests/async/synchronization-context/synchronization-context.cs | New test verifying ConfigureAwait behavior with SynchronizationContext |
| src/coreclr/vm/jitinterface.cpp | Expose new CaptureContinuationContext helper to the JIT |
| src/coreclr/vm/corelib.h | Define new METHOD__ASYNC_HELPERS__CAPTURE_CONTINUATION_CONTEXT |
| src/coreclr/jit/inline.def | Add CONTINUATION_HANDLING inline observation |
| src/coreclr/jit/importercalls.cpp | Handle and dump continuation-context flags, block inlining |
| src/coreclr/jit/importer.cpp | Wire up a new prefix flag for “continue on captured context” |
| src/coreclr/jit/gentree.h | Introduce ContinuationContextHandling enum |
| src/coreclr/jit/compiler.h | Add PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT flag |
| src/coreclr/jit/async.h | Reserve a GC slot index for continuation context |
| src/coreclr/jit/async.cpp | Allocate and populate continuation-context GC data and flags |
| src/coreclr/inc/corinfo.h | Define new CorInfoContinuationFlags bits |
| src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs | Implement CaptureContinuationContext and ThunkTask runtime support |
Comments suppressed due to low confidence (2)
src/tests/async/synchronization-context/synchronization-context.cs:51
- Consider adding a test scenario that verifies continuation capture with a custom TaskScheduler (e.g., by setting
TaskScheduler.Currentto a non-default scheduler) to ensure the TaskScheduler branch inCaptureContinuationContextis exercised.
Assert.Null(SynchronizationContext.Current);
src/coreclr/jit/importercalls.cpp:7900
- The inline observation enum value is defined as CONTINUATION_HANDLING (in inline.def), not CALLSITE_CONTINUATION_HANDLING. You should use InlineObservation::CONTINUATION_HANDLING to avoid a compile error.
inlineResult->NoteFatal(InlineObservation::CALLSITE_CONTINUATION_HANDLING);
|
There are still a few TODOs, but I think this is ready for review. cc @dotnet/jit-contrib for the JIT parts and @VSadov for other parts. @stephentoub would you also take a look at the managed parts? |
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
Show resolved
Hide resolved
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
Outdated
Show resolved
Hide resolved
This comment was marked as outdated.
This comment was marked as outdated.
WrappedYieldToThreadPool may be finished when TestSyncContext queues its continuation on it in which case no post will happen.
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
Show resolved
Hide resolved
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
Show resolved
Hide resolved
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
Show resolved
Hide resolved
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
Outdated
Show resolved
Hide resolved
Do you plan to add this in this PR? |
No, I want to do that in a follow-up. |
|
@AndyAyersMS can you please look at the JIT parts? |
|
|
||
| public object GetContinuationContext() | ||
| { | ||
| int index = 0; |
There was a problem hiding this comment.
Could this (and similar) index calculations be simplified to int index = BitOperations.PopCount((uint)(Flags & (CorInfoContinuationFlags.CORINFO_CONTINUATION_RESULT_IN_GCDATA | CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION)))?
This adds support for moving continuations after task awaits to the right context:
ConfigureAwait(true), continuations are posted to the capturedSynchronizationContextorTaskSchedulerby the BCLConfigureAwait(false)continuations are moved to the thread poolThe JIT inserts a call to a new
AsyncHelpers.CaptureContinuationContextwhen suspending because of a task await. That function either captures the currentSynchronizationContextorTaskSchedulerinto the continuation in a known place. Later the BCL will fetch the context from the continuation object based on a flag to determine whether the continuation needs to be posted somewhere else or whether it can be inlined. There is a newTask-derivedThunkTaskthat replaces the existing C#-async-function-with-a-loop and which handles the continuations explicitly.Inlining of continuations currently has no check on stack depth. This might be necessary to add. There are some differences compared to async 1 with how continuations are executed. In particular some common cases that resulted in unbounded stack usage with async 1 does not result in increasing stack usage with runtime async. Hence I am not 100% sure whether we need a mitigation or not, but more investigation is needed from my side. I'd like to leave that as a follow-up.
To get to the right behavior I have disabled inlining of task-awaited functions. This is a large concession, but getting the behavior right in the face of inlining proved to be quite tricky with the JIT's current internal representation. It is near top priority for me to reenable the inlining support, but I would like to work on it separately.
There is still one missing piece for correctness: saving and restoring the
Thread._synchronizationContextfield around async calls. It runs into similar trickiness when trying to represent the expected behavior (which is that the field is restored around synchronous calls, but not restored on resumption).