feat: support maxSize in PbjGrpcClient (#753)

Signed-off-by: Anthony Petrov <anthony@swirldslabs.com>

Author: anthony-swirldslabs

pbj-core/pbj-grpc-client-helidon/src/main/java/com/hedera/pbj/grpc/client/helidon/PbjGrpcCall.java

@@ -208,7 +208,12 @@ private void receiveRepliesLoop() {
                                 datagram.compressedFlag() == 1 ? decompressor.decompress(bytes) : bytes;
 
                         try {
-                            final ReplyT reply = replyCodec.parse(replyBytes);
+                            final ReplyT reply = replyCodec.parse(
+                                    replyBytes.toReadableSequentialData(),
+                                    false,
+                                    false,
+                                    Codec.DEFAULT_MAX_DEPTH,
+                                    grpcClient.getConfig().maxSize());
                             pipeline.onNext(reply);
                         } catch (ParseException e) {
                             pipeline.onError(e);

pbj-core/pbj-grpc-client-helidon/src/main/java/com/hedera/pbj/grpc/client/helidon/PbjGrpcClientConfig.java

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.hedera.pbj.grpc.client.helidon;
 
+import com.hedera.pbj.runtime.Codec;
 import com.hedera.pbj.runtime.grpc.GrpcCompression;
 import io.helidon.common.tls.Tls;
 import java.time.Duration;
@@ -9,6 +10,7 @@
 
 /**
  * Configuration for PBJ GRPC client.
+ * @param maxSize the maximum size of messages that the client is able to receive, defaults to Codec.DEFAULT_MAX_SIZE.
  */
 public record PbjGrpcClientConfig(
         /** A read timeout. Duration.ofSeconds(10) is a good default. */
@@ -34,15 +36,29 @@ public record PbjGrpcClientConfig(
          * e.g. "identity", "gzip", etc.
          * Note that the encoding must be registered as a `Decompressor` with `GrpcCompression` to actually be supported.
          */
-        Set<String> acceptEncodings) {
+        Set<String> acceptEncodings,
+        int maxSize) {
 
+    /** For backward compatibility before encodings were introduced. */
     public PbjGrpcClientConfig(Duration readTimeout, Tls tls, Optional<String> authority, String contentType) {
         this(
                 readTimeout,
                 tls,
                 authority,
                 contentType,
                 GrpcCompression.IDENTITY,
-                GrpcCompression.getDecompressorNames());
+                GrpcCompression.getDecompressorNames(),
+                Codec.DEFAULT_MAX_SIZE);
+    }
+
+    /** For backward compatibility before maxSize was introduced. */
+    public PbjGrpcClientConfig(
+            Duration readTimeout,
+            Tls tls,
+            Optional<String> authority,
+            String contentType,
+            String encoding,
+            Set<String> acceptEncodings) {
+        this(readTimeout, tls, authority, contentType, encoding, acceptEncodings, Codec.DEFAULT_MAX_SIZE);
     }
 }

pbj-core/pbj-grpc-client-helidon/src/test/java/com/hedera/pbj/grpc/client/helidon/PbjGrpcCallTest.java

@@ -21,6 +21,7 @@
 import com.hedera.pbj.runtime.grpc.GrpcStatus;
 import com.hedera.pbj.runtime.grpc.Pipeline;
 import com.hedera.pbj.runtime.grpc.ServiceInterface;
+import com.hedera.pbj.runtime.io.ReadableSequentialData;
 import com.hedera.pbj.runtime.io.buffer.Bytes;
 import io.helidon.common.buffers.BufferData;
 import io.helidon.common.tls.Tls;
@@ -269,7 +270,14 @@ public void testReceiveRepliesLoopSingleReply(final boolean isTimeout) throws Ex
         doReturn(bufferData).when(data).data();
 
         final Object reply = mock(Object.class);
-        doReturn(reply).when(replyCodec).parse(eq(Bytes.wrap(new byte[] {6})));
+        doReturn(reply)
+                .when(replyCodec)
+                .parse(
+                        any(ReadableSequentialData.class),
+                        eq(false),
+                        eq(false),
+                        eq(Codec.DEFAULT_MAX_DEPTH),
+                        eq(Codec.DEFAULT_MAX_SIZE));
 
         runnable.run();
 
@@ -324,7 +332,14 @@ public void testReceiveRepliesLoopParseException() throws Exception {
         doReturn(bufferData).when(data).data();
 
         final ParseException exception = new ParseException("test");
-        doThrow(exception).when(replyCodec).parse(eq(Bytes.wrap(new byte[] {6})));
+        doThrow(exception)
+                .when(replyCodec)
+                .parse(
+                        any(ReadableSequentialData.class),
+                        eq(false),
+                        eq(false),
+                        eq(Codec.DEFAULT_MAX_DEPTH),
+                        eq(Codec.DEFAULT_MAX_SIZE));
 
         runnable.run();
 

pbj-integration-tests/src/jmh/java/com/hedera/pbj/integration/jmh/grpc/PbjGrpcBench.java

@@ -5,6 +5,7 @@
 import com.hedera.pbj.grpc.helidon.PbjRouting;
 import com.hedera.pbj.integration.grpc.GrpcTestUtils;
 import com.hedera.pbj.integration.grpc.PortsAllocator;
+import com.hedera.pbj.runtime.Codec;
 import com.hedera.pbj.runtime.grpc.GrpcClient;
 import com.hedera.pbj.runtime.grpc.Pipeline;
 import com.hedera.pbj.runtime.grpc.ServiceInterface;
@@ -78,8 +79,8 @@ static GreeterInterface.GreeterClient createClient(final int port, final String[
         if (encodings == null || encodings.length == 0) {
             grpcClient = GrpcTestUtils.createGrpcClient(port, GrpcTestUtils.PROTO_OPTIONS);
         } else {
-            grpcClient =
-                    GrpcTestUtils.createGrpcClient(port, GrpcTestUtils.PROTO_OPTIONS, encodings[0], Set.of(encodings));
+            grpcClient = GrpcTestUtils.createGrpcClient(
+                    port, GrpcTestUtils.PROTO_OPTIONS, encodings[0], Set.of(encodings), Codec.DEFAULT_MAX_SIZE);
         }
 
         return new GreeterInterface.GreeterClient(grpcClient, GrpcTestUtils.PROTO_OPTIONS);

pbj-integration-tests/src/main/java/com/hedera/pbj/integration/grpc/GrpcTestUtils.java

@@ -3,6 +3,7 @@
 
 import com.hedera.pbj.grpc.client.helidon.PbjGrpcClient;
 import com.hedera.pbj.grpc.client.helidon.PbjGrpcClientConfig;
+import com.hedera.pbj.runtime.Codec;
 import com.hedera.pbj.runtime.grpc.GrpcClient;
 import com.hedera.pbj.runtime.grpc.GrpcCompression;
 import com.hedera.pbj.runtime.grpc.ServiceInterface;
@@ -35,14 +36,20 @@ public static void sleep(final GrpcClient grpcClient) {
     }
 
     public static GrpcClient createGrpcClient(final int port, final ServiceInterface.RequestOptions requestOptions) {
-        return createGrpcClient(port, requestOptions, GrpcCompression.IDENTITY, GrpcCompression.getDecompressorNames());
+        return createGrpcClient(
+                port,
+                requestOptions,
+                GrpcCompression.IDENTITY,
+                GrpcCompression.getDecompressorNames(),
+                Codec.DEFAULT_MAX_SIZE);
     }
 
     public static GrpcClient createGrpcClient(
             final int port,
             final ServiceInterface.RequestOptions requestOptions,
             String encoding,
-            Set<String> acceptEncodings) {
+            Set<String> acceptEncodings,
+            int maxSize) {
         final Tls tls = Tls.builder().enabled(false).build();
         final WebClient webClient =
                 WebClient.builder().baseUri("http://localhost:" + port).tls(tls).build();
@@ -52,7 +59,7 @@ public static GrpcClient createGrpcClient(
                 requestOptions.authority().isPresent() ? requestOptions.authority() : Optional.of("localhost:" + port);
 
         final PbjGrpcClientConfig config = new PbjGrpcClientConfig(
-                READ_TIMEOUT, tls, authority, requestOptions.contentType(), encoding, acceptEncodings);
+                READ_TIMEOUT, tls, authority, requestOptions.contentType(), encoding, acceptEncodings, maxSize);
 
         return new PbjGrpcClient(webClient, config);
     }

pbj-integration-tests/src/test/java/com/hedera/pbj/integration/test/grpc/GrpcClientComprehensiveTest.java

@@ -7,7 +7,9 @@
 
 import com.hedera.pbj.integration.grpc.GrpcTestUtils;
 import com.hedera.pbj.integration.grpc.PortsAllocator;
+import com.hedera.pbj.runtime.Codec;
 import com.hedera.pbj.runtime.grpc.GrpcClient;
+import com.hedera.pbj.runtime.grpc.GrpcCompression;
 import com.hedera.pbj.runtime.grpc.GrpcException;
 import com.hedera.pbj.runtime.grpc.Pipeline;
 import java.util.ArrayList;
@@ -108,6 +110,32 @@ void testUnaryMethodHappyCase() {
         }
     }
 
+    @Test
+    void testUnaryMethodReceivingExtraLargePayload() {
+        try (final PortsAllocator.Port port = GrpcTestUtils.PORTS.acquire();
+                final GrpcServerGreeterHandle server = serverFactory.apply(port.port())) {
+            server.start();
+            final String extraLongString = "a".repeat(Codec.DEFAULT_MAX_SIZE * 2);
+            server.setSayHello(
+                    request -> HelloReply.newBuilder().message(extraLongString).build());
+
+            final GrpcClient grpcClient = GrpcTestUtils.createGrpcClient(
+                    port.port(),
+                    GrpcTestUtils.PROTO_OPTIONS,
+                    GrpcCompression.IDENTITY,
+                    GrpcCompression.getDecompressorNames(),
+                    Codec.DEFAULT_MAX_SIZE * 2);
+            final GreeterInterface.GreeterClient client =
+                    new GreeterInterface.GreeterClient(grpcClient, GrpcTestUtils.PROTO_OPTIONS);
+
+            final HelloRequest request =
+                    HelloRequest.newBuilder().name("test name").build();
+            final HelloReply reply = client.sayHello(request);
+
+            assertEquals(extraLongString, reply.message());
+        }
+    }
+
     @Test
     void testServerStreamingMethodHappyCase() {
         try (final PortsAllocator.Port port = GrpcTestUtils.PORTS.acquire();

pbj-integration-tests/src/test/java/com/hedera/pbj/integration/test/grpc/GrpcCompressionTest.java

@@ -7,6 +7,7 @@
 import com.hedera.pbj.grpc.helidon.PbjRouting;
 import com.hedera.pbj.integration.grpc.GrpcTestUtils;
 import com.hedera.pbj.integration.grpc.PortsAllocator;
+import com.hedera.pbj.runtime.Codec;
 import com.hedera.pbj.runtime.grpc.GrpcClient;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
@@ -28,8 +29,8 @@ void testPbjGrpcClientCompressionWithGoogleGrpcServer() {
             server.setSayHello(request ->
                     HelloReply.newBuilder().message("Hello " + request.name()).build());
 
-            try (final GrpcClient grpcClient =
-                    GrpcTestUtils.createGrpcClient(port.port(), GrpcTestUtils.PROTO_OPTIONS, "gzip", Set.of("gzip"))) {
+            try (final GrpcClient grpcClient = GrpcTestUtils.createGrpcClient(
+                    port.port(), GrpcTestUtils.PROTO_OPTIONS, "gzip", Set.of("gzip"), Codec.DEFAULT_MAX_SIZE)) {
                 final GreeterInterface.GreeterClient client =
                         new GreeterInterface.GreeterClient(grpcClient, GrpcTestUtils.PROTO_OPTIONS);
 
@@ -107,8 +108,8 @@ void testPbjGrpcServerCompressionWithPbjClient() {
             handle.setSayHello(request ->
                     HelloReply.newBuilder().message("Hello " + request.name()).build());
 
-            try (final GrpcClient grpcClient =
-                    GrpcTestUtils.createGrpcClient(port.port(), GrpcTestUtils.PROTO_OPTIONS, "gzip", Set.of("gzip"))) {
+            try (final GrpcClient grpcClient = GrpcTestUtils.createGrpcClient(
+                    port.port(), GrpcTestUtils.PROTO_OPTIONS, "gzip", Set.of("gzip"), Codec.DEFAULT_MAX_SIZE)) {
                 final GreeterInterface.GreeterClient client =
                         new GreeterInterface.GreeterClient(grpcClient, GrpcTestUtils.PROTO_OPTIONS);