Skip to content

Commit f1ba10e

Browse files
kannanjgithubkannanjgithub
authored
core: Delayed deserialization for unary/server-streaming calls (v1.85.x backport) (#13132)
Backport of #13113 to v1.85.x. --- Restore delayed deserialization and lazy decompression for unary and server-streaming calls (originally introduced in #13004 and reverted in #13107). Additionally, convert ServerCall.request(1) to stream.request(2) for unary and server-streaming calls (clientSendsOneMessage()) to prevent transport deadlocks if multiple messages arrive on the stream.
1 parent ca2f8c1 commit f1ba10e

4 files changed

Lines changed: 478 additions & 26 deletions

File tree

‎core/src/main/java/io/grpc/internal/MessageDeframer.java‎

Lines changed: 85 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -437,14 +437,91 @@ private InputStream getCompressedBody() {
437437
.asRuntimeException();
438438
}
439439

440-
try {
441-
// Enforce the maxMessageSize limit on the returned stream.
442-
InputStream unlimitedStream =
443-
decompressor.decompress(ReadableBuffers.openStream(nextFrame, true));
444-
return new SizeEnforcingInputStream(
445-
unlimitedStream, maxInboundMessageSize, statsTraceCtx);
446-
} catch (IOException e) {
447-
throw new RuntimeException(e);
440+
return new LazyDecompressingInputStream(
441+
ReadableBuffers.openStream(nextFrame, true),
442+
maxInboundMessageSize,
443+
statsTraceCtx,
444+
decompressor);
445+
}
446+
447+
/**
448+
* An {@link InputStream} that delays decompressing a compressed frame until data is first read.
449+
*/
450+
@VisibleForTesting
451+
static final class LazyDecompressingInputStream extends FilterInputStream {
452+
private final Decompressor decompressor;
453+
private final int maxMessageSize;
454+
private final StatsTraceContext statsTraceCtx;
455+
private boolean initialized;
456+
private boolean closed;
457+
458+
LazyDecompressingInputStream(
459+
InputStream rawStream,
460+
int maxMessageSize,
461+
StatsTraceContext statsTraceCtx,
462+
Decompressor decompressor) {
463+
super(rawStream);
464+
this.decompressor = decompressor;
465+
this.maxMessageSize = maxMessageSize;
466+
this.statsTraceCtx = statsTraceCtx;
467+
}
468+
469+
private synchronized void ensureInitialized() throws IOException {
470+
if (closed) {
471+
throw new IOException("Stream closed");
472+
}
473+
if (!initialized) {
474+
InputStream decompressed = decompressor.decompress(in);
475+
in = new SizeEnforcingInputStream(decompressed, maxMessageSize, statsTraceCtx);
476+
initialized = true;
477+
}
478+
}
479+
480+
@Override
481+
public int read() throws IOException {
482+
ensureInitialized();
483+
return super.read();
484+
}
485+
486+
@Override
487+
public int read(byte[] b, int off, int len) throws IOException {
488+
ensureInitialized();
489+
return super.read(b, off, len);
490+
}
491+
492+
@Override
493+
public long skip(long n) throws IOException {
494+
ensureInitialized();
495+
return super.skip(n);
496+
}
497+
498+
@Override
499+
public int available() throws IOException {
500+
ensureInitialized();
501+
return super.available();
502+
}
503+
504+
@Override
505+
public synchronized void close() throws IOException {
506+
if (!closed) {
507+
closed = true;
508+
super.close();
509+
}
510+
}
511+
512+
@Override
513+
public synchronized void mark(int readlimit) {
514+
// Mark is not supported
515+
}
516+
517+
@Override
518+
public synchronized void reset() throws IOException {
519+
throw new IOException("Mark not supported");
520+
}
521+
522+
@Override
523+
public boolean markSupported() {
524+
return false;
448525
}
449526
}
450527

‎core/src/main/java/io/grpc/internal/ServerCallImpl.java‎

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,9 @@ final class ServerCallImpl<ReqT, RespT> extends ServerCall<ReqT, RespT> {
9393
public void request(int numMessages) {
9494
try (TaskCloseable ignore = PerfMark.traceTask("ServerCall.request")) {
9595
PerfMark.attachTag(tag);
96+
if (method.getType().clientSendsOneMessage() && numMessages == 1) {
97+
numMessages = 2;
98+
}
9699
stream.request(numMessages);
97100
}
98101
}
@@ -296,6 +299,7 @@ static final class ServerStreamListenerImpl<ReqT> implements ServerStreamListene
296299
private final ServerCallImpl<ReqT, ?> call;
297300
private final ServerCall.Listener<ReqT> listener;
298301
private final Context.CancellableContext context;
302+
private InputStream delayedMessage;
299303

300304
public ServerStreamListenerImpl(
301305
ServerCallImpl<ReqT, ?> call, ServerCall.Listener<ReqT> listener,
@@ -338,13 +342,27 @@ private void messagesAvailableInternal(final MessageProducer producer) {
338342
InputStream message;
339343
try {
340344
while ((message = producer.next()) != null) {
341-
try {
342-
listener.onMessage(call.method.parseRequest(message));
343-
} catch (Throwable t) {
344-
GrpcUtil.closeQuietly(message);
345-
throw t;
345+
// TODO: Consider forcing this check to be done in the transport (MessageDeframer)
346+
// https://github.com/grpc/grpc-java/pull/13004/changes#r3939373996
347+
if (call.method.getType().clientSendsOneMessage()) {
348+
if (delayedMessage != null) {
349+
GrpcUtil.closeQuietly(message);
350+
call.stream.cancel(Status.INTERNAL.withDescription("Too many requests"));
351+
GrpcUtil.closeQuietly(delayedMessage);
352+
delayedMessage = null;
353+
call.cancelled = true;
354+
return;
355+
}
356+
delayedMessage = message;
357+
} else {
358+
try {
359+
listener.onMessage(call.method.parseRequest(message));
360+
} catch (Throwable t) {
361+
GrpcUtil.closeQuietly(message);
362+
throw t;
363+
}
364+
message.close();
346365
}
347-
message.close();
348366
}
349367
} catch (Throwable t) {
350368
GrpcUtil.closeQuietly(producer);
@@ -361,6 +379,19 @@ public void halfClosed() {
361379
return;
362380
}
363381

382+
if (delayedMessage != null) {
383+
InputStream message = delayedMessage;
384+
delayedMessage = null;
385+
try {
386+
listener.onMessage(call.method.parseRequest(message));
387+
} catch (Throwable t) {
388+
GrpcUtil.closeQuietly(message);
389+
Throwables.throwIfUnchecked(t);
390+
throw new RuntimeException(t);
391+
}
392+
GrpcUtil.closeQuietly(message);
393+
}
394+
364395
listener.onHalfClose();
365396
}
366397
}
@@ -374,6 +405,10 @@ public void closed(Status status) {
374405
}
375406

376407
private void closedInternal(Status status) {
408+
if (delayedMessage != null) {
409+
GrpcUtil.closeQuietly(delayedMessage);
410+
delayedMessage = null;
411+
}
377412
Throwable cancelCause = null;
378413
try {
379414
if (status.isOk()) {

‎core/src/test/java/io/grpc/internal/MessageDeframerTest.java‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import static com.google.common.truth.Truth.assertThat;
2020
import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
2121
import static org.junit.Assert.assertEquals;
22+
import static org.junit.Assert.assertFalse;
23+
import static org.junit.Assert.assertNotNull;
2224
import static org.junit.Assert.assertNull;
2325
import static org.junit.Assert.assertThrows;
2426
import static org.junit.Assert.assertTrue;
@@ -35,9 +37,11 @@
3537
import com.google.common.io.ByteStreams;
3638
import com.google.common.primitives.Bytes;
3739
import io.grpc.Codec;
40+
import io.grpc.Decompressor;
3841
import io.grpc.InternalChannelz.TransportStats;
3942
import io.grpc.StatusRuntimeException;
4043
import io.grpc.StreamTracer;
44+
import io.grpc.internal.MessageDeframer.LazyDecompressingInputStream;
4145
import io.grpc.internal.MessageDeframer.Listener;
4246
import io.grpc.internal.MessageDeframer.SizeEnforcingInputStream;
4347
import io.grpc.internal.testing.TestStreamTracer.TestBaseStreamTracer;
@@ -52,6 +56,7 @@
5256
import java.util.List;
5357
import java.util.Locale;
5458
import java.util.concurrent.TimeUnit;
59+
import java.util.concurrent.atomic.AtomicBoolean;
5560
import java.util.zip.GZIPOutputStream;
5661
import org.junit.Before;
5762
import org.junit.Test;
@@ -313,6 +318,75 @@ public void compressed() {
313318
verifyNoMoreInteractions(listener);
314319
}
315320

321+
@Test
322+
public void compressed_lazyDecompression() throws IOException {
323+
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
324+
Decompressor countingDecompressor = new Decompressor() {
325+
@Override
326+
public String getMessageEncoding() {
327+
return "gzip";
328+
}
329+
330+
@Override
331+
public InputStream decompress(InputStream is) throws IOException {
332+
decompressCalled.set(true);
333+
return new Codec.Gzip().decompress(is);
334+
}
335+
};
336+
337+
deframer = new MessageDeframer(listener, countingDecompressor, DEFAULT_MAX_MESSAGE_SIZE,
338+
statsTraceCtx, transportTracer);
339+
deframer.request(1);
340+
341+
byte[] payload = compress(new byte[1000]);
342+
byte[] header = new byte[]{1, 0, 0, 0, (byte) payload.length};
343+
deframer.deframe(buffer(Bytes.concat(header, payload)));
344+
345+
verify(listener).messagesAvailable(producer.capture());
346+
InputStream stream = producer.getValue().next();
347+
assertNotNull(stream);
348+
349+
// Decompressor should not be invoked before bytes are read
350+
assertFalse(decompressCalled.get());
351+
352+
// Reading a byte triggers decompression
353+
assertEquals(0, stream.read());
354+
assertTrue(decompressCalled.get());
355+
}
356+
357+
@Test
358+
public void compressed_closeWithoutReading_noDecompression() throws IOException {
359+
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
360+
Decompressor countingDecompressor = new Decompressor() {
361+
@Override
362+
public String getMessageEncoding() {
363+
return "gzip";
364+
}
365+
366+
@Override
367+
public InputStream decompress(InputStream is) throws IOException {
368+
decompressCalled.set(true);
369+
return new Codec.Gzip().decompress(is);
370+
}
371+
};
372+
373+
deframer = new MessageDeframer(listener, countingDecompressor, DEFAULT_MAX_MESSAGE_SIZE,
374+
statsTraceCtx, transportTracer);
375+
deframer.request(1);
376+
377+
byte[] payload = compress(new byte[1000]);
378+
byte[] header = new byte[]{1, 0, 0, 0, (byte) payload.length};
379+
deframer.deframe(buffer(Bytes.concat(header, payload)));
380+
381+
verify(listener).messagesAvailable(producer.capture());
382+
InputStream stream = producer.getValue().next();
383+
assertNotNull(stream);
384+
385+
// Closing without reading should not decompress
386+
stream.close();
387+
assertFalse(decompressCalled.get());
388+
}
389+
316390
@Test
317391
public void deliverIsReentrantSafe() {
318392
doAnswer(
@@ -493,6 +567,82 @@ public void sizeEnforcingInputStream_markReset() throws IOException {
493567
}
494568
}
495569

570+
@RunWith(JUnit4.class)
571+
public static class LazyDecompressingInputStreamTests {
572+
private TestBaseStreamTracer tracer = new TestBaseStreamTracer();
573+
private StatsTraceContext statsTraceCtx = new StatsTraceContext(new StreamTracer[]{tracer});
574+
575+
@Test
576+
public void lazyDecompressingInputStream_doesNotInitializeUntilRead() throws IOException {
577+
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
578+
Decompressor countingDecompressor = new Decompressor() {
579+
@Override
580+
public String getMessageEncoding() {
581+
return "gzip";
582+
}
583+
584+
@Override
585+
public InputStream decompress(InputStream is) throws IOException {
586+
decompressCalled.set(true);
587+
return new Codec.Gzip().decompress(is);
588+
}
589+
};
590+
591+
ByteArrayInputStream in =
592+
new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8)));
593+
LazyDecompressingInputStream stream = new LazyDecompressingInputStream(
594+
in, 100, statsTraceCtx, countingDecompressor);
595+
596+
assertFalse(stream.markSupported());
597+
stream.mark(10);
598+
assertThrows(IOException.class, () -> stream.reset());
599+
assertFalse(decompressCalled.get());
600+
byte[] buf = new byte[5];
601+
int read = stream.read(buf);
602+
assertEquals(5, read);
603+
assertEquals("hello", new String(buf, StandardCharsets.UTF_8));
604+
assertTrue(decompressCalled.get());
605+
stream.close();
606+
}
607+
608+
@Test
609+
public void lazyDecompressingInputStream_closeWithoutRead() throws IOException {
610+
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
611+
final AtomicBoolean inClosed = new AtomicBoolean(false);
612+
Decompressor countingDecompressor = new Decompressor() {
613+
@Override
614+
public String getMessageEncoding() {
615+
return "gzip";
616+
}
617+
618+
@Override
619+
public InputStream decompress(InputStream is) throws IOException {
620+
decompressCalled.set(true);
621+
return new Codec.Gzip().decompress(is);
622+
}
623+
};
624+
625+
ByteArrayInputStream in =
626+
new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8))) {
627+
@Override
628+
public void close() throws IOException {
629+
inClosed.set(true);
630+
super.close();
631+
}
632+
};
633+
LazyDecompressingInputStream stream = new LazyDecompressingInputStream(
634+
in, 100, statsTraceCtx, countingDecompressor);
635+
636+
assertFalse(decompressCalled.get());
637+
stream.close();
638+
assertTrue(inClosed.get());
639+
assertFalse(decompressCalled.get());
640+
641+
// Reading after close should throw IOException
642+
assertThrows(IOException.class, () -> stream.read());
643+
}
644+
}
645+
496646
/**
497647
* Verify stats were published through the tracer.
498648
*

0 commit comments

Comments
 (0)