core: Rework retry memory leak fix in #9360 to send fewer FIN packets.

Without the explicit stream flush after writes in RetriableStream, buffer data is being
orphaned on RetriableStream cancellation since calls to writeRaw() can happen after
dispose() calls. We should verify the framer is not closed before continuing with
writes since writes and dispose can be interleaved.

core/src/main/java/io/grpc/internal/AbstractClientStream.java

@@ -199,6 +199,7 @@ public final void halfClose() {
   public final void cancel(Status reason) {
     Preconditions.checkArgument(!reason.isOk(), "Should not cancel with OK status");
     cancelled = true;
+    framer().dispose();
     abstractClientStreamSink().cancel(reason);
   }
 

core/src/main/java/io/grpc/internal/RetriableStream.java

@@ -34,11 +34,7 @@
 import io.grpc.internal.ClientStreamListener.RpcProgress;
 import java.io.InputStream;
 import java.lang.Thread.UncaughtExceptionHandler;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-import java.util.Random;
+import java.util.*;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
@@ -504,17 +500,34 @@ public void run() {
 
   @Override
   public final void cancel(final Status reason) {
-    Substream noopSubstream = new Substream(0 /* previousAttempts doesn't matter here */);
-    noopSubstream.stream = new NoopClientStream();
-    Runnable runnable = commit(noopSubstream);
+    if (isHedging) {
+      Substream noopSubstream = new Substream(0 /* previousAttempts doesn't matter here */);
+      noopSubstream.stream = new NoopClientStream();
+      Runnable runnable = commit(noopSubstream);
 
-    if (runnable != null) {
+      if (runnable != null) {
+        synchronized (lock) {
+          state = state.substreamDrained(noopSubstream);
+        }
+        runnable.run();
+        safeCloseMasterListener(reason, RpcProgress.PROCESSED, new Metadata());
+        return;
+      }
+    } else {
+      Optional<Substream> substreamToCommit;
       synchronized (lock) {
-        state = state.substreamDrained(noopSubstream);
+        substreamToCommit = state.drainedSubstreams.stream().findFirst();
+      }
+      if (substreamToCommit.isPresent()) {
+        Substream substream = substreamToCommit.get();
+        Runnable runnable = commit(substream);
+        if (runnable != null) {
+          runnable.run();
+//          substream.stream.cancel(reason);
+          safeCloseMasterListener(reason, RpcProgress.PROCESSED, new Metadata());
+          return;
+        }
       }
-      runnable.run();
-      safeCloseMasterListener(reason, RpcProgress.PROCESSED, new Metadata());
-      return;
     }
 
     Substream winningSubstreamToCancel = null;
@@ -565,10 +578,12 @@ class SendMessageEntry implements BufferEntry {
       @Override
       public void runWith(Substream substream) {
         substream.stream.writeMessage(method.streamRequest(message));
-        // TODO(ejona): Workaround Netty memory leak. Message writes always need to be followed by
-        // flushes (or half close), but retry appears to have a code path that the flushes may
-        // not happen. The code needs to be fixed and this removed. See #9340.
-        substream.stream.flush();
+        if (isHedging) {
+          // TODO(ejona): Workaround Netty memory leak. Message writes always need to be followed by
+          // flushes (or half close), but retry appears to have a code path that the flushes may
+          // not happen. The code needs to be fixed and this removed. See #9340.
+          substream.stream.flush();
+        }
       }
     }
 
@@ -1129,7 +1144,7 @@ private static final class State {
 
     /**
      * Unmodifiable collection of all the open substreams that are drained. Singleton once
-     * passThrough; Empty if committed but not passTrough.
+     * passThrough; Empty if committed but not passThrough.
      */
     final Collection<Substream> drainedSubstreams;
 
@@ -1180,7 +1195,6 @@ private static final class State {
               || (drainedSubstreams.size() == 1 && drainedSubstreams.contains(winningSubstream))
               || (drainedSubstreams.size() == 0 && winningSubstream.closed),
           "passThrough should imply winningSubstream is drained");
-      checkState(!cancelled || winningSubstream != null, "cancelled should imply committed");
     }
 
     @CheckReturnValue

core/src/test/java/io/grpc/internal/AbstractClientStreamTest.java

@@ -52,6 +52,8 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.SocketAddress;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import org.junit.Before;
 import org.junit.Rule;
@@ -158,6 +160,18 @@ public void cancel(Status errorStatus) {
     verify(mockListener).closed(any(Status.class), same(PROCESSED), any(Metadata.class));
   }
 
+  @Test
+  public void cancel_closesFramerAndReleasesBuffers() {
+    TrackingWritableBufferAllocator trackingAllocator = new TrackingWritableBufferAllocator();
+    AbstractClientStream stream =
+            new BaseAbstractClientStream(trackingAllocator, statsTraceCtx, transportTracer);
+    stream.start(mockListener);
+    stream.writeMessage(new ByteArrayInputStream(new byte[1]));
+    stream.cancel(Status.DEADLINE_EXCEEDED);
+    assertTrue(trackingAllocator.allocatedBuffersReleased());
+    assertTrue(stream.framer().isClosed());
+  }
+
   @Test
   public void startFailsOnNullListener() {
     AbstractClientStream stream =
@@ -584,4 +598,37 @@ public void runOnTransportThread(Runnable r) {
       r.run();
     }
   }
+
+  private static class TrackingWritableBufferAllocator implements WritableBufferAllocator {
+    List<ReleaseVerifyingBuffer> allocatedBuffers = new ArrayList<>();
+
+    @Override
+    public WritableBuffer allocate(int capacityHint) {
+      ReleaseVerifyingBuffer buf = new ReleaseVerifyingBuffer(capacityHint);
+      allocatedBuffers.add(buf);
+      return buf;
+    }
+
+    boolean allocatedBuffersReleased() {
+      return allocatedBuffers.stream().allMatch(ReleaseVerifyingBuffer::isReleased);
+    }
+  }
+
+  private static class ReleaseVerifyingBuffer extends ByteWritableBuffer {
+    boolean isReleased;
+
+    ReleaseVerifyingBuffer(int maxFrameSize) {
+      super(maxFrameSize);
+    }
+
+    @Override
+    public void release() {
+      super.release();
+      isReleased = true;
+    }
+
+    boolean isReleased() {
+      return isReleased;
+    }
+  }
 }

core/src/test/java/io/grpc/internal/RetriableStreamTest.java

@@ -279,14 +279,10 @@ public Void answer(InvocationOnMock in) {
     retriableStream.sendMessage("msg3");
     retriableStream.request(456);
 
-    inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
-    inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
+    inOrder.verify(mockStream1, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream1).request(345);
     inOrder.verify(mockStream1, times(2)).flush();
     inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
     inOrder.verify(mockStream1).request(456);
     inOrder.verifyNoMoreInteractions();
 
@@ -319,19 +315,12 @@ public Void answer(InvocationOnMock in) {
     ArgumentCaptor<ClientStreamListener> sublistenerCaptor2 =
         ArgumentCaptor.forClass(ClientStreamListener.class);
     inOrder.verify(mockStream2).start(sublistenerCaptor2.capture());
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
+    inOrder.verify(mockStream2, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream2).request(345);
     inOrder.verify(mockStream2, times(2)).flush();
     inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
     inOrder.verify(mockStream2).request(456);
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
+    inOrder.verify(mockStream2, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream2).isReady();
     inOrder.verifyNoMoreInteractions();
 
@@ -341,10 +330,7 @@ public Void answer(InvocationOnMock in) {
 
     // mockStream1 is closed so it is not in the drainedSubstreams
     verifyNoMoreInteractions(mockStream1);
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
+    inOrder.verify(mockStream2, times(2)).writeMessage(any(InputStream.class));
 
     // retry2
     doReturn(mockStream3).when(retriableStreamRecorder).newSubstream(2);
@@ -378,19 +364,12 @@ public Void answer(InvocationOnMock in) {
     ArgumentCaptor<ClientStreamListener> sublistenerCaptor3 =
         ArgumentCaptor.forClass(ClientStreamListener.class);
     inOrder.verify(mockStream3).start(sublistenerCaptor3.capture());
-    inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
-    inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
+    inOrder.verify(mockStream3, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream3).request(345);
     inOrder.verify(mockStream3, times(2)).flush();
     inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
     inOrder.verify(mockStream3).request(456);
-    for (int i = 0; i < 7; i++) {
-      inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-      inOrder.verify(mockStream3).flush(); // Memory leak workaround
-    }
+    inOrder.verify(mockStream3, times(7)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream3).isReady();
     inOrder.verifyNoMoreInteractions();
 
@@ -2103,14 +2082,10 @@ public Void answer(InvocationOnMock in) {
     hedgingStream.sendMessage("msg3");
     hedgingStream.request(456);
 
-    inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
-    inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
+    inOrder.verify(mockStream1, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream1).request(345);
     inOrder.verify(mockStream1, times(2)).flush();
     inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
     inOrder.verify(mockStream1).request(456);
     inOrder.verifyNoMoreInteractions();
 
@@ -2133,14 +2108,10 @@ public Void answer(InvocationOnMock in) {
     ArgumentCaptor<ClientStreamListener> sublistenerCaptor2 =
         ArgumentCaptor.forClass(ClientStreamListener.class);
     inOrder.verify(mockStream2).start(sublistenerCaptor2.capture());
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
-    inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
+    inOrder.verify(mockStream2, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream2).request(345);
     inOrder.verify(mockStream2, times(2)).flush();
     inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
     inOrder.verify(mockStream2).request(456);
     inOrder.verify(mockStream1).isReady();
     inOrder.verify(mockStream2).isReady();
@@ -2151,13 +2122,9 @@ public Void answer(InvocationOnMock in) {
     hedgingStream.sendMessage("msg2 after hedge2 starts");
 
     inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
     inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
     inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
     inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
     inOrder.verifyNoMoreInteractions();
 
 
@@ -2179,19 +2146,12 @@ public Void answer(InvocationOnMock in) {
     ArgumentCaptor<ClientStreamListener> sublistenerCaptor3 =
         ArgumentCaptor.forClass(ClientStreamListener.class);
     inOrder.verify(mockStream3).start(sublistenerCaptor3.capture());
-    inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
-    inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
+    inOrder.verify(mockStream3, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream3).request(345);
     inOrder.verify(mockStream3, times(2)).flush();
     inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
     inOrder.verify(mockStream3).request(456);
-    inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
-    inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
+    inOrder.verify(mockStream3, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream1).isReady();
     inOrder.verify(mockStream2).isReady();
     inOrder.verify(mockStream3).isReady();
@@ -2200,11 +2160,8 @@ public Void answer(InvocationOnMock in) {
     // send one more message
     hedgingStream.sendMessage("msg1 after hedge3 starts");
     inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
     inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
     inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
 
     // hedge3 receives nonFatalStatus
     sublistenerCaptor3.getValue().closed(
@@ -2214,9 +2171,7 @@ public Void answer(InvocationOnMock in) {
     // send one more message
     hedgingStream.sendMessage("msg1 after hedge3 fails");
     inOrder.verify(mockStream1).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream1).flush(); // Memory leak workaround
     inOrder.verify(mockStream2).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream2).flush(); // Memory leak workaround
 
     // the hedge mockStream4 starts
     fakeClock.forwardTime(HEDGING_DELAY_IN_SECONDS, TimeUnit.SECONDS);
@@ -2236,19 +2191,12 @@ public Void answer(InvocationOnMock in) {
     ArgumentCaptor<ClientStreamListener> sublistenerCaptor4 =
         ArgumentCaptor.forClass(ClientStreamListener.class);
     inOrder.verify(mockStream4).start(sublistenerCaptor4.capture());
-    inOrder.verify(mockStream4).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream4).flush(); // Memory leak workaround
-    inOrder.verify(mockStream4).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream4).flush(); // Memory leak workaround
+    inOrder.verify(mockStream4, times(2)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream4).request(345);
     inOrder.verify(mockStream4, times(2)).flush();
     inOrder.verify(mockStream4).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream4).flush(); // Memory leak workaround
     inOrder.verify(mockStream4).request(456);
-    for (int i = 0; i < 4; i++) {
-      inOrder.verify(mockStream4).writeMessage(any(InputStream.class));
-      inOrder.verify(mockStream4).flush(); // Memory leak workaround
-    }
+    inOrder.verify(mockStream4, times(4)).writeMessage(any(InputStream.class));
     inOrder.verify(mockStream1).isReady();
     inOrder.verify(mockStream2).isReady();
     inOrder.verify(mockStream4).isReady();
@@ -2371,7 +2319,6 @@ public void hedging_maxAttempts() {
 
     hedgingStream.sendMessage("msg1 after commit");
     inOrder.verify(mockStream3).writeMessage(any(InputStream.class));
-    inOrder.verify(mockStream3).flush(); // Memory leak workaround
     inOrder.verifyNoMoreInteractions();
 
     Metadata heders = new Metadata();