core: Delayed deserialization for unary/server-streaming calls (v1.84.x backport)

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

@@ -437,14 +437,91 @@ 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) {
+      // Mark is not supported
+    }
+
+    @Override
+    public synchronized void reset() throws IOException {
+      throw new IOException("Mark not supported");
+    }
+
+    @Override
+    public boolean markSupported() {
+      return false;
     }
   }
 

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

@@ -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,
@@ -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;
+              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);
@@ -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();
       }
     }
@@ -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()) {

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

@@ -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;
@@ -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;
@@ -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;
@@ -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(
@@ -493,6 +567,82 @@ 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(stream.markSupported());
+      stream.mark(10);
+      assertThrows(IOException.class, () -> stream.reset());
+      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.
    *