binder: Defer server close until suffix send completes

binder/src/main/java/io/grpc/binder/internal/Inbound.java

@@ -251,7 +251,7 @@ private final void deliverInternal() {
           break;
         }
         onDeliveryState(State.ALL_MESSAGES_DELIVERED);
-        // Fall-through.
+        // fall through
       case ALL_MESSAGES_DELIVERED:
         if (suffixAvailable) {
           onDeliveryState(State.SUFFIX_DELIVERED);
@@ -303,16 +303,21 @@ private final void closeAbnormal(
   @GuardedBy("this")
   protected abstract void deliverCloseAbnormal(Status status);
 
-  final void onTransportReady() {
+  void onTransportReady() {
     // Report transport readiness to the listener, and the outbound data.
     Outbound outbound = null;
-    StreamListener listener = null;
     synchronized (this) {
+      if (isClosed()) {
+        return;
+      }
+      if (listener != null) {
+        listener.onReady();
+      }
+      // The listener may have cancelled or closed the stream in onReady().
+      if (isClosed()) {
+        return;
+      }
       outbound = this.outbound;
-      listener = this.listener;
-    }
-    if (listener != null) {
-      listener.onReady();
     }
     if (outbound != null) {
       try {
@@ -680,12 +685,34 @@ void triggerEvent(Object event) {
       }
     }
 
-    @GuardedBy("this")
-    void onCloseSent(Status status) {
-      if (!isClosed()) {
-        onDeliveryState(State.CLOSED);
-        statsTraceContext.streamClosed(status);
-        listener.closed(Status.OK);
+    @Override
+    void onTransportReady() {
+      super.onTransportReady();
+      closeIfSuffixSent();
+    }
+
+    void closeIfSuffixSent() {
+      Outbound.ServerOutbound outbound;
+      synchronized (this) {
+        if (isClosed()) {
+          return;
+        }
+        outbound = (Outbound.ServerOutbound) this.outbound;
+      }
+      if (outbound == null) {
+        return;
+      }
+      Status status;
+      synchronized (outbound) {
+        status = outbound.getCloseStatusIfSent();
+      }
+      synchronized (this) {
+        if (status != null && !isClosed()) {
+          onDeliveryState(State.CLOSED);
+          unregister();
+          statsTraceContext.streamClosed(status);
+          listener.closed(Status.OK);
+        }
       }
     }
   }

binder/src/main/java/io/grpc/binder/internal/MultiMessageServerStream.java

@@ -114,9 +114,7 @@ public void close(Status status, Metadata trailers) {
       synchronized (outbound) {
         outbound.sendClose(status, trailers);
       }
-      synchronized (inbound) {
-        inbound.onCloseSent(status);
-      }
+      inbound.closeIfSuffixSent();
     } catch (StatusException se) {
       synchronized (inbound) {
         inbound.closeAbnormal(se.getStatus());

binder/src/main/java/io/grpc/binder/internal/Outbound.java

@@ -198,6 +198,11 @@ final boolean isReady() {
     return transport.isReady();
   }
 
+  @GuardedBy("this")
+  protected final boolean isSuffixSent() {
+    return outboundState == State.SUFFIX_SENT;
+  }
+
   @GuardedBy("this")
   final void onTransportReady() throws StatusException {
     // The transport has become ready, attempt sending.
@@ -232,7 +237,7 @@ protected final void sendInternal() throws StatusException {
           if (!messageAvailable() && !suffixReady) {
             break;
           }
-          // Fall-through.
+          // fall through
         case PREFIX_SENT:
           InputStream messageStream = peekNextMessage();
           if (messageStream != null) {
@@ -247,29 +252,27 @@ protected final void sendInternal() throws StatusException {
             // There's still more message data to deliver, break out.
             break;
           }
-          // Fall-through.
+          // fall through
         case ALL_MESSAGES_SENT:
           flags |= TransactionUtils.FLAG_SUFFIX;
           flags |= writeSuffix(parcel.get());
-          onOutboundState(State.SUFFIX_SENT);
           break;
         default:
           throw new AssertionError();
       }
       TransactionUtils.fillInFlags(parcel.get(), flags);
       int dataSize = parcel.get().dataSize();
       transport.sendTransaction(callId, parcel);
+      if (TransactionUtils.hasFlag(flags, TransactionUtils.FLAG_SUFFIX)) {
+        onOutboundState(State.SUFFIX_SENT);
+      }
       statsTraceContext.outboundWireSize(dataSize);
       statsTraceContext.outboundUncompressedSize(dataSize);
     } catch (IOException e) {
       throw Status.INTERNAL.withCause(e).asException();
     }
   }
 
-  protected final void unregister() {
-    transport.unregisterCall(callId);
-  }
-
   @Override
   public synchronized String toString() {
     return getClass().getSimpleName()
@@ -470,13 +473,17 @@ void sendClose(Status closeStatus, Metadata trailers) throws StatusException {
       send();
     }
 
+    @Nullable
+    @GuardedBy("this")
+    Status getCloseStatusIfSent() {
+      return isSuffixSent() ? closeStatus : null;
+    }
+
     @Override
     @GuardedBy("this")
     protected int writeSuffix(Parcel parcel) throws IOException, StatusException {
       int flags = TransactionUtils.writeStatus(parcel, closeStatus);
       MetadataHelper.writeMetadata(parcel, trailers);
-      // TODO: This is an ugly place for this side-effect.
-      unregister();
       return flags;
     }
   }

binder/src/main/java/io/grpc/binder/internal/SingleMessageServerStream.java

@@ -106,9 +106,7 @@ public void close(Status status, Metadata trailers) {
       synchronized (outbound) {
         outbound.sendSingleMessageAndClose(pendingHeaders, pendingSingleMessage, status, trailers);
       }
-      synchronized (inbound) {
-        inbound.onCloseSent(status);
-      }
+      inbound.closeIfSuffixSent();
     } catch (StatusException se) {
       synchronized (inbound) {
         inbound.closeAbnormal(se.getStatus());