Skip to content

Commit f237d2f

Browse files
kannanjgithubkannanjgithub
committed
Address review comments.
1 parent 50197f1 commit f237d2f

4 files changed

Lines changed: 288 additions & 136 deletions

File tree

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

Lines changed: 96 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -437,14 +437,102 @@ 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+
try {
515+
ensureInitialized();
516+
super.mark(readlimit);
517+
} catch (IOException e) {
518+
throw new RuntimeException(e);
519+
}
520+
}
521+
522+
@Override
523+
public synchronized void reset() throws IOException {
524+
ensureInitialized();
525+
super.reset();
526+
}
527+
528+
@Override
529+
public boolean markSupported() {
530+
try {
531+
ensureInitialized();
532+
return super.markSupported();
533+
} catch (IOException e) {
534+
throw new RuntimeException(e);
535+
}
448536
}
449537
}
450538

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

Lines changed: 4 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import io.grpc.CompressorRegistry;
3535
import io.grpc.Context;
3636
import io.grpc.DecompressorRegistry;
37-
import io.grpc.Detachable;
3837
import io.grpc.InternalDecompressorRegistry;
3938
import io.grpc.InternalStatus;
4039
import io.grpc.Metadata;
@@ -46,9 +45,6 @@
4645
import io.perfmark.PerfMark;
4746
import io.perfmark.Tag;
4847
import io.perfmark.TaskCloseable;
49-
import java.io.ByteArrayInputStream;
50-
import java.io.ByteArrayOutputStream;
51-
import java.io.IOException;
5248
import java.io.InputStream;
5349
import java.util.logging.Level;
5450
import java.util.logging.Logger;
@@ -325,20 +321,6 @@ public void messagesAvailable(MessageProducer producer) {
325321
}
326322
}
327323

328-
private static InputStream bufferMessage(InputStream is) throws IOException {
329-
if (is instanceof Detachable) {
330-
return ((Detachable) is).detach();
331-
}
332-
// Fallback: copy to byte array
333-
ByteArrayOutputStream baos = new ByteArrayOutputStream();
334-
byte[] buffer = new byte[4096];
335-
int bytesRead;
336-
while ((bytesRead = is.read(buffer)) != -1) {
337-
baos.write(buffer, 0, bytesRead);
338-
}
339-
return new ByteArrayInputStream(baos.toByteArray());
340-
}
341-
342324
@SuppressWarnings("Finally") // The code avoids suppressing the exception thrown from try
343325
private void messagesAvailableInternal(final MessageProducer producer) {
344326
if (call.cancelled) {
@@ -349,22 +331,18 @@ private void messagesAvailableInternal(final MessageProducer producer) {
349331
InputStream message;
350332
try {
351333
while ((message = producer.next()) != null) {
334+
// TODO: Consider forcing this check to be done in the transport (MessageDeframer)
335+
// https://github.com/grpc/grpc-java/pull/13004/changes#r3939373996
352336
if (call.method.getType().clientSendsOneMessage()) {
353337
if (delayedMessage != null) {
354338
GrpcUtil.closeQuietly(message);
355339
call.stream.cancel(Status.INTERNAL.withDescription("Too many requests"));
356340
GrpcUtil.closeQuietly(delayedMessage);
357341
delayedMessage = null;
358-
closedInternal(Status.INTERNAL.withDescription("Too many requests"));
342+
call.cancelled = true;
359343
return;
360344
}
361-
try {
362-
delayedMessage = bufferMessage(message);
363-
} catch (Throwable t) {
364-
GrpcUtil.closeQuietly(message);
365-
throw t;
366-
}
367-
message.close();
345+
delayedMessage = message;
368346
} else {
369347
try {
370348
listener.onMessage(call.method.parseRequest(message));

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

Lines changed: 147 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,79 @@ 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(decompressCalled.get());
597+
byte[] buf = new byte[5];
598+
int read = stream.read(buf);
599+
assertEquals(5, read);
600+
assertEquals("hello", new String(buf, StandardCharsets.UTF_8));
601+
assertTrue(decompressCalled.get());
602+
stream.close();
603+
}
604+
605+
@Test
606+
public void lazyDecompressingInputStream_closeWithoutRead() throws IOException {
607+
final AtomicBoolean decompressCalled = new AtomicBoolean(false);
608+
final AtomicBoolean inClosed = new AtomicBoolean(false);
609+
Decompressor countingDecompressor = new Decompressor() {
610+
@Override
611+
public String getMessageEncoding() {
612+
return "gzip";
613+
}
614+
615+
@Override
616+
public InputStream decompress(InputStream is) throws IOException {
617+
decompressCalled.set(true);
618+
return new Codec.Gzip().decompress(is);
619+
}
620+
};
621+
622+
ByteArrayInputStream in =
623+
new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8))) {
624+
@Override
625+
public void close() throws IOException {
626+
inClosed.set(true);
627+
super.close();
628+
}
629+
};
630+
LazyDecompressingInputStream stream = new LazyDecompressingInputStream(
631+
in, 100, statsTraceCtx, countingDecompressor);
632+
633+
assertFalse(decompressCalled.get());
634+
stream.close();
635+
assertTrue(inClosed.get());
636+
assertFalse(decompressCalled.get());
637+
638+
// Reading after close should throw IOException
639+
assertThrows(IOException.class, () -> stream.read());
640+
}
641+
}
642+
496643
/**
497644
* Verify stats were published through the tracer.
498645
*

0 commit comments

Comments
 (0)