Skip to content

Commit ee804b5

Browse files
committed
Specialcase OperationCanceledException in Tasj-returning thunks to produce cancelled task.
1 parent 6a6de63 commit ee804b5

File tree

5 files changed

+85
-77
lines changed

5 files changed

+85
-77
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: 13 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
7070
// try
7171
// {
7272
// Continuation cont;
73+
// Exception ex;
7374
// try
7475
// {
7576
// T result = Inner(args);
@@ -78,12 +79,12 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
7879
// if (cont == null)
7980
// return Task.FromResult(result);
8081
// }
81-
// catch (Exception ex)
82+
// catch (Exception ex1)
8283
// {
83-
// return Task.FromException(ex);
84+
// ex = ex1;
8485
// }
8586
//
86-
// return FinalizeTaskReturningThunk(cont);
87+
// return FinalizeTaskReturningThunk(cont, ex);
8788
// }
8889
// finally
8990
// {
@@ -93,6 +94,7 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
9394
ILCodeStream* pCode = pSL->NewCodeStream(ILStubLinker::kDispatch);
9495

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

9799
TypeHandle thTaskRet = thunkMsig.GetRetTypeHandleThrowing();
98100

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

114116
ILCodeLabel* returnTaskLabel = pCode->NewCodeLabel();
115-
ILCodeLabel* suspendedLabel = pCode->NewCodeLabel();
117+
ILCodeLabel* finalizeTaskLabel = pCode->NewCodeLabel();
116118
ILCodeLabel* finishedLabel = pCode->NewCodeLabel();
117119

118120
pCode->EmitLDLOCA(executionAndSyncBlockStoreLocal);
@@ -199,7 +201,7 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
199201
pCode->EmitLDLOC(continuationLocal);
200202
pCode->EmitBRFALSE(finishedLabel);
201203

202-
pCode->EmitLEAVE(suspendedLabel);
204+
pCode->EmitLEAVE(finalizeTaskLabel);
203205

204206
pCode->EmitLabel(finishedLabel);
205207
if (logicalResultLocal != UINT_MAX)
@@ -230,37 +232,12 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
230232
// Catch
231233
{
232234
pCode->BeginCatchBlock(pCode->GetToken(CoreLibBinder::GetClass(CLASS__EXCEPTION)));
233-
234-
int fromExceptionToken;
235-
if (logicalResultLocal != UINT_MAX)
236-
{
237-
MethodDesc* fromExceptionMD;
238-
if (isValueTask)
239-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__VALUETASK__FROM_EXCEPTION_1);
240-
else
241-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__TASK__FROM_EXCEPTION_1);
242-
243-
fromExceptionMD = FindOrCreateAssociatedMethodDesc(fromExceptionMD, fromExceptionMD->GetMethodTable(), FALSE, Instantiation(&thLogicalRetType, 1), FALSE);
244-
245-
fromExceptionToken = GetTokenForGenericMethodCallWithAsyncReturnType(pCode, fromExceptionMD);
246-
}
247-
else
248-
{
249-
MethodDesc* fromExceptionMD;
250-
if (isValueTask)
251-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__VALUETASK__FROM_EXCEPTION);
252-
else
253-
fromExceptionMD = CoreLibBinder::GetMethod(METHOD__TASK__FROM_EXCEPTION);
254-
255-
fromExceptionToken = pCode->GetToken(fromExceptionMD);
256-
}
257-
pCode->EmitCALL(fromExceptionToken, 1, 1);
258-
pCode->EmitSTLOC(returnTaskLocal);
259-
pCode->EmitLEAVE(returnTaskLabel);
235+
pCode->EmitSTLOC(exceptionLocal);
236+
pCode->EmitLEAVE(finalizeTaskLabel);
260237
pCode->EndCatchBlock();
261238
}
262239

263-
pCode->EmitLabel(suspendedLabel);
240+
pCode->EmitLabel(finalizeTaskLabel);
264241

265242
int finalizeTaskReturningThunkToken;
266243
if (logicalResultLocal != UINT_MAX)
@@ -283,8 +260,10 @@ void MethodDesc::EmitTaskReturningThunk(MethodDesc* pAsyncOtherVariant, MetaSig&
283260
md = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__FINALIZE_TASK_RETURNING_THUNK);
284261
finalizeTaskReturningThunkToken = pCode->GetToken(md);
285262
}
263+
286264
pCode->EmitLDLOC(continuationLocal);
287-
pCode->EmitCALL(finalizeTaskReturningThunkToken, 1, 1);
265+
pCode->EmitLDLOC(exceptionLocal);
266+
pCode->EmitCALL(finalizeTaskReturningThunkToken, 2, 1);
288267
pCode->EmitSTLOC(returnTaskLocal);
289268
pCode->EmitLEAVE(returnTaskLabel);
290269

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)