server and channel adjustments

4 files changed, 23 insertions, 17 deletions

diff --git a/src/csharp/Grpc.Core/Channel.cs b/src/csharp/Grpc.Core/Channel.cs
index 1803920662..f9925a8a76 100644
--- a/src/csharp/Grpc.Core/Channel.cs
+++ b/src/csharp/Grpc.Core/Channel.cs
@@ -127,6 +127,20 @@ namespace Grpc.Core
             }
         }
 
+        // cached handler for watch connectivity state
+        static readonly BatchCompletionDelegate WatchConnectivityStateHandler = (success, ctx, state) =>
+        {
+            var tcs = (TaskCompletionSource<object>) state;
+            if (success)
+            {
+                tcs.SetResult(null);
+            }
+            else
+            {
+                tcs.SetCanceled();
+            }
+        };
+
         /// <summary>
         /// Returned tasks completes once channel state has become different from 
         /// given lastObservedState. 
@@ -138,18 +152,8 @@ namespace Grpc.Core
                 "Shutdown is a terminal state. No further state changes can occur.");
             var tcs = new TaskCompletionSource<object>();
             var deadlineTimespec = deadline.HasValue ? Timespec.FromDateTime(deadline.Value) : Timespec.InfFuture;
-            var handler = new BatchCompletionDelegate((success, ctx) =>
-            {
-                if (success)
-                {
-                    tcs.SetResult(null);
-                }
-                else
-                {
-                    tcs.SetCanceled();
-                }
-            });
-            handle.WatchConnectivityState(lastObservedState, deadlineTimespec, completionQueue, handler);
+            // pass "tcs" as "state" for WatchConnectivityStateHandler.
+            handle.WatchConnectivityState(lastObservedState, deadlineTimespec, completionQueue, WatchConnectivityStateHandler, tcs);
             return tcs.Task;
         }
 
diff --git a/src/csharp/Grpc.Core/Internal/ChannelSafeHandle.cs b/src/csharp/Grpc.Core/Internal/ChannelSafeHandle.cs
index f826a17bad..1eeb0e3d97 100644
--- a/src/csharp/Grpc.Core/Internal/ChannelSafeHandle.cs
+++ b/src/csharp/Grpc.Core/Internal/ChannelSafeHandle.cs
@@ -64,10 +64,10 @@ namespace Grpc.Core.Internal
             return Native.grpcsharp_channel_check_connectivity_state(this, tryToConnect ? 1 : 0);
         }
 
-        public void WatchConnectivityState(ChannelState lastObservedState, Timespec deadline, CompletionQueueSafeHandle cq, BatchCompletionDelegate callback)
+        public void WatchConnectivityState(ChannelState lastObservedState, Timespec deadline, CompletionQueueSafeHandle cq, BatchCompletionDelegate callback, object callbackState)
         {
             var ctx = BatchContextSafeHandle.Create();
-            cq.CompletionRegistry.RegisterBatchCompletion(ctx, callback);
+            cq.CompletionRegistry.RegisterBatchCompletion(ctx, callback, callbackState);
             Native.grpcsharp_channel_watch_connectivity_state(this, lastObservedState, deadline, cq, ctx);
         }
 
diff --git a/src/csharp/Grpc.Core/Internal/ServerSafeHandle.cs b/src/csharp/Grpc.Core/Internal/ServerSafeHandle.cs
index 63000e9a22..a308890cde 100644
--- a/src/csharp/Grpc.Core/Internal/ServerSafeHandle.cs
+++ b/src/csharp/Grpc.Core/Internal/ServerSafeHandle.cs
@@ -59,13 +59,15 @@ namespace Grpc.Core.Internal
         {
             Native.grpcsharp_server_start(this);
         }
-    
+
         public void ShutdownAndNotify(BatchCompletionDelegate callback, CompletionQueueSafeHandle completionQueue)
         {
             using (completionQueue.NewScope())
             {
                 var ctx = BatchContextSafeHandle.Create();
-                completionQueue.CompletionRegistry.RegisterBatchCompletion(ctx, callback);
+                // TODO(jtattermusch): delegate allocation by caller can be avoided by utilizing the "state" object,
+                // but server shutdown isn't worth optimizing right now.
+                completionQueue.CompletionRegistry.RegisterBatchCompletion(ctx, callback, null);
                 Native.grpcsharp_server_shutdown_and_notify_callback(this, completionQueue, ctx);
             }
         }
diff --git a/src/csharp/Grpc.Core/Server.cs b/src/csharp/Grpc.Core/Server.cs
index 77ad876bdf..71c7f108f3 100644
--- a/src/csharp/Grpc.Core/Server.cs
+++ b/src/csharp/Grpc.Core/Server.cs
@@ -387,7 +387,7 @@ namespace Grpc.Core
         /// <summary>
         /// Handles native callback.
         /// </summary>
-        private void HandleServerShutdown(bool success, BatchContextSafeHandle ctx)
+        private void HandleServerShutdown(bool success, BatchContextSafeHandle ctx, object state)
         {
             shutdownTcs.SetResult(null);
         }