use just one response batch for unary response serverside calls

2 files changed, 25 insertions, 18 deletions

diff --git a/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs b/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs
index c5d900b134..b1566b44a7 100644
--- a/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs
+++ b/src/csharp/Grpc.Core/Internal/AsyncCallServer.cs
@@ -139,8 +139,11 @@ namespace Grpc.Core.Internal
         /// Sends call result status, indicating we are done with writes.
         /// Sending a status different from StatusCode.OK will also implicitly cancel the call.
         /// </summary>
-        public Task SendStatusFromServerAsync(Status status, Metadata trailers)
+        public Task SendStatusFromServerAsync(Status status, Metadata trailers, Tuple<TResponse, WriteFlags> optionalWrite)
         {
+            byte[] payload = optionalWrite != null ? UnsafeSerialize(optionalWrite.Item1) : null;
+            var writeFlags = optionalWrite != null ? optionalWrite.Item2 : default(WriteFlags);
+
             lock (myLock)
             {
                 GrpcPreconditions.CheckState(started);
@@ -150,11 +153,15 @@ namespace Grpc.Core.Internal
                 using (var metadataArray = MetadataArraySafeHandle.Create(trailers))
                 {
                     call.StartSendStatusFromServer(HandleSendStatusFromServerFinished, status, metadataArray, !initialMetadataSent,
-                        null, new WriteFlags());
+                        payload, writeFlags);
                 }
                 halfcloseRequested = true;
                 initialMetadataSent = true;
                 sendStatusFromServerTcs = new TaskCompletionSource<object>();
+                if (optionalWrite != null)
+                {
+                    streamingWritesCounter++;
+                }
                 return sendStatusFromServerTcs.Task;
             }
         }
diff --git a/src/csharp/Grpc.Core/Internal/ServerCallHandler.cs b/src/csharp/Grpc.Core/Internal/ServerCallHandler.cs
index 00d82d51e8..bbbefd0699 100644
--- a/src/csharp/Grpc.Core/Internal/ServerCallHandler.cs
+++ b/src/csharp/Grpc.Core/Internal/ServerCallHandler.cs
@@ -75,14 +75,15 @@ namespace Grpc.Core.Internal
             var responseStream = new ServerResponseStream<TRequest, TResponse>(asyncCall);
 
             Status status;
+            Tuple<TResponse,WriteFlags> responseTuple = null;
             var context = HandlerUtils.NewContext(newRpc, asyncCall.Peer, responseStream, asyncCall.CancellationToken);
             try
             {
                 GrpcPreconditions.CheckArgument(await requestStream.MoveNext().ConfigureAwait(false));
                 var request = requestStream.Current;
-                var result = await handler(request, context).ConfigureAwait(false);
+                var response = await handler(request, context).ConfigureAwait(false);
                 status = context.Status;
-                await responseStream.WriteAsync(result).ConfigureAwait(false);
+                responseTuple = Tuple.Create(response, HandlerUtils.GetWriteFlags(context.WriteOptions));
             } 
             catch (Exception e)
             {
@@ -91,7 +92,7 @@ namespace Grpc.Core.Internal
             }
             try
             {
-                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers).ConfigureAwait(false);
+                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers, responseTuple).ConfigureAwait(false);
             }
             catch (OperationCanceledException)
             {
@@ -145,7 +146,7 @@ namespace Grpc.Core.Internal
 
             try
             {
-                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers).ConfigureAwait(false);
+                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers, null).ConfigureAwait(false);
             }
             catch (OperationCanceledException)
             {
@@ -183,19 +184,13 @@ namespace Grpc.Core.Internal
             var responseStream = new ServerResponseStream<TRequest, TResponse>(asyncCall);
 
             Status status;
+            Tuple<TResponse,WriteFlags> responseTuple = null;
             var context = HandlerUtils.NewContext(newRpc, asyncCall.Peer, responseStream, asyncCall.CancellationToken);
             try
             {
-                var result = await handler(requestStream, context).ConfigureAwait(false);
+                var response = await handler(requestStream, context).ConfigureAwait(false);
                 status = context.Status;
-                try
-                {
-                    await responseStream.WriteAsync(result).ConfigureAwait(false);
-                }
-                catch (OperationCanceledException)
-                {
-                    status = Status.DefaultCancelled;
-                }
+                responseTuple = Tuple.Create(response, HandlerUtils.GetWriteFlags(context.WriteOptions));
             }
             catch (Exception e)
             {
@@ -205,7 +200,7 @@ namespace Grpc.Core.Internal
 
             try
             {
-                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers).ConfigureAwait(false);
+                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers, responseTuple).ConfigureAwait(false);
             }
             catch (OperationCanceledException)
             {
@@ -256,7 +251,7 @@ namespace Grpc.Core.Internal
             }
             try
             {
-                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers).ConfigureAwait(false);
+                await asyncCall.SendStatusFromServerAsync(status, context.ResponseTrailers, null).ConfigureAwait(false);
             }
             catch (OperationCanceledException)
             {
@@ -278,7 +273,7 @@ namespace Grpc.Core.Internal
             
             asyncCall.Initialize(newRpc.Call);
             var finishedTask = asyncCall.ServerSideCallAsync();
-            await asyncCall.SendStatusFromServerAsync(new Status(StatusCode.Unimplemented, ""), Metadata.Empty).ConfigureAwait(false);
+            await asyncCall.SendStatusFromServerAsync(new Status(StatusCode.Unimplemented, ""), Metadata.Empty, null).ConfigureAwait(false);
             await finishedTask.ConfigureAwait(false);
         }
     }
@@ -297,6 +292,11 @@ namespace Grpc.Core.Internal
             return new Status(StatusCode.Unknown, "Exception was thrown by handler.");
         }
 
+        public static WriteFlags GetWriteFlags(WriteOptions writeOptions)
+        {
+            return writeOptions != null ? writeOptions.Flags : default(WriteFlags);
+        }
+
         public static ServerCallContext NewContext<TRequest, TResponse>(ServerRpcNew newRpc, string peer, ServerResponseStream<TRequest, TResponse> serverResponseStream, CancellationToken cancellationToken)
             where TRequest : class
             where TResponse : class