diff --git a/core/src/main/java/io/grpc/internal/RetriableStream.java b/core/src/main/java/io/grpc/internal/RetriableStream.java index 0c37a0beaca..eec53360853 100644 --- a/core/src/main/java/io/grpc/internal/RetriableStream.java +++ b/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); diff --git a/core/src/test/java/io/grpc/internal/RetriableStreamTest.java b/core/src/test/java/io/grpc/internal/RetriableStreamTest.java index afbdaa395b0..a958065f4c7 100644 --- a/core/src/test/java/io/grpc/internal/RetriableStreamTest.java +++ b/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 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 stream) { + Status reason = Status.CANCELLED.withDescription("cancel before registration"); + List> 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);