stub: Ignore unary response on server if status is not OK

Fixes grpc#5969

Author: NaveenPrasannaV

stub/src/main/java/io/grpc/stub/ServerCalls.java

@@ -335,6 +335,7 @@ private static final class ServerCallStreamObserverImpl<ReqT, RespT>
     private boolean aborted = false;
     private boolean completed = false;
     private Runnable onCloseHandler;
+    private RespT unaryResponse;
 
     // Non private to avoid synthetic class
     ServerCallStreamObserverImpl(ServerCall<ReqT, RespT> call, boolean serverStreamingOrBidi) {
@@ -373,15 +374,22 @@ public void onNext(RespT response) {
       }
       checkState(!aborted, "Stream was terminated by error, no further calls are allowed");
       checkState(!completed, "Stream is already completed, no further calls are allowed");
-      if (!sentHeaders) {
-        call.sendHeaders(new Metadata());
-        sentHeaders = true;
+      if (serverStreamingOrBidi) {
+        if (!sentHeaders) {
+          call.sendHeaders(new Metadata());
+          sentHeaders = true;
+        }
+        call.sendMessage(response);
+      } else {
+        unaryResponse = response;
       }
-      call.sendMessage(response);
     }
 
     @Override
     public void onError(Throwable t) {
+      if (!serverStreamingOrBidi) {
+        unaryResponse = null;
+      }
       Metadata metadata = Status.trailersFromThrowable(t);
       if (metadata == null) {
         metadata = new Metadata();
@@ -392,6 +400,14 @@ public void onError(Throwable t) {
 
     @Override
     public void onCompleted() {
+      if (!serverStreamingOrBidi && unaryResponse != null) {
+        if (!sentHeaders) {
+          call.sendHeaders(new Metadata());
+          sentHeaders = true;
+        }
+        call.sendMessage(unaryResponse);
+        unaryResponse = null;
+      }
       call.close(Status.OK, new Metadata());
       completed = true;
     }

stub/src/test/java/io/grpc/stub/ServerCallsTest.java

@@ -33,6 +33,7 @@
 import io.grpc.ServerServiceDefinition;
 import io.grpc.ServiceDescriptor;
 import io.grpc.Status;
+import io.grpc.Status.Code;
 import io.grpc.StatusRuntimeException;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
@@ -620,6 +621,59 @@ public void onClose(Status status, Metadata trailers) {
     assertArrayEquals(new int[]{0, 1, 1, 2, 2, 2}, receivedMessages);
   }
 
+  @Test
+  public void serverUnaryResponseMsgWithOkStatus() {
+    ServerCallHandler<Integer, Integer> serverCallHandler =
+        ServerCalls.asyncUnaryCall(
+            new ServerCalls.UnaryMethod<Integer, Integer>() {
+              @Override
+              public void invoke(Integer request, StreamObserver<Integer> responseObserver) {
+                responseObserver.onNext(request);
+                responseObserver.onCompleted();
+              }
+            });
+    ServerCall.Listener<Integer> callListener =
+        serverCallHandler.startCall(serverCall, new Metadata());
+    serverCall.isReady = true;
+    serverCall.isCancelled = false;
+    callListener.onReady();
+    callListener.onMessage(1);
+    callListener.onHalfClose();
+
+    assertThat(serverCall.status.getCode()).isEqualTo(Code.OK);
+    assertThat(serverCall.responses).containsExactly(1);
+  }
+
+  @Test
+  public void serverUnaryResponseMsgWithNotOkStatus() {
+    ServerCallHandler<Integer, Integer> serverCallHandler =
+        ServerCalls.asyncUnaryCall(
+            new ServerCalls.UnaryMethod<Integer, Integer>() {
+              @Override
+              public void invoke(Integer request, StreamObserver<Integer> responseObserver) {
+                responseObserver.onNext(request);
+                responseObserver.onError(
+                    Status.INTERNAL
+                        .withDescription("Response message is null for unary call")
+                        .asRuntimeException());
+              }
+            });
+
+    ServerCall.Listener<Integer> callListener =
+        serverCallHandler.startCall(serverCall, new Metadata());
+
+    serverCall.isReady = true;
+    serverCall.isCancelled = false;
+    callListener.onReady();
+    callListener.onMessage(1);
+    callListener.onHalfClose();
+
+    assertThat(serverCall.status.getCode()).isEqualTo(Code.INTERNAL);
+    assertThat(serverCall.status.getDescription())
+        .isEqualTo("Response message is null for unary call");
+    assertThat(serverCall.responses).isEmpty();
+  }
+
   public static class IntegerMarshaller implements MethodDescriptor.Marshaller<Integer> {
     @Override
     public InputStream stream(Integer value) {