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 21, 2023
1 parent 478f30e commit e3a54b4
Show file tree
Hide file tree
Showing 4 changed files with 16 additions and 69 deletions.
3 changes: 3 additions & 0 deletions core/src/main/java/io/grpc/internal/MessageFramer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,7 @@ public void writePayload(InputStream message) {
}

private int writeUncompressed(InputStream message, int messageLength) throws IOException {
verifyNotClosed();
if (messageLength != -1) {
currentMessageWireSize = messageLength;
return writeKnownLengthUncompressed(message, messageLength);
Expand Down Expand Up @@ -215,6 +216,7 @@ private int getKnownLength(InputStream inputStream) throws IOException {
*/
private int writeKnownLengthUncompressed(InputStream message, int messageLength)
throws IOException {
verifyNotClosed();
if (maxOutboundMessageSize >= 0 && messageLength > maxOutboundMessageSize) {
throw Status.RESOURCE_EXHAUSTED
.withDescription(
Expand Down Expand Up @@ -280,6 +282,7 @@ private static int writeToOutputStream(InputStream message, OutputStream outputS
}

private void writeRaw(byte[] b, int off, int len) {
verifyNotClosed();
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
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
1 change: 1 addition & 0 deletions netty/src/main/java/io/grpc/netty/NettyClientStream.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,7 @@ public void writeFrame(
@Override
public void cancel(Status status) {
try (TaskCloseable ignore = PerfMark.traceTask("NettyClientStream$Sink.cancel")) {
framer().dispose();
writeQueue.enqueue(new CancelClientStreamCommand(transportState(), status), true);
}
}
Expand Down

0 comments on commit e3a54b4

Please sign in to comment.