1 files changed, 30 insertions, 55 deletions

diff --git a/src/csharp/Grpc.Core/Internal/CompletionRegistry.cs b/src/csharp/Grpc.Core/Internal/CompletionRegistry.cs
index 1102c8d14f..b68655b33c 100644
--- a/src/csharp/Grpc.Core/Internal/CompletionRegistry.cs
+++ b/src/csharp/Grpc.Core/Internal/CompletionRegistry.cs
@@ -19,15 +19,15 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Runtime.InteropServices;
+using System.Threading;
 using Grpc.Core.Logging;
 using Grpc.Core.Utils;
 
 namespace Grpc.Core.Internal
 {
-    internal delegate void OpCompletionDelegate(bool success);
-
-    internal delegate void BatchCompletionDelegate(bool success, BatchContextSafeHandle ctx);
+    internal delegate void BatchCompletionDelegate(bool success, BatchContextSafeHandle ctx, object state);
 
     internal delegate void RequestCallCompletionDelegate(bool success, RequestCallContextSafeHandle ctx);
 
@@ -36,8 +36,8 @@ namespace Grpc.Core.Internal
         static readonly ILogger Logger = GrpcEnvironment.Logger.ForType<CompletionRegistry>();
 
         readonly GrpcEnvironment environment;
-        readonly Dictionary<IntPtr, OpCompletionDelegate> dict = new Dictionary<IntPtr, OpCompletionDelegate>(new IntPtrComparer());
-        readonly object myLock = new object();
+        readonly Dictionary<IntPtr, IOpCompletionCallback> dict = new Dictionary<IntPtr, IOpCompletionCallback>(new IntPtrComparer());
+        SpinLock spinLock = new SpinLock(Debugger.IsAttached);
         IntPtr lastRegisteredKey;  // only for testing
 
         public CompletionRegistry(GrpcEnvironment environment)
@@ -45,38 +45,51 @@ namespace Grpc.Core.Internal
             this.environment = environment;
         }
 
-        public void Register(IntPtr key, OpCompletionDelegate callback)
+        public void Register(IntPtr key, IOpCompletionCallback callback)
         {
             environment.DebugStats.PendingBatchCompletions.Increment();
-            lock (myLock)
+
+            bool lockTaken = false;
+            try
             {
+                spinLock.Enter(ref lockTaken);
+
                 dict.Add(key, callback);
                 this.lastRegisteredKey = key;
             }
+            finally
+            {
+                if (lockTaken) spinLock.Exit();
+            }
         }
 
-        public void RegisterBatchCompletion(BatchContextSafeHandle ctx, BatchCompletionDelegate callback)
+        public void RegisterBatchCompletion(BatchContextSafeHandle ctx, BatchCompletionDelegate callback, object state)
         {
-            // TODO(jtattermusch): get rid of new delegate creation here
-            OpCompletionDelegate opCallback = ((success) => HandleBatchCompletion(success, ctx, callback));
-            Register(ctx.Handle, opCallback);
+            ctx.SetCompletionCallback(callback, state);
+            Register(ctx.Handle, ctx);
         }
 
         public void RegisterRequestCallCompletion(RequestCallContextSafeHandle ctx, RequestCallCompletionDelegate callback)
         {
-            // TODO(jtattermusch): get rid of new delegate creation here
-            OpCompletionDelegate opCallback = ((success) => HandleRequestCallCompletion(success, ctx, callback));
-            Register(ctx.Handle, opCallback);
+            ctx.CompletionCallback = callback;
+            Register(ctx.Handle, ctx);
         }
 
-        public OpCompletionDelegate Extract(IntPtr key)
+        public IOpCompletionCallback Extract(IntPtr key)
         {
-            OpCompletionDelegate value = null;
-            lock (myLock)
+            IOpCompletionCallback value = null;
+            bool lockTaken = false;
+            try
             {
+                spinLock.Enter(ref lockTaken);
+
                 value = dict[key];
                 dict.Remove(key);
             }
+            finally
+            {
+                if (lockTaken) spinLock.Exit();
+            }
             environment.DebugStats.PendingBatchCompletions.Decrement();
             return value;
         }
@@ -89,44 +102,6 @@ namespace Grpc.Core.Internal
             get { return this.lastRegisteredKey; }
         }
 
-        private static void HandleBatchCompletion(bool success, BatchContextSafeHandle ctx, BatchCompletionDelegate callback)
-        {
-            try
-            {
-                callback(success, ctx);
-            }
-            catch (Exception e)
-            {
-                Logger.Error(e, "Exception occured while invoking batch completion delegate.");
-            }
-            finally
-            {
-                if (ctx != null)
-                {
-                    ctx.Dispose();
-                }
-            }
-        }
-
-        private static void HandleRequestCallCompletion(bool success, RequestCallContextSafeHandle ctx, RequestCallCompletionDelegate callback)
-        {
-            try
-            {
-                callback(success, ctx);
-            }
-            catch (Exception e)
-            {
-                Logger.Error(e, "Exception occured while invoking request call completion delegate.");
-            }
-            finally
-            {
-                if (ctx != null)
-                {
-                    ctx.Dispose();
-                }
-            }
-        }
-
         /// <summary>
         /// IntPtr doesn't implement <c>IEquatable{IntPtr}</c> so we need to use custom comparer to avoid boxing.
         /// </summary>