Skip to content

core: Delayed deserialization for unary/server-streaming calls #13004

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.

Sign up for GitHub

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

Open
kannanjgithub wants to merge 4 commits into
base: master
Choose a base branch
from
+446 −25
Open

core: Delayed deserialization for unary/server-streaming calls #13004

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
09d8c90
Fix remote OOM via delayed deserialization for unary/server-streaming…
kannanjgithub Aug 25, 2026
553b33a
Address review comments.
kannanjgithub Aug 26, 2026
50197f1
Synchronously mark the call as cancelled by calling `closedInternal` …
kannanjgithub Aug 27, 2026
f237d2f
Address review comments.
kannanjgithub Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 96 additions & 8 deletions core/src/main/java/io/grpc/internal/MessageDeframer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -437,14 +437,102 @@ private InputStream getCompressedBody() {
.asRuntimeException();
}

try {
// Enforce the maxMessageSize limit on the returned stream.
InputStream unlimitedStream =
decompressor.decompress(ReadableBuffers.openStream(nextFrame, true));
return new SizeEnforcingInputStream(
unlimitedStream, maxInboundMessageSize, statsTraceCtx);
} catch (IOException e) {
throw new RuntimeException(e);
return new LazyDecompressingInputStream(
ReadableBuffers.openStream(nextFrame, true),
maxInboundMessageSize,
statsTraceCtx,
decompressor);
}

/**
* An {@link InputStream} that delays decompressing a compressed frame until data is first read.
*/
@VisibleForTesting
static final class LazyDecompressingInputStream extends FilterInputStream {
private final Decompressor decompressor;
private final int maxMessageSize;
private final StatsTraceContext statsTraceCtx;
private boolean initialized;
private boolean closed;

LazyDecompressingInputStream(
InputStream rawStream,
int maxMessageSize,
StatsTraceContext statsTraceCtx,
Decompressor decompressor) {
super(rawStream);
this.decompressor = decompressor;
this.maxMessageSize = maxMessageSize;
this.statsTraceCtx = statsTraceCtx;
}

private synchronized void ensureInitialized() throws IOException {
if (closed) {
throw new IOException("Stream closed");
}
if (!initialized) {
InputStream decompressed = decompressor.decompress(in);
in = new SizeEnforcingInputStream(decompressed, maxMessageSize, statsTraceCtx);
initialized = true;
}
}

@Override
public int read() throws IOException {
ensureInitialized();
return super.read();
}

@Override
public int read(byte[] b, int off, int len) throws IOException {
ensureInitialized();
return super.read(b, off, len);
}

@Override
public long skip(long n) throws IOException {
ensureInitialized();
return super.skip(n);
}

@Override
public int available() throws IOException {
ensureInitialized();
return super.available();
}

@Override
public synchronized void close() throws IOException {
if (!closed) {
closed = true;
super.close();
}
}

@Override
public synchronized void mark(int readlimit) {
try {
ensureInitialized();
super.mark(readlimit);
} catch (IOException e) {
throw new RuntimeException(e);
}
}

@Override
public synchronized void reset() throws IOException {
ensureInitialized();
super.reset();
}

@Override
public boolean markSupported() {
try {
ensureInitialized();
return super.markSupported();
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}

Expand Down
44 changes: 38 additions & 6 deletions core/src/main/java/io/grpc/internal/ServerCallImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,7 @@ static final class ServerStreamListenerImpl<ReqT> implements ServerStreamListene
private final ServerCallImpl<ReqT, ?> call;
private final ServerCall.Listener<ReqT> listener;
private final Context.CancellableContext context;
private InputStream delayedMessage;

public ServerStreamListenerImpl(
ServerCallImpl<ReqT, ?> call, ServerCall.Listener<ReqT> listener,
Expand Down Expand Up @@ -330,13 +331,27 @@ private void messagesAvailableInternal(final MessageProducer producer) {
InputStream message;
try {
while ((message = producer.next()) != null) {
try {
listener.onMessage(call.method.parseRequest(message));
} catch (Throwable t) {
GrpcUtil.closeQuietly(message);
throw t;
// TODO: Consider forcing this check to be done in the transport (MessageDeframer)
// https://github.com/grpc/grpc-java/pull/13004/changes#r3939373996
if (call.method.getType().clientSendsOneMessage()) {
if (delayedMessage != null) {
GrpcUtil.closeQuietly(message);
call.stream.cancel(Status.INTERNAL.withDescription("Too many requests"));
GrpcUtil.closeQuietly(delayedMessage);
delayedMessage = null;

@ejona86 ejona86 Aug 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This puts the call back into a normal state, so if other events happen after this one (e.g., message, or half close), that could end up propagating to the application before the cancel is processed. I don't know the easiest way to handle that though; obviously we could set some more state/booleans. It is probably worth looking into the exception handling in the executor see what would happen if we throw here.

@kannanjgithub kannanjgithub Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If an exception is thrown from messagesAvailableInternal, the catch block in the wrapped code submitted to the call executor catches it and calls internalClose(t) eventually leading to an asynchronous callback from the transport. It still does not handle the race you mentioned. Instead I'm now invoking closedInternal synchronously when the error is detected. This synchronously sets call.cancelled = true and cancels the context before returning from the executor task.

@ejona86 ejona86 Sep 5, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now doesn't this code call the application's onCancel() twice? It calls it once here with the call to closedInternal(), and then later when the cancellation is actually processed.

@kannanjgithub kannanjgithub Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. All we need to do instead of calling closedInternal() that already notifies the application onCancel is to set call.cancelled = true;. This avoids the double cancellation callback.

call.cancelled = true;
return;
}
delayedMessage = message;
} else {
try {
listener.onMessage(call.method.parseRequest(message));
} catch (Throwable t) {
GrpcUtil.closeQuietly(message);
throw t;
}
message.close();
}
message.close();
}
} catch (Throwable t) {
GrpcUtil.closeQuietly(producer);
Expand All @@ -353,6 +368,19 @@ public void halfClosed() {
return;
}

if (delayedMessage != null) {
InputStream message = delayedMessage;
delayedMessage = null;
try {
listener.onMessage(call.method.parseRequest(message));
} catch (Throwable t) {
GrpcUtil.closeQuietly(message);
Throwables.throwIfUnchecked(t);
throw new RuntimeException(t);
}
GrpcUtil.closeQuietly(message);
}

listener.onHalfClose();
}
}
Expand All @@ -366,6 +394,10 @@ public void closed(Status status) {
}

private void closedInternal(Status status) {
if (delayedMessage != null) {
GrpcUtil.closeQuietly(delayedMessage);
delayedMessage = null;
}
Throwable cancelCause = null;
try {
if (status.isOk()) {
Expand Down
147 changes: 147 additions & 0 deletions core/src/test/java/io/grpc/internal/MessageDeframerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
import static com.google.common.truth.Truth.assertThat;
import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
Expand All @@ -35,9 +37,11 @@
import com.google.common.io.ByteStreams;
import com.google.common.primitives.Bytes;
import io.grpc.Codec;
import io.grpc.Decompressor;
import io.grpc.InternalChannelz.TransportStats;
import io.grpc.StatusRuntimeException;
import io.grpc.StreamTracer;
import io.grpc.internal.MessageDeframer.LazyDecompressingInputStream;
import io.grpc.internal.MessageDeframer.Listener;
import io.grpc.internal.MessageDeframer.SizeEnforcingInputStream;
import io.grpc.internal.testing.TestStreamTracer.TestBaseStreamTracer;
Expand All @@ -52,6 +56,7 @@
import java.util.List;
import java.util.Locale;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.zip.GZIPOutputStream;
import org.junit.Before;
import org.junit.Test;
Expand Down Expand Up @@ -313,6 +318,75 @@ public void compressed() {
verifyNoMoreInteractions(listener);
}

@Test
public void compressed_lazyDecompression() throws IOException {
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
Decompressor countingDecompressor = new Decompressor() {
@Override
public String getMessageEncoding() {
return "gzip";
}

@Override
public InputStream decompress(InputStream is) throws IOException {
decompressCalled.set(true);
return new Codec.Gzip().decompress(is);
}
};

deframer = new MessageDeframer(listener, countingDecompressor, DEFAULT_MAX_MESSAGE_SIZE,
statsTraceCtx, transportTracer);
deframer.request(1);

byte[] payload = compress(new byte[1000]);
byte[] header = new byte[]{1, 0, 0, 0, (byte) payload.length};
deframer.deframe(buffer(Bytes.concat(header, payload)));

verify(listener).messagesAvailable(producer.capture());
InputStream stream = producer.getValue().next();
assertNotNull(stream);

// Decompressor should not be invoked before bytes are read
assertFalse(decompressCalled.get());

// Reading a byte triggers decompression
assertEquals(0, stream.read());
assertTrue(decompressCalled.get());
}

@Test
public void compressed_closeWithoutReading_noDecompression() throws IOException {
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
Decompressor countingDecompressor = new Decompressor() {
@Override
public String getMessageEncoding() {
return "gzip";
}

@Override
public InputStream decompress(InputStream is) throws IOException {
decompressCalled.set(true);
return new Codec.Gzip().decompress(is);
}
};

deframer = new MessageDeframer(listener, countingDecompressor, DEFAULT_MAX_MESSAGE_SIZE,
statsTraceCtx, transportTracer);
deframer.request(1);

byte[] payload = compress(new byte[1000]);
byte[] header = new byte[]{1, 0, 0, 0, (byte) payload.length};
deframer.deframe(buffer(Bytes.concat(header, payload)));

verify(listener).messagesAvailable(producer.capture());
InputStream stream = producer.getValue().next();
assertNotNull(stream);

// Closing without reading should not decompress
stream.close();
assertFalse(decompressCalled.get());
}

@Test
public void deliverIsReentrantSafe() {
doAnswer(
Expand Down Expand Up @@ -493,6 +567,79 @@ public void sizeEnforcingInputStream_markReset() throws IOException {
}
}

@RunWith(JUnit4.class)
public static class LazyDecompressingInputStreamTests {
private TestBaseStreamTracer tracer = new TestBaseStreamTracer();
private StatsTraceContext statsTraceCtx = new StatsTraceContext(new StreamTracer[]{tracer});

@Test
public void lazyDecompressingInputStream_doesNotInitializeUntilRead() throws IOException {
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
Decompressor countingDecompressor = new Decompressor() {
@Override
public String getMessageEncoding() {
return "gzip";
}

@Override
public InputStream decompress(InputStream is) throws IOException {
decompressCalled.set(true);
return new Codec.Gzip().decompress(is);
}
};

ByteArrayInputStream in =
new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8)));
LazyDecompressingInputStream stream = new LazyDecompressingInputStream(
in, 100, statsTraceCtx, countingDecompressor);

assertFalse(decompressCalled.get());
byte[] buf = new byte[5];
int read = stream.read(buf);
assertEquals(5, read);
assertEquals("hello", new String(buf, StandardCharsets.UTF_8));
assertTrue(decompressCalled.get());
stream.close();
}

@Test
public void lazyDecompressingInputStream_closeWithoutRead() throws IOException {
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
final AtomicBoolean inClosed = new AtomicBoolean(false);
Decompressor countingDecompressor = new Decompressor() {
@Override
public String getMessageEncoding() {
return "gzip";
}

@Override
public InputStream decompress(InputStream is) throws IOException {
decompressCalled.set(true);
return new Codec.Gzip().decompress(is);
}
};

ByteArrayInputStream in =
new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8))) {
@Override
public void close() throws IOException {
inClosed.set(true);
super.close();
}
};
LazyDecompressingInputStream stream = new LazyDecompressingInputStream(
in, 100, statsTraceCtx, countingDecompressor);

assertFalse(decompressCalled.get());
stream.close();
assertTrue(inClosed.get());
assertFalse(decompressCalled.get());

// Reading after close should throw IOException
assertThrows(IOException.class, () -> stream.read());
}
}

/**
* Verify stats were published through the tracer.
*
Expand Down
Loading
Loading