From 1159126138d752f21c8a02b74002d83aec5f2c88 Mon Sep 17 00:00:00 2001
From: James Newton-King <james@newtonking.com>
Date: Tue, 25 Feb 2020 10:03:16 +1300
Subject: [PATCH] Avoid CancellationToken related allocations (#787)

---
 .../Internal/Deadline/ServerCallDeadlineManager.cs   | 11 ++++++-----
 .../Internal/HttpContextServerCallContext.cs         | 12 ++++++------
 .../HttpContextServerCallContextTests.cs             |  4 ++--
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/src/Grpc.AspNetCore.Server/Internal/Deadline/ServerCallDeadlineManager.cs b/src/Grpc.AspNetCore.Server/Internal/Deadline/ServerCallDeadlineManager.cs
index 07626aef7..5a84a268b 100644
--- a/src/Grpc.AspNetCore.Server/Internal/Deadline/ServerCallDeadlineManager.cs
+++ b/src/Grpc.AspNetCore.Server/Internal/Deadline/ServerCallDeadlineManager.cs
@@ -35,7 +35,7 @@ internal abstract class ServerCallDeadlineManager : IAsyncDisposable
         // Lock is to ensure deadline doesn't execute as call is completing
         internal SemaphoreSlim Lock { get; }
         // Internal for testing
-        internal bool _callComplete;
+        public bool CallComplete { get; private set; }
 
         public CancellationToken CancellationToken => _deadlineCts!.Token;
 
@@ -79,16 +79,16 @@ public void Initialize(ISystemClock clock, TimeSpan timeout, CancellationToken r
             _requestAbortedRegistration = requestAborted.Register(() =>
             {
                 // Call is complete if the request has aborted
-                _callComplete = true;
+                CallComplete = true;
                 _deadlineCts?.Cancel();
             });
         }
 
         protected abstract CancellationTokenSource CreateCancellationTokenSource(TimeSpan timeout, ISystemClock clock);
 
-        public void SetCallComplete()
+        public void SetCallEnded()
         {
-            _callComplete = true;
+            CallComplete = true;
         }
 
         protected void DeadlineExceeded()
@@ -116,6 +116,7 @@ private async Task DeadlineExceededAsync()
                 }
 
                 await ServerCallContext.DeadlineExceededAsync();
+                CallComplete = true;
             }
             finally
             {
@@ -127,7 +128,7 @@ private bool CanExceedDeadline()
         {
             // Deadline callback could be raised by the CTS after call has been completed (either successfully, with error, or aborted)
             // but before deadline exceeded registration has been disposed
-            return !_callComplete;
+            return !CallComplete;
         }
 
         public ValueTask DisposeAsync()
diff --git a/src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs b/src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs
index 42296cb84..f91892221 100644
--- a/src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs
+++ b/src/Grpc.AspNetCore.Server/Internal/HttpContextServerCallContext.cs
@@ -180,14 +180,14 @@ private void ProcessHandlerError(Exception ex, string method)
             }
 
             // Don't update trailers if request has exceeded deadline/aborted
-            if (!CancellationToken.IsCancellationRequested)
+            if (DeadlineManager == null || !DeadlineManager.CallComplete)
             {
                 HttpContext.Response.ConsolidateTrailers(this);
             }
 
-            LogCallEnd();
+            DeadlineManager?.SetCallEnded();
 
-            DeadlineManager?.SetCallComplete();
+            LogCallEnd();
         }
 
         // If there is a deadline then we need to have our own cancellation token.
@@ -266,14 +266,14 @@ private async Task EndCallAsyncCore(Task lockTask)
         private void EndCallCore()
         {
             // Don't set trailers if deadline exceeded or request aborted
-            if (!CancellationToken.IsCancellationRequested)
+            if (DeadlineManager == null || !DeadlineManager.CallComplete)
             {
                 HttpContext.Response.ConsolidateTrailers(this);
             }
 
-            LogCallEnd();
+            DeadlineManager?.SetCallEnded();
 
-            DeadlineManager?.SetCallComplete();
+            LogCallEnd();
         }
 
         private void LogCallEnd()
diff --git a/test/Grpc.AspNetCore.Server.Tests/HttpContextServerCallContextTests.cs b/test/Grpc.AspNetCore.Server.Tests/HttpContextServerCallContextTests.cs
index 9332d6f24..9c3e544d7 100644
--- a/test/Grpc.AspNetCore.Server.Tests/HttpContextServerCallContextTests.cs
+++ b/test/Grpc.AspNetCore.Server.Tests/HttpContextServerCallContextTests.cs
@@ -608,7 +608,7 @@ private async Task LongRunningDeadline_WaitsUntilDeadlineIsFinished(string metho
                 Assert.Fail($"{methodName} did not wait on lock taken by deadline cancellation.");
             }
 
-            Assert.IsFalse(serverCallContext.DeadlineManager!._callComplete);
+            Assert.IsFalse(serverCallContext.DeadlineManager!.CallComplete);
 
             // Wait for dispose to finish
             syncPoint.Continue();
@@ -616,7 +616,7 @@ private async Task LongRunningDeadline_WaitsUntilDeadlineIsFinished(string metho
 
             Assert.AreEqual(GrpcProtocolConstants.ResetStreamNoError, httpResetFeature.ErrorCode);
 
-            Assert.IsTrue(serverCallContext.DeadlineManager!._callComplete);
+            Assert.IsTrue(serverCallContext.DeadlineManager!.CallComplete);
         }
 
         [Test]