From 39775cf30fe2b29693652184e01a480a8596a89c Mon Sep 17 00:00:00 2001
From: mgravell <marc.gravell@gmail.com>
Date: Thu, 11 Jul 2019 13:07:50 +0100
Subject: [PATCH] remove delegate allocation and boxing from cancellation
 registrations

---
 src/csharp/Grpc.Core.Tests/CallCancellationTest.cs    | 11 +++++++++++
 src/csharp/Grpc.Core/Internal/AsyncCall.cs            | 10 +++-------
 src/csharp/Grpc.Core/Internal/AsyncCallBase.cs        |  8 ++++++++
 src/csharp/Grpc.Core/Internal/ClientResponseStream.cs |  3 +--
 src/csharp/Grpc.Core/Internal/ServerRequestStream.cs  |  4 +---
 5 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/csharp/Grpc.Core.Tests/CallCancellationTest.cs b/src/csharp/Grpc.Core.Tests/CallCancellationTest.cs
index e040f52380a..753b8b227b1 100644
--- a/src/csharp/Grpc.Core.Tests/CallCancellationTest.cs
+++ b/src/csharp/Grpc.Core.Tests/CallCancellationTest.cs
@@ -178,5 +178,16 @@ namespace Grpc.Core.Tests
                 Assert.AreEqual(StatusCode.Cancelled, ex.Status.StatusCode);
             }
         }
+
+        [Test]
+        public void CanDisposeDefaultCancellationRegistration()
+        {
+            // prove that we're fine to dispose default CancellationTokenRegistration
+            // values without boxing them to IDisposable for a null-check
+            var obj = default(CancellationTokenRegistration);
+            obj.Dispose();
+
+            using (obj) {}
+        }
     }
 }
diff --git a/src/csharp/Grpc.Core/Internal/AsyncCall.cs b/src/csharp/Grpc.Core/Internal/AsyncCall.cs
index a1c688140d1..ce3508d398e 100644
--- a/src/csharp/Grpc.Core/Internal/AsyncCall.cs
+++ b/src/csharp/Grpc.Core/Internal/AsyncCall.cs
@@ -38,7 +38,7 @@ namespace Grpc.Core.Internal
         bool registeredWithChannel;
 
         // Dispose of to de-register cancellation token registration
-        IDisposable cancellationTokenRegistration;
+        CancellationTokenRegistration cancellationTokenRegistration;
 
         // Completion of a pending unary response if not null.
         TaskCompletionSource<TResponse> unaryResponseTcs;
@@ -409,7 +409,7 @@ namespace Grpc.Core.Internal
             // deadlock.
             // See https://github.com/grpc/grpc/issues/14777
             // See https://github.com/dotnet/corefx/issues/14903
-            cancellationTokenRegistration?.Dispose();
+            cancellationTokenRegistration.Dispose();
         }
 
         protected override bool IsClient
@@ -509,11 +509,7 @@ namespace Grpc.Core.Internal
         // Make sure that once cancellationToken for this call is cancelled, Cancel() will be called.
         private void RegisterCancellationCallback()
         {
-            var token = details.Options.CancellationToken;
-            if (token.CanBeCanceled)
-            {
-                cancellationTokenRegistration = token.Register(() => this.Cancel());
-            }
+            cancellationTokenRegistration = RegisterCancellationCallbackForToken(details.Options.CancellationToken);
         }
 
         /// <summary>
diff --git a/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs b/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs
index 0b99aaf3e21..ac1f9066516 100644
--- a/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs
+++ b/src/csharp/Grpc.Core/Internal/AsyncCallBase.cs
@@ -394,5 +394,13 @@ namespace Grpc.Core.Internal
         {
             HandleReadFinished(success, receivedMessageReader);
         }
+
+        internal CancellationTokenRegistration RegisterCancellationCallbackForToken(CancellationToken cancellationToken)
+        {
+            if (cancellationToken.CanBeCanceled) return cancellationToken.Register(CancelCallFromToken, this);
+            return default(CancellationTokenRegistration);
+        }
+
+        private static readonly Action<object> CancelCallFromToken = state => ((AsyncCallBase<TWrite, TRead>)state).Cancel();
     }
 }
diff --git a/src/csharp/Grpc.Core/Internal/ClientResponseStream.cs b/src/csharp/Grpc.Core/Internal/ClientResponseStream.cs
index ab649ee7662..ff73b26dff4 100644
--- a/src/csharp/Grpc.Core/Internal/ClientResponseStream.cs
+++ b/src/csharp/Grpc.Core/Internal/ClientResponseStream.cs
@@ -49,8 +49,7 @@ namespace Grpc.Core.Internal
 
         public async Task<bool> MoveNext(CancellationToken token)
         {
-            var cancellationTokenRegistration = token.CanBeCanceled ? token.Register(() => call.Cancel()) : (IDisposable) null;
-            using (cancellationTokenRegistration)
+            using (call.RegisterCancellationCallbackForToken(token))
             {
                 var result = await call.ReadMessageAsync().ConfigureAwait(false);
                 this.current = result;
diff --git a/src/csharp/Grpc.Core/Internal/ServerRequestStream.cs b/src/csharp/Grpc.Core/Internal/ServerRequestStream.cs
index 058dddb7ebb..f3655139e47 100644
--- a/src/csharp/Grpc.Core/Internal/ServerRequestStream.cs
+++ b/src/csharp/Grpc.Core/Internal/ServerRequestStream.cs
@@ -49,9 +49,7 @@ namespace Grpc.Core.Internal
 
         public async Task<bool> MoveNext(CancellationToken token)
         {
-            
-            var cancellationTokenRegistration = token.CanBeCanceled ? token.Register(() => call.Cancel()) : (IDisposable) null;
-            using (cancellationTokenRegistration)
+            using (call.RegisterCancellationCallbackForToken(token))
             {
                 var result = await call.ReadMessageAsync().ConfigureAwait(false);
                 this.current = result;