Skip to content

Commit e36e121

Browse files
JamesNKJamesNK
authored
Fix race condition between deadline and call finishing (#1082)
1 parent 96c06f4 commit e36e121

File tree

6 files changed

+79
-117
lines changed

6 files changed

+79
-117
lines changed

‎src/Grpc.AspNetCore.Server/Internal/CallHandlers/ClientStreamingServerCallHandler.cs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,11 +53,17 @@ protected override async Task HandleCallAsyncCore(HttpContext httpContext, HttpC
5353
throw new RpcException(new Status(StatusCode.Cancelled, "No message returned from method."));
5454
}
5555

56-
if (serverCallContext.DeadlineManager != null && serverCallContext.DeadlineManager.CallComplete)
56+
if (serverCallContext.DeadlineManager != null && serverCallContext.DeadlineManager.CancellationToken.IsCancellationRequested)
5757
{
58-
// The deadline has been exceeded and the call has been completed.
58+
// The cancellation token has been raised. Ensure that any DeadlineManager tasks have
59+
// been completed before continuing.
60+
await serverCallContext.DeadlineManager.CancellationProcessedTask;
61+
5962
// There is no point trying to write to the response because it has been finished.
60-
return;
63+
if (serverCallContext.DeadlineManager.CallComplete)
64+
{
65+
return;
66+
}
6167
}
6268

6369
var responseBodyWriter = httpContext.Response.BodyWriter;

‎src/Grpc.AspNetCore.Server/Internal/CallHandlers/UnaryServerCallHandler.cs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,17 @@ protected override async Task HandleCallAsyncCore(HttpContext httpContext, HttpC
5454
throw new RpcException(new Status(StatusCode.Cancelled, "No message returned from method."));
5555
}
5656

57-
if (serverCallContext.DeadlineManager != null && serverCallContext.DeadlineManager.CallComplete)
57+
if (serverCallContext.DeadlineManager != null && serverCallContext.DeadlineManager.CancellationToken.IsCancellationRequested)
5858
{
59-
// The deadline has been exceeded and the call has been completed.
59+
// The cancellation token has been raised. Ensure that any DeadlineManager tasks have
60+
// been completed before continuing.
61+
await serverCallContext.DeadlineManager.CancellationProcessedTask;
62+
6063
// There is no point trying to write to the response because it has been finished.
61-
return;
64+
if (serverCallContext.DeadlineManager.CallComplete)
65+
{
66+
return;
67+
}
6268
}
6369

6470
var responseBodyWriter = httpContext.Response.BodyWriter;

‎src/Grpc.AspNetCore.Server/Internal/Deadline/ServerCallDeadlineManager.cs‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ internal abstract class ServerCallDeadlineManager : IAsyncDisposable
2929
private Task? _deadlineExceededTask;
3030
private CancellationTokenRegistration _deadlineExceededRegistration;
3131
private CancellationTokenRegistration _requestAbortedRegistration;
32+
private readonly TaskCompletionSource<object?> _cancellationProcessedTcs;
3233
protected HttpContextServerCallContext ServerCallContext { get; }
3334

3435
internal DateTime Deadline { get; private set; }
@@ -39,6 +40,12 @@ internal abstract class ServerCallDeadlineManager : IAsyncDisposable
3940

4041
public CancellationToken CancellationToken => _deadlineCts!.Token;
4142

43+
// Task to wait for when a call is being completed to ensure that registered deadline cancellation
44+
// events have finished processing.
45+
// - Avoids a race condition between deadline being raised and the call completing.
46+
// - Required because OCE error thrown from token happens before registered events.
47+
public Task CancellationProcessedTask => _cancellationProcessedTcs.Task;
48+
4249
public static ServerCallDeadlineManager Create(
4350
HttpContextServerCallContext serverCallContext,
4451
ISystemClock clock,
@@ -68,6 +75,7 @@ protected ServerCallDeadlineManager(HttpContextServerCallContext serverCallConte
6875
ServerCallContext = serverCallContext;
6976

7077
Lock = new SemaphoreSlim(1, 1);
78+
_cancellationProcessedTcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
7179
}
7280

7381
public void Initialize(ISystemClock clock, TimeSpan timeout, CancellationToken requestAborted)
@@ -76,6 +84,7 @@ public void Initialize(ISystemClock clock, TimeSpan timeout, CancellationToken r
7684

7785
_deadlineCts = CreateCancellationTokenSource(timeout, clock);
7886
_deadlineExceededRegistration = _deadlineCts.Token.Register(DeadlineExceeded);
87+
7988
_requestAbortedRegistration = requestAborted.Register(() =>
8089
{
8190
// Call is complete if the request has aborted
@@ -98,29 +107,36 @@ protected void DeadlineExceeded()
98107

99108
private async Task DeadlineExceededAsync()
100109
{
101-
if (!CanExceedDeadline())
102-
{
103-
return;
104-
}
105-
106-
Debug.Assert(Lock != null, "Lock has not been created.");
107-
108-
await Lock.WaitAsync();
109-
110110
try
111111
{
112-
// Double check after lock is aquired
113112
if (!CanExceedDeadline())
114113
{
115114
return;
116115
}
117116

118-
await ServerCallContext.DeadlineExceededAsync();
119-
CallComplete = true;
117+
Debug.Assert(Lock != null, "Lock has not been created.");
118+
119+
await Lock.WaitAsync();
120+
121+
try
122+
{
123+
// Double check after lock is aquired
124+
if (!CanExceedDeadline())
125+
{
126+
return;
127+
}
128+
129+
await ServerCallContext.DeadlineExceededAsync();
130+
CallComplete = true;
131+
}
132+
finally
133+
{
134+
Lock.Release();
135+
}
120136
}
121137
finally
122138
{
123-
Lock.Release();
139+
_cancellationProcessedTcs.TrySetResult(null);
124140
}
125141
}
126142

‎src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,13 @@ private async Task ProcessHandlerErrorAsyncCore(Exception ex, string method)
153153
{
154154
Debug.Assert(DeadlineManager != null, "Deadline manager should have been created.");
155155

156+
if (DeadlineManager.CancellationToken.IsCancellationRequested)
157+
{
158+
// The cancellation token has been raised. Ensure that any DeadlineManager tasks have
159+
// been completed before continuing.
160+
await DeadlineManager.CancellationProcessedTask;
161+
}
162+
156163
await DeadlineManager.Lock.WaitAsync();
157164

158165
try

‎test/FunctionalTests/Server/DeadlineTests.cs‎

Lines changed: 12 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -184,79 +184,21 @@ static async Task<HelloReply> ThrowErrorWithinDeadline(HelloRequest request, Ser
184184
Assert.True(HasLog(LogLevel.Trace, "DeadlineStopped", "Request deadline stopped."));
185185
}
186186

187-
[Test]
188-
public async Task UnaryMethodErrorAfterExceedDeadline()
187+
[TestCase(true)]
188+
[TestCase(false)]
189+
public async Task UnaryMethodDeadlineExceeded(bool throwErrorOnCancellation)
189190
{
190-
static async Task<HelloReply> ThrowErrorExceedDeadline(HelloRequest request, ServerCallContext context)
191-
{
192-
while (!context.CancellationToken.IsCancellationRequested)
193-
{
194-
await Task.Delay(50);
195-
}
196-
197-
throw new InvalidOperationException("An error.");
198-
}
199-
200-
// Arrange
201-
SetExpectedErrorsFilter(writeContext =>
191+
async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCallContext context)
202192
{
203-
if (writeContext.LoggerName == TestConstants.ServerCallHandlerTestName)
193+
try
204194
{
205-
// Deadline happened before write
206-
if (writeContext.EventId.Name == "ErrorExecutingServiceMethod" &&
207-
writeContext.State.ToString() == "Error when executing service method 'ThrowErrorExceedDeadline'." &&
208-
writeContext.Exception!.Message == "An error.")
209-
{
210-
return true;
211-
}
195+
await Task.Delay(1000, context.CancellationToken);
212196
}
213-
214-
return false;
215-
});
216-
217-
var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>(ThrowErrorExceedDeadline, nameof(ThrowErrorExceedDeadline));
218-
219-
var requestMessage = new HelloRequest
220-
{
221-
Name = "World"
222-
};
223-
224-
var requestStream = new MemoryStream();
225-
MessageHelpers.WriteMessage(requestStream, requestMessage);
226-
227-
var httpRequest = GrpcHttpHelper.Create(method.FullName);
228-
httpRequest.Headers.Add(GrpcProtocolConstants.TimeoutHeader, "200m");
229-
httpRequest.Content = new GrpcStreamContent(requestStream);
230-
231-
// Act
232-
var response = await Fixture.Client.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead).DefaultTimeout();
233-
234-
// Assert
235-
response.AssertIsSuccessfulGrpcRequest();
236-
response.AssertTrailerStatus(StatusCode.DeadlineExceeded, "Deadline Exceeded");
237-
238-
Assert.True(HasLog(LogLevel.Debug, "DeadlineExceeded", "Request with timeout of 00:00:00.2000000 has exceeded its deadline."));
239-
240-
// Ensure follow up error is logged.
241-
await TestHelpers.AssertIsTrueRetryAsync(
242-
() => HasLog(LogLevel.Error, "ErrorExecutingServiceMethod", "Error when executing service method 'ThrowErrorExceedDeadline'."),
243-
"Missing follow up error log.").DefaultTimeout();
244-
245-
await TestHelpers.AssertIsTrueRetryAsync(
246-
() => HasLog(LogLevel.Trace, "DeadlineStopped", "Request deadline stopped."),
247-
"Missing deadline stopped log.").DefaultTimeout();
248-
}
249-
250-
[Test]
251-
public async Task UnaryMethodDeadlineExceeded()
252-
{
253-
static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCallContext context)
254-
{
255-
while (!context.CancellationToken.IsCancellationRequested)
197+
catch (OperationCanceledException) when (!throwErrorOnCancellation)
256198
{
257-
await Task.Delay(50);
199+
// nom nom nom
258200
}
259-
201+
260202
return new HelloReply();
261203
}
262204

@@ -265,17 +207,8 @@ static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCall
265207
{
266208
if (writeContext.LoggerName == TestConstants.ServerCallHandlerTestName)
267209
{
268-
// Deadline happened before write
269210
if (writeContext.EventId.Name == "ErrorExecutingServiceMethod" &&
270-
writeContext.State.ToString() == "Error when executing service method 'WriteUntilError'." &&
271-
writeContext.Exception!.Message == "Cannot write message after request is complete.")
272-
{
273-
return true;
274-
}
275-
276-
// Deadline happened during write (error raised from pipeline writer)
277-
if (writeContext.Exception is InvalidOperationException &&
278-
writeContext.Exception.Message == "Writing is not allowed after writer was completed.")
211+
writeContext.State.ToString() == "Error when executing service method 'WaitUntilDeadline-True'.")
279212
{
280213
return true;
281214
}
@@ -284,7 +217,7 @@ static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCall
284217
return false;
285218
});
286219

287-
var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>(WaitUntilDeadline, nameof(WaitUntilDeadline));
220+
var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>(WaitUntilDeadline, $"{nameof(WaitUntilDeadline)}-{throwErrorOnCancellation}");
288221

289222
var requestMessage = new HelloRequest
290223
{
@@ -335,8 +268,7 @@ static async Task WriteUntilError(HelloRequest request, IServerStreamWriter<Hell
335268
{
336269
// Deadline happened before write
337270
if (writeContext.EventId.Name == "ErrorExecutingServiceMethod" &&
338-
writeContext.State.ToString() == "Error when executing service method 'WriteUntilError'." &&
339-
writeContext.Exception!.Message == "Cannot write message after request is complete.")
271+
writeContext.State.ToString() == "Error when executing service method 'WriteUntilError'.")
340272
{
341273
return true;
342274
}

‎test/FunctionalTests/Web/Server/DeadlineTests.cs‎

Lines changed: 13 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -44,17 +44,20 @@ public DeadlineTests(GrpcTestMode grpcTestMode, TestServerEndpointName endpointN
4444
{
4545
}
4646

47-
[Test]
48-
public async Task UnaryMethodDeadlineExceeded()
47+
[TestCase(true)]
48+
[TestCase(false)]
49+
public async Task UnaryMethodDeadlineExceeded(bool throwErrorOnCancellation)
4950
{
50-
static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCallContext context)
51+
async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCallContext context)
5152
{
52-
while (!context.CancellationToken.IsCancellationRequested)
53+
try
5354
{
54-
await Task.Delay(10);
55+
await Task.Delay(1000, context.CancellationToken);
56+
}
57+
catch (OperationCanceledException) when (!throwErrorOnCancellation)
58+
{
59+
// nom nom nom
5560
}
56-
57-
await Task.Delay(50);
5861

5962
return new HelloReply();
6063
}
@@ -65,15 +68,7 @@ static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCall
6568
{
6669
// Deadline happened before write
6770
if (writeContext.EventId.Name == "ErrorExecutingServiceMethod" &&
68-
writeContext.State.ToString() == "Error when executing service method 'WriteUntilError'." &&
69-
writeContext.Exception!.Message == "Cannot write message after request is complete.")
70-
{
71-
return true;
72-
}
73-
74-
// Deadline happened during write (error raised from pipeline writer)
75-
if (writeContext.Exception is InvalidOperationException &&
76-
writeContext.Exception.Message == "Writing is not allowed after writer was completed.")
71+
writeContext.State.ToString() == "Error when executing service method 'WaitUntilDeadline-True'.")
7772
{
7873
return true;
7974
}
@@ -82,7 +77,7 @@ static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCall
8277
return false;
8378
});
8479

85-
var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>(WaitUntilDeadline, nameof(WaitUntilDeadline));
80+
var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>(WaitUntilDeadline, $"{nameof(WaitUntilDeadline)}-{throwErrorOnCancellation}");
8681

8782
var grpcWebClient = CreateGrpcWebClient();
8883

@@ -95,7 +90,7 @@ static async Task<HelloReply> WaitUntilDeadline(HelloRequest request, ServerCall
9590
MessageHelpers.WriteMessage(requestStream, requestMessage);
9691

9792
var httpRequest = GrpcHttpHelper.Create(method.FullName);
98-
httpRequest.Headers.Add(GrpcProtocolConstants.TimeoutHeader, "50m");
93+
httpRequest.Headers.Add(GrpcProtocolConstants.TimeoutHeader, "100m");
9994
httpRequest.Content = new GrpcStreamContent(requestStream);
10095

10196
// Act

0 commit comments

Comments
 (0)