Skip to content

Commit 6c8dbe5

Browse files
committed
Base on Task instead of ValueTask
1 parent 8825898 commit 6c8dbe5

File tree

2 files changed

+110
-68
lines changed

2 files changed

+110
-68
lines changed

tracer/src/Datadog.Trace/ClrProfiler/CallTarget/Handlers/Continuations/ValueTaskContinuationGenerator.NetFx.cs

+48-32
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
#if !NETCOREAPP3_1_OR_GREATER
88
using System;
9+
using System.Runtime.ExceptionServices;
910
using System.Threading.Tasks;
1011
using Datadog.Trace.DuckTyping;
1112
using Datadog.Trace.Vendors.Serilog.Events;
@@ -86,39 +87,50 @@ public SyncCallbackHandler(ObjectContinuationMethodDelegate continuation, bool p
8687
return ValueTaskActivator<TReturn>.CreateInstance(secondTask);
8788
}
8889

89-
private async Task ContinuationAction(Task previousValueTask, TTarget? target, CallTargetState state)
90+
private async Task ContinuationAction(Task previousTask, TTarget? target, CallTargetState state)
9091
{
91-
try
92+
if (!previousTask.IsCompleted)
9293
{
93-
await previousValueTask.ConfigureAwait(_preserveContext);
94+
await new NoThrowAwaiter(previousTask, _preserveContext);
9495
}
95-
catch (Exception ex)
96+
97+
Exception? exception = null;
98+
99+
if (previousTask.Status == TaskStatus.Faulted)
100+
{
101+
exception = previousTask.Exception?.GetBaseException();
102+
}
103+
else if (previousTask.Status == TaskStatus.Canceled)
96104
{
97105
try
98106
{
99-
// *
100-
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
101-
// *
102-
_continuation(target, default, ex, in state);
107+
// The only supported way to extract the cancellation exception is to await the task
108+
await previousTask.ConfigureAwait(_preserveContext);
103109
}
104-
catch (Exception contEx)
110+
catch (Exception ex)
105111
{
106-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
112+
exception = ex;
107113
}
108-
109-
throw;
110114
}
111115

112116
try
113117
{
114118
// *
115119
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
116120
// *
117-
_continuation(target, default, default, in state);
121+
_continuation(target, null, exception, in state);
118122
}
119-
catch (Exception contEx)
123+
catch (Exception ex)
120124
{
121-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
125+
IntegrationOptions<TIntegration, TTarget>.LogException(ex);
126+
}
127+
128+
// *
129+
// If the original task throws an exception we rethrow it here.
130+
// *
131+
if (exception != null)
132+
{
133+
ExceptionDispatchInfo.Capture(exception).Throw();
122134
}
123135
}
124136
}
@@ -160,44 +172,48 @@ public override TReturn ExecuteCallback(TTarget? instance, TReturn? returnValue,
160172
return ValueTaskActivator<TReturn>.CreateInstance(secondTask);
161173
}
162174

163-
private async Task ContinuationAction(Task previousValueTask, TTarget? target, CallTargetState state, Exception? exception)
175+
private async Task ContinuationAction(Task previousTask, TTarget? target, CallTargetState state, Exception? exception)
164176
{
165-
if (exception != null)
177+
if (!previousTask.IsCompleted)
166178
{
167-
await _asyncContinuation(target, default, exception, in state).ConfigureAwait(_preserveContext);
179+
await new NoThrowAwaiter(previousTask, _preserveContext);
168180
}
169181

170-
try
182+
if (previousTask.Status == TaskStatus.Faulted)
171183
{
172-
await previousValueTask.ConfigureAwait(_preserveContext);
184+
exception ??= previousTask.Exception?.GetBaseException();
173185
}
174-
catch (Exception ex)
186+
else if (previousTask.Status == TaskStatus.Canceled)
175187
{
176188
try
177189
{
178-
// *
179-
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
180-
// *
181-
await _asyncContinuation(target, default, ex, in state).ConfigureAwait(_preserveContext);
190+
// The only supported way to extract the cancellation exception is to await the task
191+
await previousTask.ConfigureAwait(_preserveContext);
182192
}
183-
catch (Exception contEx)
193+
catch (Exception ex)
184194
{
185-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
195+
exception ??= ex;
186196
}
187-
188-
throw;
189197
}
190198

191199
try
192200
{
193201
// *
194202
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
195203
// *
196-
await _asyncContinuation(target, default, default, in state).ConfigureAwait(_preserveContext);
204+
await _asyncContinuation(target, null, exception, in state).ConfigureAwait(_preserveContext);
197205
}
198-
catch (Exception contEx)
206+
catch (Exception ex)
207+
{
208+
IntegrationOptions<TIntegration, TTarget>.LogException(ex);
209+
}
210+
211+
// *
212+
// If the original task throws an exception we rethrow it here.
213+
// *
214+
if (exception != null)
199215
{
200-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
216+
ExceptionDispatchInfo.Capture(exception).Throw();
201217
}
202218
}
203219
}

tracer/src/Datadog.Trace/ClrProfiler/CallTarget/Handlers/Continuations/ValueTaskContinuationGenerator`1.NetFx.cs

+62-36
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
#if !NETCOREAPP3_1_OR_GREATER
88
using System;
9+
using System.Runtime.ExceptionServices;
910
using System.Threading.Tasks;
1011
using Datadog.Trace.DuckTyping;
1112
using Datadog.Trace.Vendors.Serilog.Events;
@@ -89,43 +90,59 @@ public SyncCallbackHandler(ContinuationMethodDelegate continuation, bool preserv
8990
return ValueTaskActivator<TReturn, TResult?>.CreateInstance(secondTask);
9091
}
9192

92-
private async Task<TResult?> ContinuationAction(Task<TResult?> previousValueTask, TTarget? target, CallTargetState state)
93+
private async Task<TResult?> ContinuationAction(Task<TResult?> previousTask, TTarget? target, CallTargetState state)
9394
{
94-
TResult? result = default;
95-
try
95+
if (!previousTask.IsCompleted)
9696
{
97-
result = await previousValueTask.ConfigureAwait(_preserveContext);
97+
await new NoThrowAwaiter(previousTask, _preserveContext);
9898
}
99-
catch (Exception ex)
99+
100+
TResult? taskResult = default;
101+
Exception? exception = null;
102+
TResult? continuationResult = default;
103+
104+
if (previousTask.Status == TaskStatus.RanToCompletion)
105+
{
106+
taskResult = previousTask.Result;
107+
}
108+
else if (previousTask.Status == TaskStatus.Faulted)
109+
{
110+
exception = previousTask.Exception?.GetBaseException();
111+
}
112+
else if (previousTask.Status == TaskStatus.Canceled)
100113
{
101114
try
102115
{
103-
// *
104-
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
105-
// *
106-
_continuation(target, result, ex, in state);
116+
// The only supported way to extract the cancellation exception is to await the task
117+
await previousTask.ConfigureAwait(_preserveContext);
107118
}
108-
catch (Exception contEx)
119+
catch (Exception ex)
109120
{
110-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
121+
exception = ex;
111122
}
112-
113-
throw;
114123
}
115124

116125
try
117126
{
118127
// *
119128
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
120129
// *
121-
return _continuation(target, result, null, in state);
130+
continuationResult = _continuation(target, taskResult, exception, in state);
122131
}
123-
catch (Exception contEx)
132+
catch (Exception ex)
124133
{
125-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
134+
IntegrationOptions<TIntegration, TTarget>.LogException(ex);
126135
}
127136

128-
return result;
137+
// *
138+
// If the original task throws an exception we rethrow it here.
139+
// *
140+
if (exception != null)
141+
{
142+
ExceptionDispatchInfo.Capture(exception).Throw();
143+
}
144+
145+
return continuationResult;
129146
}
130147
}
131148

@@ -166,48 +183,57 @@ public override TReturn ExecuteCallback(TTarget? instance, TReturn? returnValue,
166183
return ValueTaskActivator<TReturn, TResult?>.CreateInstance(secondTask);
167184
}
168185

169-
private async Task<TResult?> ContinuationAction(Task<TResult?> previousValueTask, TTarget? target, CallTargetState state, Exception? exception)
186+
private async Task<TResult?> ContinuationAction(Task<TResult?> previousTask, TTarget? target, CallTargetState state, Exception? exception)
170187
{
171-
if (exception != null)
188+
TResult? taskResult = default;
189+
if (!previousTask.IsCompleted)
172190
{
173-
return await _asyncContinuation(target, default, exception, in state).ConfigureAwait(_preserveContext);
191+
await new NoThrowAwaiter(previousTask, _preserveContext);
174192
}
175193

176-
TResult? result = default;
177-
try
194+
if (previousTask.Status == TaskStatus.RanToCompletion)
178195
{
179-
result = await previousValueTask.ConfigureAwait(_preserveContext);
196+
taskResult = previousTask.Result;
180197
}
181-
catch (Exception ex)
198+
else if (previousTask.Status == TaskStatus.Faulted)
199+
{
200+
exception ??= previousTask.Exception?.GetBaseException();
201+
}
202+
else if (previousTask.Status == TaskStatus.Canceled)
182203
{
183204
try
184205
{
185-
// *
186-
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
187-
// *
188-
await _asyncContinuation(target, result, ex, in state).ConfigureAwait(_preserveContext);
206+
// The only supported way to extract the cancellation exception is to await the task
207+
await previousTask.ConfigureAwait(_preserveContext);
189208
}
190-
catch (Exception contEx)
209+
catch (Exception ex)
191210
{
192-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
211+
exception ??= ex;
193212
}
194-
195-
throw;
196213
}
197214

215+
TResult? continuationResult = default;
198216
try
199217
{
200218
// *
201219
// Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
202220
// *
203-
return await _asyncContinuation(target, result, null, in state).ConfigureAwait(_preserveContext);
221+
continuationResult = await _asyncContinuation(target, taskResult, exception, in state).ConfigureAwait(_preserveContext);
222+
}
223+
catch (Exception ex)
224+
{
225+
IntegrationOptions<TIntegration, TTarget>.LogException(ex);
204226
}
205-
catch (Exception contEx)
227+
228+
// *
229+
// If the original task throws an exception we rethrow it here.
230+
// *
231+
if (exception != null)
206232
{
207-
IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
233+
ExceptionDispatchInfo.Capture(exception).Throw();
208234
}
209235

210-
return result;
236+
return continuationResult;
211237
}
212238
}
213239

0 commit comments

Comments
 (0)