Skip to content

Commit

Permalink
core: Rework retry memory leak fix in #9360 to send fewer FIN packets.
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@jtk54
jtk54 committed May 24, 2023
1 parent 478f30e commit f66e8df
Show file tree
Hide file tree
Showing 5 changed files with 69 additions and 72 deletions.
1 change: 1 addition & 0 deletions core/src/main/java/io/grpc/internal/AbstractClientStream.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
19 changes: 16 additions & 3 deletions core/src/main/java/io/grpc/internal/MessageFramer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import static com.google.common.base.Preconditions.checkState;
import static java.lang.Math.min;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.io.ByteStreams;
import io.grpc.Codec;
import io.grpc.Compressor;
Expand Down Expand Up @@ -163,7 +164,11 @@ public void writePayload(InputStream message) {
statsTraceCtx.outboundMessageSent(currentMessageSeqNo, currentMessageWireSize, written);
}

private int writeUncompressed(InputStream message, int messageLength) throws IOException {
@VisibleForTesting
int writeUncompressed(InputStream message, int messageLength) throws IOException {
if (isClosed()) {
return -1;
}
if (messageLength != -1) {
currentMessageWireSize = messageLength;
return writeKnownLengthUncompressed(message, messageLength);
Expand Down Expand Up @@ -213,8 +218,12 @@ private int getKnownLength(InputStream inputStream) throws IOException {
/**
* Write an unserialized message with a known length, uncompressed.
*/
private int writeKnownLengthUncompressed(InputStream message, int messageLength)
@VisibleForTesting
int writeKnownLengthUncompressed(InputStream message, int messageLength)
throws IOException {
if (isClosed()) {
return -1;
}
if (maxOutboundMessageSize >= 0 && messageLength > maxOutboundMessageSize) {
throw Status.RESOURCE_EXHAUSTED
.withDescription(
Expand Down Expand Up @@ -279,7 +288,11 @@ private static int writeToOutputStream(InputStream message, OutputStream outputS
}
}

private void writeRaw(byte[] b, int off, int len) {
@VisibleForTesting
void writeRaw(byte[] b, int off, int len) {
if (isClosed()) {
return;
}
while (len > 0) {
if (buffer != null && buffer.writableBytes() == 0) {
commitToSink(false, false);
Expand Down
4 changes: 0 additions & 4 deletions core/src/main/java/io/grpc/internal/RetriableStream.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -565,10 +565,6 @@ 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();
}
}

Expand Down
40 changes: 40 additions & 0 deletions core/src/test/java/io/grpc/internal/MessageFramerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import io.grpc.internal.testing.TestStreamTracer.TestBaseStreamTracer;
import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.Arrays;
Expand Down Expand Up @@ -323,6 +324,45 @@ public void dontCompressIfNotRequested() {
checkStats(1000, 1000);
}

@Test
public void writeUncompressedNoopIfDisposed() {
byte[] bytes = new byte[]{3, 14};
int written;
framer.dispose();
try {
written = framer.writeUncompressed(new ByteArrayInputStream(bytes), bytes.length);
} catch (IOException e) {
throw new AssertionError("framer.writeUncompressed failed", e);
}
assertEquals(-1, written);
verifyNoMoreInteractions(sink);
assertEquals(0, allocator.allocCount);
}

@Test
public void writeKnownLengthUncompressedNoopIfDisposed() {
byte[] bytes = new byte[]{3, 14};
int written;
framer.dispose();
try {
written = framer.writeKnownLengthUncompressed(new ByteArrayInputStream(bytes), bytes.length);
} catch (IOException e) {
throw new AssertionError("framer.writeUncompressed failed", e);
}
assertEquals(-1, written);
verifyNoMoreInteractions(sink);
assertEquals(0, allocator.allocCount);
}

@Test
public void writeRawNoopIfDisposed() {
byte[] bytes = new byte[]{3, 14};
framer.dispose();
framer.writeRaw(bytes, 0, bytes.length);
verifyNoMoreInteractions(sink);
assertEquals(0, allocator.allocCount);
}

@Test
public void closeIsRentrantSafe() {
MessageFramer.Sink reentrant = new MessageFramer.Sink() {
Expand Down
77 changes: 12 additions & 65 deletions core/src/test/java/io/grpc/internal/RetriableStreamTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down Expand Up @@ -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();

Expand All @@ -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);
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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();

Expand All @@ -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();
Expand All @@ -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();


Expand All @@ -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();
Expand All @@ -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(
Expand All @@ -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);
Expand All @@ -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();
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit f66e8df

Please sign in to comment.