Skip to content

Commit 606c124

Browse files
VSadovVSadov
committed
Special-case OperationCanceledException in Task-returning thunks to produce canceled task.
1 parent e96e206 commit 606c124

File tree

5 files changed

+116
-78
lines changed

5 files changed

+116
-78
lines changed

‎src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs‎

Lines changed: 58 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -558,54 +558,83 @@ private static bool QueueContinuationFollowUpActionIfNecessary<T, TOps>(T task,
558558
#pragma warning disable CA1859
559559
// When a Task-returning thunk gets a continuation result
560560
// it calls here to make a Task that awaits on the current async state.
561-
private static Task<T?> FinalizeTaskReturningThunk<T>(Continuation continuation)
561+
private static Task<T?> FinalizeTaskReturningThunk<T>(Continuation continuation, Exception ex)
562562
{
563-
Continuation finalContinuation = new Continuation();
564-
565-
// Note that the exact location the return value is placed is tied
566-
// into getAsyncResumptionStub in the VM, so do not change this
567-
// without also changing that code (and the JIT).
568-
if (RuntimeHelpers.IsReferenceOrContainsReferences<T>())
563+
if (continuation is not null)
569564
{
570-
finalContinuation.Flags = CorInfoContinuationFlags.CORINFO_CONTINUATION_RESULT_IN_GCDATA | CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION;
571-
finalContinuation.GCData = new object[1];
565+
Continuation finalContinuation = new Continuation();
566+
567+
// Note that the exact location the return value is placed is tied
568+
// into getAsyncResumptionStub in the VM, so do not change this
569+
// without also changing that code (and the JIT).
570+
if (RuntimeHelpers.IsReferenceOrContainsReferences<T>())
571+
{
572+
finalContinuation.Flags = CorInfoContinuationFlags.CORINFO_CONTINUATION_RESULT_IN_GCDATA | CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION;
573+
finalContinuation.GCData = new object[1];
574+
}
575+
else
576+
{
577+
finalContinuation.Flags = CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION;
578+
finalContinuation.Data = new byte[Unsafe.SizeOf<T>()];
579+
}
580+
581+
continuation.Next = finalContinuation;
582+
583+
ThunkTask<T?> result = new();
584+
result.HandleSuspended(continuation);
585+
return result;
572586
}
573587
else
574588
{
575-
finalContinuation.Flags = CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION;
576-
finalContinuation.Data = new byte[Unsafe.SizeOf<T>()];
589+
Task<T?> task = new();
590+
// Tail of AsyncTaskMethodBuilderT.SetException
591+
bool successfullySet = ex is OperationCanceledException oce ?
592+
task.TrySetCanceled(oce.CancellationToken, oce) :
593+
task.TrySetException(ex);
594+
595+
Debug.Assert(successfullySet);
596+
return task;
577597
}
578-
579-
continuation.Next = finalContinuation;
580-
581-
ThunkTask<T?> result = new();
582-
result.HandleSuspended(continuation);
583-
return result;
584598
}
585599

586-
private static Task FinalizeTaskReturningThunk(Continuation continuation)
600+
private static Task FinalizeTaskReturningThunk(Continuation continuation, Exception ex)
587601
{
588-
Continuation finalContinuation = new Continuation
602+
if (continuation is not null)
589603
{
590-
Flags = CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION,
591-
};
592-
continuation.Next = finalContinuation;
604+
Continuation finalContinuation = new Continuation
605+
{
606+
Flags = CorInfoContinuationFlags.CORINFO_CONTINUATION_NEEDS_EXCEPTION,
607+
};
608+
continuation.Next = finalContinuation;
593609

594-
ThunkTask result = new();
595-
result.HandleSuspended(continuation);
596-
return result;
610+
ThunkTask result = new();
611+
result.HandleSuspended(continuation);
612+
return result;
613+
}
614+
else
615+
{
616+
Debug.Assert(ex is not null);
617+
Task task = new();
618+
// Tail of AsyncTaskMethodBuilderT.SetException
619+
bool successfullySet = ex is OperationCanceledException oce ?
620+
task.TrySetCanceled(oce.CancellationToken, oce) :
621+
task.TrySetException(ex);
622+
623+
Debug.Assert(successfullySet);
624+
return task;
625+
}
597626
}
598627

599-
private static ValueTask<T?> FinalizeValueTaskReturningThunk<T>(Continuation continuation)
628+
private static ValueTask<T?> FinalizeValueTaskReturningThunk<T>(Continuation continuation, Exception ex)
600629
{
601630
// We only come to these methods in the expensive case (already
602631
// suspended), so ValueTask optimization here is not relevant.
603-
return new ValueTask<T?>(FinalizeTaskReturningThunk<T>(continuation));
632+
return new ValueTask<T?>(FinalizeTaskReturningThunk<T>(continuation, ex));
604633
}
605634

606-
private static ValueTask FinalizeValueTaskReturningThunk(Continuation continuation)
635+
private static ValueTask FinalizeValueTaskReturningThunk(Continuation continuation, Exception ex)
607636
{
608-
return new ValueTask(FinalizeTaskReturningThunk(continuation));
637+
return new ValueTask(FinalizeTaskReturningThunk(continuation, ex));
609638
}
610639

611640
[MethodImpl(MethodImplOptions.AggressiveInlining)]

‎src/coreclr/vm/asyncthunks.cpp‎

Lines changed: 44 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,43 @@ bool MethodDesc::TryGenerateAsyncThunk(DynamicResolver** resolver, COR_ILMETHOD_
5858
return true;
5959
}
6060

61+
// provided an async method, emits a Task-returning wrapper.
6162
void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig& thunkMsig, ILStubLinker* pSL)
6263
{
6364
_ASSERTE(!pAsyncOtherVariant->IsAsyncThunkMethod());
6465

66+
// Emits roughly the following code:
67+
//
68+
// ExecutionAndSyncBlockStore store = default;
69+
// store.Push();
70+
// try
71+
// {
72+
// Continuation cont;
73+
// Exception ex;
74+
// try
75+
// {
76+
// T result = Inner(args);
77+
// // call an intrisic to see if the call above produced a continuation
78+
// cont = StubHelpers.AsyncCallContinuation();
79+
// if (cont == null)
80+
// return Task.FromResult(result);
81+
// }
82+
// catch (Exception ex1)
83+
// {
84+
// ex = ex1;
85+
// }
86+
//
87+
// return FinalizeTaskReturningThunk(cont, ex);
88+
// }
89+
// finally
90+
// {
91+
// store.Pop();
92+
// }
93+
6594
ILCodeStream* pCode = pSL->NewCodeStream(ILStubLinker::kDispatch);
6695

6796
unsigned continuationLocal = pCode->NewLocal(LocalDesc(CoreLibBinder::GetClass(CLASS__CONTINUATION)));
97+
unsigned exceptionLocal = pCode->NewLocal(LocalDesc(CoreLibBinder::GetClass(CLASS__EXCEPTION)));
6898

6999
TypeHandle thTaskRet = thunkMsig.GetRetTypeHandleThrowing();
70100

@@ -84,7 +114,7 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
84114
DWORD executionAndSyncBlockStoreLocal = pCode->NewLocal(executionAndSyncBlockStoreLocalDesc);
85115

86116
ILCodeLabel* returnTaskLabel = pCode->NewCodeLabel();
87-
ILCodeLabel* suspendedLabel = pCode->NewCodeLabel();
117+
ILCodeLabel* finalizeTaskLabel = pCode->NewCodeLabel();
88118
ILCodeLabel* finishedLabel = pCode->NewCodeLabel();
89119

90120
pCode->EmitLDLOCA(executionAndSyncBlockStoreLocal);
@@ -171,7 +201,7 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
171201
pCode->EmitLDLOC(continuationLocal);
172202
pCode->EmitBRFALSE(finishedLabel);
173203

174-
pCode->EmitLEAVE(suspendedLabel);
204+
pCode->EmitLEAVE(finalizeTaskLabel);
175205

176206
pCode->EmitLabel(finishedLabel);
177207
if (logicalResultLocal != UINT_MAX)
@@ -202,37 +232,12 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
202232
// Catch
203233
{
204234
pCode->BeginCatchBlock(pCode->GetToken(CoreLibBinder::GetClass(CLASS__EXCEPTION)));
205-
206-
int fromExceptionToken;
207-
if (logicalResultLocal != UINT_MAX)
208-
{
209-
MethodDesc* fromExceptionMD;
210-
if (isValueTask)
211-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__VALUETASK__FROM_EXCEPTION_1);
212-
else
213-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__TASK__FROM_EXCEPTION_1);
214-
215-
fromExceptionMD = FindOrCreateAssociatedMethodDesc(fromExceptionMD, fromExceptionMD->GetMethodTable(), FALSE, Instantiation(&thLogicalRetType, 1), FALSE);
216-
217-
fromExceptionToken = GetTokenForGenericMethodCallWithAsyncReturnType(pCode, fromExceptionMD);
218-
}
219-
else
220-
{
221-
MethodDesc* fromExceptionMD;
222-
if (isValueTask)
223-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__VALUETASK__FROM_EXCEPTION);
224-
else
225-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__TASK__FROM_EXCEPTION);
226-
227-
fromExceptionToken = pCode->GetToken(fromExceptionMD);
228-
}
229-
pCode->EmitCALL(fromExceptionToken, 1, 1);
230-
pCode->EmitSTLOC(returnTaskLocal);
231-
pCode->EmitLEAVE(returnTaskLabel);
235+
pCode->EmitSTLOC(exceptionLocal);
236+
pCode->EmitLEAVE(finalizeTaskLabel);
232237
pCode->EndCatchBlock();
233238
}
234239

235-
pCode->EmitLabel(suspendedLabel);
240+
pCode->EmitLabel(finalizeTaskLabel);
236241

237242
int finalizeTaskReturningThunkToken;
238243
if (logicalResultLocal != UINT_MAX)
@@ -255,8 +260,10 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
255260
md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINALIZE_TASK_RETURNING_THUNK);
256261
finalizeTaskReturningThunkToken = pCode->GetToken(md);
257262
}
263+
258264
pCode->EmitLDLOC(continuationLocal);
259-
pCode->EmitCALL(finalizeTaskReturningThunkToken, 1, 1);
265+
pCode->EmitLDLOC(exceptionLocal);
266+
pCode->EmitCALL(finalizeTaskReturningThunkToken, 2, 1);
260267
pCode->EmitSTLOC(returnTaskLocal);
261268
pCode->EmitLEAVE(returnTaskLabel);
262269

@@ -447,15 +454,17 @@ int MethodDesc::GetTokenForGenericTypeMethodCallWithAsyncReturnType(ILCodeStream
447454
return pCode->GetToken(md, typeSigToken);
448455
}
449456

457+
// provided a Task-returning method, emits an async wrapper.
450458
void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& msig, ILStubLinker* pSL)
451459
{
452460
_ASSERTE(!pAsyncOtherVariant->IsAsyncThunkMethod());
453461
_ASSERTE(!pAsyncOtherVariant->IsVoid());
454462

455-
// TODO: (async) we may now be able to just do "AsyncHelpers.Await(other(arg))",
456-
// but would need to make sure it is not "optimized" back to calling this same thunk.
457-
458-
// Implement IL that is effectively the following
463+
// The thunk is roughly the same as "AsyncHelpers.Await(other(arg))", but without any
464+
// synchronization preferences. In a case of resumption the thunk will run on whatever the
465+
// thread/context was used by the awaiter to call us on completion.
466+
//
467+
// Implement IL that is effectively the following:
459468
/*
460469
{
461470
TaskAwaiter<RetType> awaiter = other(arg).GetAwaiter();

‎src/coreclr/vm/binder.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,12 @@ void CoreLibBinder::BuildConvertedSignature(const BYTE* pSig, SigBuilder * pSigB
431431
callConv = *pSig++;
432432
pSigBuilder->AppendData(callConv);
433433

434+
if ((callConv & IMAGE_CEE_CS_CALLCONV_GENERIC) != 0)
435+
{
436+
unsigned genericArgCount = *pSig++;
437+
pSigBuilder->AppendData(genericArgCount);
438+
}
439+
434440
if ((callConv & IMAGE_CEE_CS_CALLCONV_MASK) == IMAGE_CEE_CS_CALLCONV_DEFAULT) {
435441
// arg count
436442
argCount = *pSig++;
@@ -442,12 +448,6 @@ void CoreLibBinder::BuildConvertedSignature(const BYTE* pSig, SigBuilder * pSigB
442448
argCount = 0;
443449
}
444450

445-
if ((callConv & IMAGE_CEE_CS_CALLCONV_GENERIC) != 0)
446-
{
447-
unsigned genericArgCount = *pSig++;
448-
pSigBuilder->AppendData(genericArgCount);
449-
}
450-
451451
// <= because we want to include the return value or the field
452452
for (unsigned i = 0; i <= argCount; i++) {
453453
if (ConvertType(pSig, pSigBuilder))

‎src/coreclr/vm/corelib.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -731,10 +731,10 @@ DEFINE_METHOD(ASYNC_HELPERS, ALLOC_CONTINUATION, AllocContinuation,
731731
DEFINE_METHOD(ASYNC_HELPERS, ALLOC_CONTINUATION_METHOD, AllocContinuationMethod, NoSig)
732732
DEFINE_METHOD(ASYNC_HELPERS, ALLOC_CONTINUATION_CLASS, AllocContinuationClass, NoSig)
733733
DEFINE_METHOD(ASYNC_HELPERS, ALLOC_CONTINUATION_RESULT_BOX, AllocContinuationResultBox, SM_VoidPtr_RetObj)
734-
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_TASK_RETURNING_THUNK, FinalizeTaskReturningThunk, SM_Continuation_RetTask)
735-
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_TASK_RETURNING_THUNK_1, FinalizeTaskReturningThunk, GM_Continuation_RetTaskOfT)
736-
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_VALUETASK_RETURNING_THUNK, FinalizeValueTaskReturningThunk, SM_Continuation_RetValueTask)
737-
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_VALUETASK_RETURNING_THUNK_1, FinalizeValueTaskReturningThunk, GM_Continuation_RetValueTaskOfT)
734+
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_TASK_RETURNING_THUNK, FinalizeTaskReturningThunk, SM_Continuation_Exception_RetTask)
735+
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_TASK_RETURNING_THUNK_1, FinalizeTaskReturningThunk, GM_Continuation_Exception_RetTaskOfT)
736+
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_VALUETASK_RETURNING_THUNK, FinalizeValueTaskReturningThunk, SM_Continuation_Exception_RetValueTask)
737+
DEFINE_METHOD(ASYNC_HELPERS, FINALIZE_VALUETASK_RETURNING_THUNK_1, FinalizeValueTaskReturningThunk, GM_Continuation_Exception_RetValueTaskOfT)
738738
DEFINE_METHOD(ASYNC_HELPERS, UNSAFE_AWAIT_AWAITER_1, UnsafeAwaitAwaiter, GM_T_RetVoid)
739739
DEFINE_METHOD(ASYNC_HELPERS, CAPTURE_EXECUTION_CONTEXT, CaptureExecutionContext, NoSig)
740740
DEFINE_METHOD(ASYNC_HELPERS, RESTORE_EXECUTION_CONTEXT, RestoreExecutionContext, NoSig)

‎src/coreclr/vm/metasig.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -626,10 +626,10 @@ DEFINE_METASIG(SM(PtrByte_RetVoid, P(b), v))
626626

627627
DEFINE_METASIG_T(SM(RetContinuation, , C(CONTINUATION)))
628628
DEFINE_METASIG(GM(T_RetVoid, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, M(0), v))
629-
DEFINE_METASIG_T(SM(Continuation_RetTask, C(CONTINUATION), C(TASK)))
630-
DEFINE_METASIG_T(GM(Continuation_RetTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, C(CONTINUATION), GI(C(TASK_1), 1, M(0))))
631-
DEFINE_METASIG_T(SM(Continuation_RetValueTask, C(CONTINUATION), g(VALUETASK)))
632-
DEFINE_METASIG_T(GM(Continuation_RetValueTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, C(CONTINUATION), GI(g(VALUETASK_1), 1, M(0))))
629+
DEFINE_METASIG_T(SM(Continuation_Exception_RetTask, C(CONTINUATION) C(EXCEPTION), C(TASK)))
630+
DEFINE_METASIG_T(GM(Continuation_Exception_RetTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, C(CONTINUATION) C(EXCEPTION), GI(C(TASK_1), 1, M(0))))
631+
DEFINE_METASIG_T(SM(Continuation_Exception_RetValueTask, C(CONTINUATION) C(EXCEPTION), g(VALUETASK)))
632+
DEFINE_METASIG_T(GM(Continuation_Exception_RetValueTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, C(CONTINUATION) C(EXCEPTION), GI(g(VALUETASK_1), 1, M(0))))
633633

634634
// Undefine macros in case we include the file again in the compilation unit
635635

0 commit comments

Comments
 (0)