Skip to content

Commit 232785c

Browse files
authored
Used TaskToApm everywhere (#31874)
1 parent 998b230 commit 232785c

9 files changed

+18
-231
lines changed

src/Middleware/ResponseCaching/src/Microsoft.AspNetCore.ResponseCaching.csproj

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,4 +19,8 @@
1919
<Reference Include="Microsoft.Extensions.Logging.Abstractions" />
2020
</ItemGroup>
2121

22+
<ItemGroup>
23+
<Compile Include="$(RepoRoot)src\Shared\TaskToApm.cs" Link="Streams\TaskToApm.cs" />
24+
</ItemGroup>
25+
2226
</Project>

src/Middleware/ResponseCaching/src/Streams/ResponseCachingStream.cs

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -182,17 +182,9 @@ public override void WriteByte(byte value)
182182
}
183183

184184
public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
185-
{
186-
return StreamUtilities.ToIAsyncResult(WriteAsync(buffer, offset, count), callback, state);
187-
}
185+
=> TaskToApm.Begin(WriteAsync(buffer, offset, count, CancellationToken.None), callback, state);
188186

189187
public override void EndWrite(IAsyncResult asyncResult)
190-
{
191-
if (asyncResult == null)
192-
{
193-
throw new ArgumentNullException(nameof(asyncResult));
194-
}
195-
((Task)asyncResult).GetAwaiter().GetResult();
196-
}
188+
=> TaskToApm.End(asyncResult);
197189
}
198190
}

src/Middleware/ResponseCaching/src/Streams/SegmentWriteStream.cs

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -190,17 +190,9 @@ public override void WriteByte(byte value)
190190
}
191191

192192
public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
193-
{
194-
return StreamUtilities.ToIAsyncResult(WriteAsync(buffer, offset, count), callback, state);
195-
}
193+
=> TaskToApm.Begin(WriteAsync(buffer, offset, count, CancellationToken.None), callback, state);
196194

197195
public override void EndWrite(IAsyncResult asyncResult)
198-
{
199-
if (asyncResult == null)
200-
{
201-
throw new ArgumentNullException(nameof(asyncResult));
202-
}
203-
((Task)asyncResult).GetAwaiter().GetResult();
204-
}
196+
=> TaskToApm.End(asyncResult);
205197
}
206198
}
Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,6 @@
11
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

4-
using System;
5-
using System.Threading;
6-
using System.Threading.Tasks;
7-
84
namespace Microsoft.AspNetCore.ResponseCaching
95
{
106
internal static class StreamUtilities
@@ -14,28 +10,5 @@ internal static class StreamUtilities
1410
/// </summary>
1511
// Internal for testing
1612
internal static int BodySegmentSize { get; set; } = 81920;
17-
18-
internal static IAsyncResult ToIAsyncResult(Task task, AsyncCallback? callback, object? state)
19-
{
20-
var tcs = new TaskCompletionSource<int>(state);
21-
task.ContinueWith(t =>
22-
{
23-
if (t.IsFaulted)
24-
{
25-
tcs.TrySetException(t.Exception!.InnerExceptions);
26-
}
27-
else if (t.IsCanceled)
28-
{
29-
tcs.TrySetCanceled();
30-
}
31-
else
32-
{
33-
tcs.TrySetResult(0);
34-
}
35-
36-
callback?.Invoke(tcs.Task);
37-
}, CancellationToken.None, TaskContinuationOptions.None, TaskScheduler.Default);
38-
return tcs.Task;
39-
}
4013
}
4114
}

src/Middleware/ResponseCompression/src/Microsoft.AspNetCore.ResponseCompression.csproj

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,8 @@
1717
<Reference Include="Microsoft.Extensions.Options" />
1818
</ItemGroup>
1919

20+
<ItemGroup>
21+
<Compile Include="$(RepoRoot)src\Shared\TaskToApm.cs" Link="TaskToApm.cs" />
22+
</ItemGroup>
23+
2024
</Project>

src/Middleware/ResponseCompression/src/ResponseCompressionBody.cs

Lines changed: 2 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -174,51 +174,10 @@ public override void Write(byte[] buffer, int offset, int count)
174174
}
175175

176176
public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
177-
{
178-
var tcs = new TaskCompletionSource(state: state, TaskCreationOptions.RunContinuationsAsynchronously);
179-
InternalWriteAsync(buffer, offset, count, callback, tcs);
180-
return tcs.Task;
181-
}
182-
183-
private async void InternalWriteAsync(byte[] buffer, int offset, int count, AsyncCallback? callback, TaskCompletionSource tcs)
184-
{
185-
try
186-
{
187-
await WriteAsync(buffer, offset, count);
188-
tcs.TrySetResult();
189-
}
190-
catch (Exception ex)
191-
{
192-
tcs.TrySetException(ex);
193-
}
194-
195-
if (callback != null)
196-
{
197-
// Offload callbacks to avoid stack dives on sync completions.
198-
var ignored = Task.Run(() =>
199-
{
200-
try
201-
{
202-
callback(tcs.Task);
203-
}
204-
catch (Exception)
205-
{
206-
// Suppress exceptions on background threads.
207-
}
208-
});
209-
}
210-
}
177+
=> TaskToApm.Begin(WriteAsync(buffer, offset, count, CancellationToken.None), callback, state);
211178

212179
public override void EndWrite(IAsyncResult asyncResult)
213-
{
214-
if (asyncResult == null)
215-
{
216-
throw new ArgumentNullException(nameof(asyncResult));
217-
}
218-
219-
var task = (Task)asyncResult;
220-
task.GetAwaiter().GetResult();
221-
}
180+
=> TaskToApm.End(asyncResult);
222181

223182
public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
224183
{

src/Servers/HttpSys/src/RequestProcessing/RequestStream.cs

Lines changed: 2 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -190,111 +190,10 @@ internal void UpdateAfterRead(uint statusCode, uint dataRead)
190190
}
191191

192192
public override unsafe IAsyncResult BeginRead(byte[] buffer, int offset, int size, AsyncCallback? callback, object? state)
193-
{
194-
ValidateReadBuffer(buffer, offset, size);
195-
CheckSizeLimit();
196-
if (_closed)
197-
{
198-
RequestStreamAsyncResult result = new RequestStreamAsyncResult(this, state, callback);
199-
result.Complete(0);
200-
return result;
201-
}
202-
// TODO: Verbose log parameters
203-
204-
RequestStreamAsyncResult? asyncResult = null;
205-
206-
uint dataRead = 0;
207-
if (_dataChunkIndex != -1)
208-
{
209-
dataRead = _requestContext.Request.GetChunks(ref _dataChunkIndex, ref _dataChunkOffset, buffer, offset, size);
210-
211-
if (dataRead > 0)
212-
{
213-
asyncResult = new RequestStreamAsyncResult(this, state, callback, buffer, offset, 0);
214-
asyncResult.Complete((int)dataRead);
215-
return asyncResult;
216-
}
217-
}
218-
219-
uint statusCode = 0;
220-
221-
// the http.sys team recommends that we limit the size to 128kb
222-
if (size > MaxReadSize)
223-
{
224-
size = MaxReadSize;
225-
}
226-
227-
asyncResult = new RequestStreamAsyncResult(this, state, callback, buffer, offset, dataRead);
228-
uint bytesReturned;
229-
230-
try
231-
{
232-
uint flags = 0;
233-
234-
statusCode =
235-
HttpApi.HttpReceiveRequestEntityBody(
236-
RequestQueueHandle,
237-
RequestId,
238-
flags,
239-
asyncResult.PinnedBuffer,
240-
(uint)size,
241-
out bytesReturned,
242-
asyncResult.NativeOverlapped!);
243-
}
244-
catch (Exception e)
245-
{
246-
Log.ErrorWhenReadBegun(Logger, e);
247-
asyncResult.Dispose();
248-
throw;
249-
}
250-
251-
if (statusCode != UnsafeNclNativeMethods.ErrorCodes.ERROR_SUCCESS && statusCode != UnsafeNclNativeMethods.ErrorCodes.ERROR_IO_PENDING)
252-
{
253-
asyncResult.Dispose();
254-
if (statusCode == UnsafeNclNativeMethods.ErrorCodes.ERROR_HANDLE_EOF)
255-
{
256-
asyncResult = new RequestStreamAsyncResult(this, state, callback, dataRead);
257-
asyncResult.Complete((int)bytesReturned);
258-
}
259-
else
260-
{
261-
Exception exception = new IOException(string.Empty, new HttpSysException((int)statusCode));
262-
Log.ErrorWhenReadBegun(Logger, exception);
263-
Abort();
264-
throw exception;
265-
}
266-
}
267-
else if (statusCode == UnsafeNclNativeMethods.ErrorCodes.ERROR_SUCCESS &&
268-
HttpSysListener.SkipIOCPCallbackOnSuccess)
269-
{
270-
// IO operation completed synchronously - callback won't be called to signal completion.
271-
asyncResult.IOCompleted(statusCode, bytesReturned);
272-
}
273-
return asyncResult;
274-
}
193+
=> TaskToApm.Begin(ReadAsync(buffer, offset, size, CancellationToken.None), callback, state);
275194

276195
public override int EndRead(IAsyncResult asyncResult)
277-
{
278-
if (asyncResult == null)
279-
{
280-
throw new ArgumentNullException("asyncResult");
281-
}
282-
RequestStreamAsyncResult? castedAsyncResult = asyncResult as RequestStreamAsyncResult;
283-
if (castedAsyncResult == null || castedAsyncResult.RequestStream != this)
284-
{
285-
throw new ArgumentException(Resources.Exception_WrongIAsyncResult, "asyncResult");
286-
}
287-
if (castedAsyncResult.EndCalled)
288-
{
289-
throw new InvalidOperationException(Resources.Exception_EndCalledMultipleTimes);
290-
}
291-
castedAsyncResult.EndCalled = true;
292-
// wait & then check for errors
293-
// Throws on failure
294-
var dataRead = castedAsyncResult.Task.GetAwaiter().GetResult();
295-
// TODO: Verbose log #dataRead.
296-
return dataRead;
297-
}
196+
=> TaskToApm.End<int>(asyncResult);
298197

299198
public override unsafe Task<int> ReadAsync(byte[] buffer, int offset, int size, CancellationToken cancellationToken)
300199
{

src/Servers/HttpSys/src/RequestProcessing/RequestStreamAsyncResult.cs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,6 @@ internal Task<int> Task
7777
get { return _tcs.Task; }
7878
}
7979

80-
internal bool EndCalled { get; set; }
81-
8280
internal void IOCompleted(uint errorCode, uint numBytes)
8381
{
8482
IOCompleted(this, errorCode, numBytes);

src/Servers/HttpSys/src/ResponseStream.cs

Lines changed: 2 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -73,43 +73,9 @@ public override async Task WriteAsync(byte[] buffer, int offset, int count, Canc
7373
}
7474

7575
public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
76-
{
77-
return ToIAsyncResult(WriteAsync(buffer, offset, count), callback, state);
78-
}
76+
=> TaskToApm.Begin(WriteAsync(buffer, offset, count, CancellationToken.None), callback, state);
7977

8078
public override void EndWrite(IAsyncResult asyncResult)
81-
{
82-
if (asyncResult == null)
83-
{
84-
throw new ArgumentNullException(nameof(asyncResult));
85-
}
86-
((Task)asyncResult).GetAwaiter().GetResult();
87-
}
88-
89-
private static IAsyncResult ToIAsyncResult(Task task, AsyncCallback? callback, object? state)
90-
{
91-
var tcs = new TaskCompletionSource<int>(state);
92-
task.ContinueWith(t =>
93-
{
94-
if (t.IsFaulted)
95-
{
96-
tcs.TrySetException(t.Exception!.InnerExceptions);
97-
}
98-
else if (t.IsCanceled)
99-
{
100-
tcs.TrySetCanceled();
101-
}
102-
else
103-
{
104-
tcs.TrySetResult(0);
105-
}
106-
107-
if (callback != null)
108-
{
109-
callback(tcs.Task);
110-
}
111-
}, CancellationToken.None, TaskContinuationOptions.None, TaskScheduler.Default);
112-
return tcs.Task;
113-
}
79+
=> TaskToApm.End(asyncResult);
11480
}
11581
}

0 commit comments

Comments
 (0)