core: prevent RetriableStream registry leak on cancel race

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

@@ -399,8 +399,19 @@ public final void start(ClientStreamListener listener) {
       return;
     }
 
+    // cancel() may have committed this stream before prestart() registered it. In that case the
+    // post-commit callback ran before registration and could not remove the stream. Run it again
+    // after registration so the channel's uncommitted stream registry cannot retain this stream.
+    boolean alreadyCommitted;
     synchronized (lock) {
-      state.buffer.add(new StartEntry());
+      alreadyCommitted = state.winningSubstream != null;
+      if (!alreadyCommitted) {
+        state.buffer.add(new StartEntry());
+      }
+    }
+    if (alreadyCommitted) {
+      postCommit();
+      return;
     }
 
     Substream substream = createSubstream(0, false, false);

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

@@ -235,6 +235,46 @@ public void tearDown() {
     assertEquals(0, fakeClock.numPendingTasks());
   }
 
+  @Test
+  public void cancelDuringStartBeforeRegistration() {
+    assertCancelDuringStart(retriableStream);
+  }
+
+  @Test
+  public void hedging_cancelBeforeRegistration() {
+    assertCancelDuringStart(hedgingStream);
+  }
+
+  @Test
+  public void transparentRetry_cancelBeforeRegistration() {
+    RetriableStream<String> stream = new RecordedRetriableStream(
+        method, new Metadata(), channelBufferUsed, PER_RPC_BUFFER_LIMIT, CHANNEL_BUFFER_LIMIT,
+        MoreExecutors.directExecutor(), fakeClock.getScheduledExecutorService(), null, null, null);
+    assertCancelDuringStart(stream);
+  }
+
+  private void assertCancelDuringStart(RetriableStream<String> stream) {
+    Status reason = Status.CANCELLED.withDescription("cancel before registration");
+    List<RetriableStream<String>> registeredStreams = new ArrayList<>();
+    doAnswer(invocation -> {
+      stream.cancel(reason);
+      registeredStreams.add(stream);
+      return null;
+    }).when(retriableStreamRecorder).prestart();
+    doAnswer(invocation -> {
+      registeredStreams.remove(stream);
+      return null;
+    }).when(retriableStreamRecorder).postCommit();
+
+    stream.start(masterListener);
+
+    assertThat(registeredStreams).isEmpty();
+    verify(retriableStreamRecorder).prestart();
+    verify(retriableStreamRecorder, times(2)).postCommit();
+    verify(masterListener).closed(same(reason), same(PROCESSED), any(Metadata.class));
+    verify(retriableStreamRecorder, never()).newSubstream(anyInt());
+  }
+
   @Test
   public void retry_everythingDrained() {
     ClientStream mockStream1 = mock(ClientStream.class);