Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

core: backport RetriableStream bug fix (backport v1.40.x) #8387

Merged
merged 2 commits into from Aug 9, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
63 changes: 40 additions & 23 deletions core/src/main/java/io/grpc/internal/RetriableStream.java
Expand Up @@ -104,6 +104,7 @@ abstract class RetriableStream<ReqT> implements ClientStream {
@GuardedBy("lock")
private FutureCanceller scheduledHedging;
private long nextBackoffIntervalNanos;
private Status cancellationStatus;

RetriableStream(
MethodDescriptor<ReqT, ?> method, Metadata headers,
Expand Down Expand Up @@ -244,15 +245,21 @@ private void drain(Substream substream) {
int index = 0;
int chunk = 0x80;
List<BufferEntry> list = null;
boolean streamStarted = false;

while (true) {
State savedState;

synchronized (lock) {
savedState = state;
if (savedState.winningSubstream != null && savedState.winningSubstream != substream) {
// committed but not me
break;
if (streamStarted) {
if (savedState.winningSubstream != null && savedState.winningSubstream != substream) {
// committed but not me, to be cancelled
break;
}
if (savedState.cancelled) {
break;
}
}
if (index == savedState.buffer.size()) { // I'm drained
state = savedState.substreamDrained(substream);
Expand All @@ -274,22 +281,25 @@ private void drain(Substream substream) {
}

for (BufferEntry bufferEntry : list) {
savedState = state;
if (savedState.winningSubstream != null && savedState.winningSubstream != substream) {
// committed but not me
break;
bufferEntry.runWith(substream);
if (bufferEntry instanceof RetriableStream.StartEntry) {
streamStarted = true;
}
if (savedState.cancelled) {
checkState(
savedState.winningSubstream == substream,
"substream should be CANCELLED_BECAUSE_COMMITTED already");
return;
if (streamStarted) {
savedState = state;
if (savedState.winningSubstream != null && savedState.winningSubstream != substream) {
// committed but not me, to be cancelled
break;
}
if (savedState.cancelled) {
break;
}
}
bufferEntry.runWith(substream);
}
}

substream.stream.cancel(CANCELLED_BECAUSE_COMMITTED);
substream.stream.cancel(
state.winningSubstream == substream ? cancellationStatus : CANCELLED_BECAUSE_COMMITTED);
}

/**
Expand All @@ -299,6 +309,13 @@ private void drain(Substream substream) {
@Nullable
abstract Status prestart();

class StartEntry implements BufferEntry {
@Override
public void runWith(Substream substream) {
substream.stream.start(new Sublistener(substream));
}
}

/** Starts the first PRC attempt. */
@Override
public final void start(ClientStreamListener listener) {
Expand All @@ -311,13 +328,6 @@ public final void start(ClientStreamListener listener) {
return;
}

class StartEntry implements BufferEntry {
@Override
public void runWith(Substream substream) {
substream.stream.start(new Sublistener(substream));
}
}

synchronized (lock) {
state.buffer.add(new StartEntry());
}
Expand Down Expand Up @@ -450,11 +460,18 @@ public final void cancel(Status reason) {
return;
}

state.winningSubstream.stream.cancel(reason);
Substream winningSubstreamToCancel = null;
synchronized (lock) {
// This is not required, but causes a short-circuit in the draining process.
if (state.drainedSubstreams.contains(state.winningSubstream)) {
winningSubstreamToCancel = state.winningSubstream;
} else { // the winningSubstream will be cancelled while draining
cancellationStatus = reason;
}
state = state.cancelled();
}
if (winningSubstreamToCancel != null) {
winningSubstreamToCancel.stream.cancel(reason);
}
}

private void delayOrExecute(BufferEntry bufferEntry) {
Expand Down
126 changes: 126 additions & 0 deletions core/src/test/java/io/grpc/internal/RetriableStreamTest.java
Expand Up @@ -749,6 +749,91 @@ public void request(int numMessages) {
inOrder.verify(mockStream2, never()).writeMessage(any(InputStream.class));
}

@Test
public void cancelWhileDraining() {
ArgumentCaptor<ClientStreamListener> sublistenerCaptor1 =
ArgumentCaptor.forClass(ClientStreamListener.class);
ClientStream mockStream1 = mock(ClientStream.class);
ClientStream mockStream2 =
mock(
ClientStream.class,
delegatesTo(
new NoopClientStream() {
@Override
public void request(int numMessages) {
retriableStream.cancel(
Status.CANCELLED.withDescription("cancelled while requesting"));
}
}));

InOrder inOrder = inOrder(retriableStreamRecorder, mockStream1, mockStream2);
doReturn(mockStream1).when(retriableStreamRecorder).newSubstream(0);
retriableStream.start(masterListener);
inOrder.verify(mockStream1).start(sublistenerCaptor1.capture());
retriableStream.request(3);
inOrder.verify(mockStream1).request(3);

// retry
doReturn(mockStream2).when(retriableStreamRecorder).newSubstream(1);
sublistenerCaptor1.getValue().closed(
Status.fromCode(RETRIABLE_STATUS_CODE_1), PROCESSED, new Metadata());
fakeClock.forwardTime((long) (INITIAL_BACKOFF_IN_SECONDS * FAKE_RANDOM), TimeUnit.SECONDS);

inOrder.verify(mockStream2).start(any(ClientStreamListener.class));
inOrder.verify(mockStream2).request(3);
inOrder.verify(retriableStreamRecorder).postCommit();
ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class);
inOrder.verify(mockStream2).cancel(statusCaptor.capture());
assertThat(statusCaptor.getValue().getCode()).isEqualTo(Code.CANCELLED);
assertThat(statusCaptor.getValue().getDescription())
.isEqualTo("Stream thrown away because RetriableStream committed");
verify(masterListener).closed(
statusCaptor.capture(), any(RpcProgress.class), any(Metadata.class));
assertThat(statusCaptor.getValue().getCode()).isEqualTo(Code.CANCELLED);
assertThat(statusCaptor.getValue().getDescription()).isEqualTo("cancelled while requesting");
}

@Test
public void cancelWhileRetryStart() {
ArgumentCaptor<ClientStreamListener> sublistenerCaptor1 =
ArgumentCaptor.forClass(ClientStreamListener.class);
ClientStream mockStream1 = mock(ClientStream.class);
ClientStream mockStream2 =
mock(
ClientStream.class,
delegatesTo(
new NoopClientStream() {
@Override
public void start(ClientStreamListener listener) {
retriableStream.cancel(
Status.CANCELLED.withDescription("cancelled while retry start"));
}
}));

InOrder inOrder = inOrder(retriableStreamRecorder, mockStream1, mockStream2);
doReturn(mockStream1).when(retriableStreamRecorder).newSubstream(0);
retriableStream.start(masterListener);
inOrder.verify(mockStream1).start(sublistenerCaptor1.capture());

// retry
doReturn(mockStream2).when(retriableStreamRecorder).newSubstream(1);
sublistenerCaptor1.getValue().closed(
Status.fromCode(RETRIABLE_STATUS_CODE_1), PROCESSED, new Metadata());
fakeClock.forwardTime((long) (INITIAL_BACKOFF_IN_SECONDS * FAKE_RANDOM), TimeUnit.SECONDS);

inOrder.verify(mockStream2).start(any(ClientStreamListener.class));
inOrder.verify(retriableStreamRecorder).postCommit();
ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class);
inOrder.verify(mockStream2).cancel(statusCaptor.capture());
assertThat(statusCaptor.getValue().getCode()).isEqualTo(Code.CANCELLED);
assertThat(statusCaptor.getValue().getDescription())
.isEqualTo("Stream thrown away because RetriableStream committed");
verify(masterListener).closed(
statusCaptor.capture(), any(RpcProgress.class), any(Metadata.class));
assertThat(statusCaptor.getValue().getCode()).isEqualTo(Code.CANCELLED);
assertThat(statusCaptor.getValue().getDescription()).isEqualTo("cancelled while retry start");
}

@Test
public void operationsAfterImmediateCommit() {
ArgumentCaptor<ClientStreamListener> sublistenerCaptor1 =
Expand Down Expand Up @@ -916,6 +1001,47 @@ public void start(ClientStreamListener listener) {
verify(mockStream3).request(1);
}

@Test
public void commitAndCancelWhileDraining() {
ClientStream mockStream1 = mock(ClientStream.class);
ClientStream mockStream2 =
mock(
ClientStream.class,
delegatesTo(
new NoopClientStream() {
@Override
public void start(ClientStreamListener listener) {
// commit while draining
listener.headersRead(new Metadata());
// cancel while draining
retriableStream.cancel(
Status.CANCELLED.withDescription("cancelled while drained"));
}
}));

when(retriableStreamRecorder.newSubstream(anyInt()))
.thenReturn(mockStream1, mockStream2);

retriableStream.start(masterListener);

ArgumentCaptor<ClientStreamListener> sublistenerCaptor1 =
ArgumentCaptor.forClass(ClientStreamListener.class);
verify(mockStream1).start(sublistenerCaptor1.capture());

ClientStreamListener listener1 = sublistenerCaptor1.getValue();

// retry
listener1.closed(
Status.fromCode(RETRIABLE_STATUS_CODE_1), PROCESSED, new Metadata());
fakeClock.forwardTime((long) (INITIAL_BACKOFF_IN_SECONDS * FAKE_RANDOM), TimeUnit.SECONDS);

verify(mockStream2).start(any(ClientStreamListener.class));
verify(retriableStreamRecorder).postCommit();
ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class);
verify(mockStream2).cancel(statusCaptor.capture());
assertThat(statusCaptor.getValue().getCode()).isEqualTo(Code.CANCELLED);
assertThat(statusCaptor.getValue().getDescription()).isEqualTo("cancelled while drained");
}

@Test
public void perRpcBufferLimitExceeded() {
Expand Down