perf: benchmark zstd GRPC compression (#749)

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

Author: anthony-swirldslabs

pbj-core/gradle/modules.properties

@@ -1,3 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
 # Module to Maven artifact mappings for JPMS
 io.helidon.common.features.codegen=io.helidon.common.features:helidon-common-features-codegen
+com.github.luben.zstd_jni=com.github.luben:zstd-jni

pbj-core/hiero-dependency-versions/build.gradle.kts

@@ -54,6 +54,7 @@ dependencies.constraints {
         excludes.add("org.antlr:antlr4")
         excludes.add("com.google.protobuf:protoc")
         excludes.add("io.grpc:protoc-gen-grpc-java")
+        excludes.add("com.github.luben:zstd-jni")
     }
 
     // Testing only
@@ -77,4 +78,5 @@ dependencies.constraints {
     api("io.helidon.webclient:helidon-webclient-http2:$helidon") {
         because("io.helidon.webclient.http2")
     }
+    api("com.github.luben:zstd-jni:1.5.5-11") { because("com.github.luben.zstd") }
 }

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

@@ -52,6 +52,24 @@ public class PbjGrpcCall<RequestT, ReplyT> implements GrpcCall<RequestT, ReplyT>
     private static final HeaderName GRPC_ENCODING = HeaderNames.createFromLowercase("grpc-encoding");
     private static final HeaderName GRPC_ACCEPT_ENCODING = HeaderNames.createFromLowercase("grpc-accept-encoding");
 
+    private static final PbjGrpcNetworkBytesInspector NO_OP_NETWORK_BYTES_INSPECTOR =
+            new PbjGrpcNetworkBytesInspector() {};
+
+    private static PbjGrpcNetworkBytesInspector networkBytesInspector = NO_OP_NETWORK_BYTES_INSPECTOR;
+
+    /**
+     * Install a PbjGrpcNetworkBytesInspector, which may be null to reset it to no-op.
+     * This is an internal API that isn't suited for general-purpose applications.
+     * See `PbjGrpcNetworkBytesInspector` javadoc for more details.
+     * This method is not thread-safe and relies on eventual consistency to take effect. So it's best to invoke it early
+     * in the application startup.
+     * @param networkBytesInspector a PbjGrpcNetworkBytesInspector instance
+     */
+    public static void setNetworkBytesInspector(PbjGrpcNetworkBytesInspector networkBytesInspector) {
+        PbjGrpcCall.networkBytesInspector =
+                networkBytesInspector != null ? networkBytesInspector : NO_OP_NETWORK_BYTES_INSPECTOR;
+    }
+
     private final PbjGrpcClient grpcClient;
     private final Codec<RequestT> requestCodec;
     private final Codec<ReplyT> replyCodec;
@@ -123,6 +141,7 @@ public class PbjGrpcCall<RequestT, ReplyT> implements GrpcCall<RequestT, ReplyT>
     public void sendRequest(final RequestT request, final boolean endOfStream) {
         final Bytes requestBytes = requestCodec.toBytes(request);
         final Bytes bytes = GrpcCompression.getCompressor(grpcOutgoingEncoding).compress(requestBytes);
+        PbjGrpcCall.networkBytesInspector.sent(bytes);
         final BufferData bufferData =
                 BufferData.create(PbjGrpcDatagramReader.PREFIX_LENGTH + Math.toIntExact(bytes.length()));
 
@@ -202,6 +221,7 @@ private void receiveRepliesLoop() {
                         BufferData data = datagram.data();
                         final byte[] array = data.readBytes();
                         final Bytes bytes = Bytes.wrap(array);
+                        PbjGrpcCall.networkBytesInspector.received(bytes);
                         // If the compressedFlag is 0, then per the specification, the message isn't compressed
                         // regardless of the grpc-encoding value:
                         // https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md

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

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: Apache-2.0
+package com.hedera.pbj.grpc.client.helidon;
+
+import com.hedera.pbj.runtime.io.buffer.Bytes;
+
+/**
+ * An interface for inspecting raw bytes of data payloads sent/received over GRPC.
+ * <p>
+ * The bytes are encoded/compressed if the GRPC encoding isn't `identity`. The bytes only represent the data payload
+ * and don't include any system, non-user bytes, such as GRPC datagram size/compression flag or any headers sent
+ * over the GRPC/HTTP2 connection.
+ * <p>
+ * The primary purpose of this interface is to help estimate the amount of user data being transferred over
+ * a network connection during GRPC interactions. Currently, this interface doesn't allow one to associate
+ * a particular Bytes object with a specific GRPC request. In the future, this interface may be extended or
+ * even modified to allow for such association. Therefore, before using this interface, application developers
+ * should very carefully consider if they're ready to keep up with the interface changes and update their
+ * applications accordingly.
+ * <p>
+ * It is STRONGLY RECOMMENDED to consider this interface being a PBJ internal implementation detail that is
+ * not generally suited for use in regular applications.
+ */
+public interface PbjGrpcNetworkBytesInspector {
+    /**
+     * Inspect bytes to be sent to a remote peer, after compressing them and prior to actually sending them.
+     * @param bytes the bytes sent, potentially compressed.
+     */
+    default void sent(Bytes bytes) {}
+
+    /**
+     * Inspect bytes just received from a remote peer, prior to decompressing them and passing them to the application.
+     * @param bytes the bytes sent, potentially compressed.
+     */
+    default void received(Bytes bytes) {}
+}

pbj-core/pbj-runtime/src/main/java/com/hedera/pbj/runtime/grpc/GrpcCompression.java

@@ -7,6 +7,8 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -84,16 +86,30 @@ public Bytes decompress(Bytes bytes) {
         }
     }
 
-    private static final Map<String, Compressor> COMPRESSOR_MAP = Map.of(
-            IdentityGrpcTransformer.NAME, IdentityGrpcTransformer.INSTANCE,
-            GzipGrpcTransformer.NAME, GzipGrpcTransformer.INSTANCE);
-    private static final Map<String, Decompressor> DECOMPRESSOR_MAP = Map.of(
-            IdentityGrpcTransformer.NAME, IdentityGrpcTransformer.INSTANCE,
-            GzipGrpcTransformer.NAME, GzipGrpcTransformer.INSTANCE);
+    private static final Map<String, Compressor> COMPRESSOR_MAP = new HashMap<>();
+    private static final Map<String, Decompressor> DECOMPRESSOR_MAP = new HashMap<>();
+
+    /** Register a Compressor, potentially overwriting an existing registration for `name`. */
+    public static synchronized void registerCompressor(String name, Compressor compressor) {
+        COMPRESSOR_MAP.put(name, compressor);
+    }
+
+    /** Register a Decompressor, potentially overwriting an existing registration for `name`. */
+    public static synchronized void registerDecompressor(String name, Decompressor decompressor) {
+        DECOMPRESSOR_MAP.put(name, decompressor);
+    }
+
+    static {
+        registerCompressor(IdentityGrpcTransformer.NAME, IdentityGrpcTransformer.INSTANCE);
+        registerCompressor(GzipGrpcTransformer.NAME, GzipGrpcTransformer.INSTANCE);
+
+        registerDecompressor(IdentityGrpcTransformer.NAME, IdentityGrpcTransformer.INSTANCE);
+        registerDecompressor(GzipGrpcTransformer.NAME, GzipGrpcTransformer.INSTANCE);
+    }
 
     /** Return names of all known compressors. */
     public static Set<String> getCompressorNames() {
-        return COMPRESSOR_MAP.keySet();
+        return Collections.unmodifiableSet(COMPRESSOR_MAP.keySet());
     }
 
     /** Return a known Compressor by its name, or null if unknown. */
@@ -103,7 +119,7 @@ public static Compressor getCompressor(String name) {
 
     /** Return names of all known decompressors. */
     public static Set<String> getDecompressorNames() {
-        return DECOMPRESSOR_MAP.keySet();
+        return Collections.unmodifiableSet(DECOMPRESSOR_MAP.keySet());
     }
 
     /** Return a known Decompressor by its name, or null if unknown. */
@@ -129,7 +145,7 @@ public static Decompressor determineDecompressor(List<String> encodingList) {
             } else {
                 throw new IllegalStateException(
                         "GRPC peer didn't provide grpc-encoding header and 'identity' is unsupported, only the following are supported: "
-                                + GrpcCompression.getDecompressorNames());
+                                + DECOMPRESSOR_MAP.keySet());
             }
         } else if (encodingList.size() > 1) {
             throw new IllegalStateException("GRPC peer specified multiple encodings at once: " + encodingList);
@@ -142,7 +158,7 @@ public static Decompressor determineDecompressor(List<String> encodingList) {
                 return GrpcCompression.getDecompressor(encoding);
             } else {
                 throw new IllegalStateException("GRPC peer uses an unsupported encoding: '" + encoding
-                        + "' while only the following are supported: " + GrpcCompression.getDecompressorNames());
+                        + "' while only the following are supported: " + DECOMPRESSOR_MAP.keySet());
             }
         }
     }
@@ -155,8 +171,8 @@ public static Decompressor determineDecompressor(List<String> encodingList) {
      */
     public static String determineCompressorName(List<String> acceptEncoding, String encoding) {
         final List<String> supportedAcceptEncodings = acceptEncoding.stream()
-                .filter(ae -> GrpcCompression.getCompressorNames().stream()
-                        .anyMatch(sae -> ae.equals(sae) || ae.startsWith(sae + ";")))
+                .filter(ae ->
+                        COMPRESSOR_MAP.keySet().stream().anyMatch(sae -> ae.equals(sae) || ae.startsWith(sae + ";")))
                 .toList();
 
         if (supportedAcceptEncodings.isEmpty()) {

pbj-integration-tests/build.gradle.kts

@@ -56,6 +56,7 @@ jmhModuleInfo {
     requires("com.google.protobuf.util")
     requires("io.helidon.common")
     requires("io.helidon.webserver")
+    requires("com.github.luben.zstd_jni")
 }
 
 // version is added to module-info.class files
@@ -69,7 +70,10 @@ configurations.testRuntimeClasspath {
 }
 
 // IMPROVE: Test code should not have a direct dependency to 'com.hedera.pbj.compiler'
-dependencies { testImplementation("com.hedera.pbj:pbj-compiler") { isTransitive = false } }
+dependencies {
+    jmhImplementation("com.github.luben:zstd-jni")
+    testImplementation("com.hedera.pbj:pbj-compiler") { isTransitive = false }
+}
 
 dependencyAnalysis { issues { all { onAny { exclude("com.hedera.pbj:pbj-compiler") } } } }
 

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

@@ -5,6 +5,7 @@
 import edu.umd.cs.findbugs.annotations.NonNull;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.Flow;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
@@ -19,6 +20,31 @@ public class GreeterService implements GreeterInterface {
     static final HelloRequest EMPTY_REQUEST = new HelloRequest("");
     static final HelloReply EMPTY_REPLY = new HelloReply("");
 
+    // For generating random payload strings. Note that benchmarks run in multiple threads,
+    // so it's impossible to make them deterministic (or it'd be too very slow.)
+    static final Random RANDOM = new Random();
+
+    /**
+     * Generate half-random/half-repeated string to allow compressors to compress something.
+     * @param length total length of the string, should be even.
+     * @return a semi-random string
+     */
+    static String generateHalfRandomString(int length) {
+        StringBuilder sb = new StringBuilder(length);
+        sb.append("a".repeat(length / 2));
+        for (int i = length / 2; i < length; i++) {
+            sb.append((char) (32 + RANDOM.nextInt(127 - 32)));
+        }
+        return sb.toString();
+    }
+
+    // We want our benchmark to measure the GRPC implementation performance. We don't want to include the string
+    // generation code into the time it takes to process a request. So let's precompute the strings here:
+    static final String NORMAL_PAYLOAD_STRING = generateHalfRandomString(256);
+    static final String HEAVY_PAYLOAD_STRING = generateHalfRandomString(8192);
+    // PBJ DEFAULT_MAX_SIZE is 2MB. See https://github.com/hashgraph/pbj/issues/748
+    static final String SUPER_PAYLOAD_STRING = generateHalfRandomString(2000 * 1024);
+
     private final PayloadWeight weight;
     private final int streamCount;
 

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

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: Apache-2.0
+package com.hedera.pbj.integration.jmh.grpc;
+
+import com.hedera.pbj.grpc.client.helidon.PbjGrpcCall;
+import com.hedera.pbj.grpc.client.helidon.PbjGrpcNetworkBytesInspector;
+import com.hedera.pbj.runtime.io.buffer.Bytes;
+
+/**
+ * Simulate network latency based on the actual amount of user data being sent/received over a network.
+ * Note: with 1Gbps, we get 8ns per byte. So to test higher speeds, we should switch to floating point math.
+ * A very fast network would make any compression look bad because we'll waste CPU time on the compression
+ * instead of leveraging the fast network.
+ * A slow network would show that compression may be useful sometimes. Specifically, for larger, compressible
+ * payloads. Smaller payloads (<8K) never benefit from compression.
+ * <p>
+ * The network latency simulator uses the PbjGrpcNetworkBytesInspector feature of the PbjGrpcCall, which is
+ * a static entity, meaning that it affects any and all PbjGrpcCall objects in the same JVM. This is by design
+ * for performance reasons. However, users of this Simulator should be aware of this and not try to use it
+ * from multiple threads with different latency parameters at once.
+ */
+public class NetworkLatencySimulator {
+    // ms 1e-3, us 1e-6, ns 1e-9:
+    private static final long NANOS_IN_MILLI = 1_000_000L;
+
+    /**
+     * Install the NetworkLatencySimulator as a PbjGrpcNetworkBytesInspector in PbjGrpcCall.
+     * @param networkSpeedMbitPerSecond the speed in Mbps, e.g. 1_000 for 1Gbps network
+     * @param printSizes if true, print a few sent/received sizes for debugging/additional information
+     */
+    public static void simulate(final long networkSpeedMbitPerSecond, final boolean printSizes) {
+        // mbit->mbyte = /8:
+        final long nanosPerByte = 1_000_000_000L * 8 / (networkSpeedMbitPerSecond * 1_000_000L);
+        PbjGrpcCall.setNetworkBytesInspector(new PbjGrpcNetworkBytesInspector() {
+            // max number of sizes to print if enabled:
+            int counter = 4;
+
+            private void sleep(long bytes) {
+                final long nanos = nanosPerByte * bytes;
+                final long deadline = System.nanoTime() + nanos;
+                while (System.nanoTime() < deadline) {
+                    Thread.onSpinWait();
+                }
+            }
+
+            @Override
+            public void sent(Bytes bytes) {
+                if (printSizes && counter-- >= 0) {
+                    System.err.println("sent: " + bytes.length() + " bytes");
+                }
+                sleep(bytes.length());
+            }
+
+            @Override
+            public void received(Bytes bytes) {
+                if (printSizes && counter-- >= 0) {
+                    System.err.println("received: " + bytes.length() + " bytes");
+                }
+                sleep(bytes.length());
+            }
+        });
+    }
+}

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

@@ -11,8 +11,9 @@
  */
 public enum PayloadWeight {
     LIGHT(() -> GreeterService.EMPTY_REQUEST, request -> GreeterService.EMPTY_REPLY),
-    NORMAL(() -> new HelloRequest("a".repeat(256)), request -> new HelloReply(request.name())),
-    HEAVY(() -> new HelloRequest("a".repeat(8192)), request -> new HelloReply(request.name()));
+    NORMAL(() -> new HelloRequest(GreeterService.NORMAL_PAYLOAD_STRING), request -> new HelloReply(request.name())),
+    HEAVY(() -> new HelloRequest(GreeterService.HEAVY_PAYLOAD_STRING), request -> new HelloReply(request.name())),
+    SUPER(() -> new HelloRequest(GreeterService.SUPER_PAYLOAD_STRING), request -> new HelloReply(request.name()));
 
     public final Supplier<HelloRequest> requestSupplier;
     public final Function<HelloRequest, HelloReply> replyProvider;

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

@@ -3,6 +3,7 @@
 
 import com.hedera.pbj.grpc.helidon.PbjGrpcServiceConfig;
 import com.hedera.pbj.grpc.helidon.PbjRouting;
+import com.hedera.pbj.grpc.helidon.config.PbjConfig;
 import com.hedera.pbj.integration.grpc.GrpcTestUtils;
 import com.hedera.pbj.integration.grpc.PortsAllocator;
 import com.hedera.pbj.runtime.Codec;
@@ -50,25 +51,39 @@
 @SuppressWarnings("unused")
 @State(Scope.Benchmark)
 @Fork(1)
-@Warmup(iterations = 2)
-@Measurement(iterations = 5)
+@Warmup(iterations = 1)
+@Measurement(iterations = 3)
 @OutputTimeUnit(TimeUnit.SECONDS)
 @BenchmarkMode(Mode.Throughput)
 public class PbjGrpcBench {
-    private static final int INVOCATIONS = 20_000;
+    private static final int INVOCATIONS = 2_000;
 
-    private record ServerHandle(WebServer server) implements AutoCloseable {
+    static {
+        new ZstdGrpcTransformer().register("zstd");
+
+        // 1Gbps network:
+        NetworkLatencySimulator.simulate(1_000, false);
+    }
+
+    public record ServerHandle(WebServer server) implements AutoCloseable {
         @Override
         public void close() {
             server.stop();
         }
 
-        static ServerHandle start(
+        public static ServerHandle start(
                 final int port, final ServiceInterface service, final PbjGrpcServiceConfig serviceConfig) {
+            final int maxPayloadSize = 20 * 1024 * 1024;
+            final PbjConfig pbjConfig = PbjConfig.builder()
+                    .name("pbj")
+                    .maxMessageSizeBytes(maxPayloadSize)
+                    .build();
+            ;
             return new ServerHandle(WebServer.builder()
                     .port(port)
+                    .addProtocol(pbjConfig)
                     .addRouting(PbjRouting.builder().service(service, serviceConfig))
-                    .maxPayloadSize(10000)
+                    .maxPayloadSize(maxPayloadSize)
                     .build()
                     .start());
         }
@@ -96,7 +111,7 @@ public static class UnaryState {
         @Param
         PayloadWeight weight;
 
-        @Param({"identity", "gzip"})
+        @Param({"identity", "gzip", "zstd"})
         String encodings;
 
         PortsAllocator.Port port;
@@ -143,10 +158,13 @@ public void tearDown() {
     // There's code duplicated from UnaryState above. It's because JMH is having troubles when states use inheritance.
     @State(Scope.Thread)
     public static class StreamingState {
-        @Param
+        // Skip the SUPER weight as it would be too slow.
+        @Param({"LIGHT", "NORMAL", "HEAVY"})
         PayloadWeight weight;
 
-        @Param({"identity", "gzip"})
+        // Streaming benchmarks are much slower than unary. Also, compression benchmarks are much slower than identity.
+        // So for streaming, only test the identity:
+        @Param({"identity"})
         String encodings;
 
         @Param({"1", "10"})