1 files changed, 50 insertions, 4 deletions

diff --git a/src/csharp/Grpc.Core/Internal/BatchContextSafeHandle.cs b/src/csharp/Grpc.Core/Internal/BatchContextSafeHandle.cs
index cd5e3d8911..1e6f1fba37 100644
--- a/src/csharp/Grpc.Core/Internal/BatchContextSafeHandle.cs
+++ b/src/csharp/Grpc.Core/Internal/BatchContextSafeHandle.cs
@@ -20,15 +20,25 @@ using System;
 using System.Runtime.InteropServices;
 using System.Text;
 using Grpc.Core;
+using Grpc.Core.Logging;
+using Grpc.Core.Utils;
 
 namespace Grpc.Core.Internal
 {
+    internal interface IOpCompletionCallback
+    {
+        void OnComplete(bool success);
+    }
+
     /// <summary>
     /// grpcsharp_batch_context
     /// </summary>
-    internal class BatchContextSafeHandle : SafeHandleZeroIsInvalid
+    internal class BatchContextSafeHandle : SafeHandleZeroIsInvalid, IOpCompletionCallback
     {
         static readonly NativeMethods Native = NativeMethods.Get();
+        static readonly ILogger Logger = GrpcEnvironment.Logger.ForType<BatchContextSafeHandle>();
+
+        CompletionCallbackData completionCallbackData;
 
         private BatchContextSafeHandle()
         {
@@ -47,19 +57,26 @@ namespace Grpc.Core.Internal
             }
         }
 
+        public void SetCompletionCallback(BatchCompletionDelegate callback, object state)
+        {
+            GrpcPreconditions.CheckState(completionCallbackData.Callback == null);
+            GrpcPreconditions.CheckNotNull(callback, nameof(callback));
+            completionCallbackData = new CompletionCallbackData(callback, state);
+        }
+
         // Gets data of recv_initial_metadata completion.
         public Metadata GetReceivedInitialMetadata()
         {
             IntPtr metadataArrayPtr = Native.grpcsharp_batch_context_recv_initial_metadata(this);
             return MetadataArraySafeHandle.ReadMetadataFromPtrUnsafe(metadataArrayPtr);
         }
-            
+
         // Gets data of recv_status_on_client completion.
         public ClientSideStatus GetReceivedStatusOnClient()
         {
             UIntPtr detailsLength;
             IntPtr detailsPtr = Native.grpcsharp_batch_context_recv_status_on_client_details(this, out detailsLength);
-            string details = MarshalUtils.PtrToStringUTF8(detailsPtr, (int) detailsLength.ToUInt32());
+            string details = MarshalUtils.PtrToStringUTF8(detailsPtr, (int)detailsLength.ToUInt32());
             var status = new Status(Native.grpcsharp_batch_context_recv_status_on_client_status(this), details);
 
             IntPtr metadataArrayPtr = Native.grpcsharp_batch_context_recv_status_on_client_trailing_metadata(this);
@@ -86,11 +103,40 @@ namespace Grpc.Core.Internal
         {
             return Native.grpcsharp_batch_context_recv_close_on_server_cancelled(this) != 0;
         }
-            
+
         protected override bool ReleaseHandle()
         {
             Native.grpcsharp_batch_context_destroy(handle);
             return true;
         }
+
+        void IOpCompletionCallback.OnComplete(bool success)
+        {
+            try
+            {
+                completionCallbackData.Callback(success, this, completionCallbackData.State);
+            }
+            catch (Exception e)
+            {
+                Logger.Error(e, "Exception occured while invoking batch completion delegate.");
+            }
+            finally
+            {
+                completionCallbackData = default(CompletionCallbackData);
+                Dispose();
+            }
+        }
+
+        struct CompletionCallbackData
+        {
+            public CompletionCallbackData(BatchCompletionDelegate callback, object state)
+            {
+                this.Callback = callback;
+                this.State = state;
+            }
+
+            public BatchCompletionDelegate Callback { get; }
+            public object State { get; }
+        }
     }
 }