Fix potential deadlock

Fix a potential deadlock where the server never stops by checking the result of `WaitToReadAsync()` and by passing the `CancellationToken` to `ReadAsync()`.
martincostello · May 19, 2024 · abf4ba8 · abf4ba8
1 parent 68ae617
commit abf4ba8
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/AwsLambdaTestServer/RuntimeHandler.cs b/src/AwsLambdaTestServer/RuntimeHandler.cs
@@ -134,8 +134,12 @@ internal async Task HandleNextAsync(HttpContext httpContext)
             using var cts = CancellationTokenSource.CreateLinkedTokenSource(httpContext.RequestAborted, _cancellationToken);
 
             // Wait until there is a request to process
-            await _requests.Reader.WaitToReadAsync(cts.Token).ConfigureAwait(false);
-            request = await _requests.Reader.ReadAsync().ConfigureAwait(false);
+            if (!await _requests.Reader.WaitToReadAsync(cts.Token).ConfigureAwait(false))
+            {
+                cts.Token.ThrowIfCancellationRequested();
+            }
+
+            request = await _requests.Reader.ReadAsync(cts.Token).ConfigureAwait(false);
         }
         catch (Exception ex) when (ex is OperationCanceledException or ChannelClosedException)
         {