From a33740da23acf464f2f69139145c34b4302c6694 Mon Sep 17 00:00:00 2001
From: William Burns <mudokonman@gmail.com>
Date: Tue, 9 Sep 2025 12:46:54 -0700
Subject: [PATCH] Add size estimate benchmarks

---
 protostream/pom.xml                           |  2 +-
 .../protostream/ProtostreamBenchmark.java     | 46 +++++++++++++++++--
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/protostream/pom.xml b/protostream/pom.xml
index 919b0b6..bed61f2 100644
--- a/protostream/pom.xml
+++ b/protostream/pom.xml
@@ -93,7 +93,7 @@ THE POSSIBILITY OF SUCH DAMAGE.
             Name of the benchmark Uber-JAR to generate.
           -->
         <uberjar.name>protostream-benchmark</uberjar.name>
-        <protostream.version>6.0.0.Dev06</protostream.version>
+        <protostream.version>6.0.0-SNAPSHOT</protostream.version>
     </properties>
 
     <build>
diff --git a/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java b/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java
index d541770..7be35f3 100644
--- a/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java
+++ b/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java
@@ -1,5 +1,10 @@
 package org.infinispan.protostream;
 
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+
 import org.infinispan.protostream.impl.RandomAccessOutputStreamImpl;
 import org.infinispan.protostream.impl.TagWriterImpl;
 import org.infinispan.protostream.state.AddressState;
@@ -21,11 +26,6 @@
 import org.openjdk.jmh.annotations.State;
 import org.openjdk.jmh.annotations.Warmup;
 
-import java.io.ByteArrayInputStream;
-import java.io.IOException;
-import java.util.UUID;
-import java.util.concurrent.TimeUnit;
-
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.NANOSECONDS)
 @Warmup(iterations = 3, time = 10)
@@ -46,6 +46,15 @@ public Object testSizedAddress(ContextState contextState, AddressState addressSt
       return sizedWriter;
    }
 
+   @Benchmark
+   public Object testSizeEstimateAddress(ContextState contextState, AddressState addressState) throws IOException {
+      SerializationContext ctx = contextState.getCtx();
+      Address address = addressState.getAddress();
+      TagWriterImpl sizedWriter = TagWriterImpl.newInstanceSizeEstimate(ctx);
+      WrappedMessage.write(ctx, sizedWriter, address);
+      return sizedWriter;
+   }
+
    @Benchmark
    public Object testMarshallAddress(ContextState contextState, AddressState addressState) throws IOException {
       RandomAccessOutputStreamImpl raos = contextState.getOutputStream();
@@ -68,6 +77,15 @@ public Object testSizedUser(ContextState contextState, UserState userState) thro
       return sizedWriter;
    }
 
+   @Benchmark
+   public Object testSizeEstimateUser(ContextState contextState, UserState userState) throws IOException {
+      SerializationContext ctx = contextState.getCtx();
+      User user = userState.getUser();
+      TagWriterImpl sizedWriter = TagWriterImpl.newInstanceSizeEstimate(ctx);
+      WrappedMessage.write(ctx, sizedWriter, user);
+      return sizedWriter;
+   }
+
    @Benchmark
    public Object testMarshallUser(ContextState contextState, UserState userState) throws IOException {
       RandomAccessOutputStreamImpl raos = contextState.getOutputStream();
@@ -90,6 +108,15 @@ public Object testSizedIracMetadata(ContextState contextState, MetadataState met
       return sizedWriter;
    }
 
+   @Benchmark
+   public Object testSizeEstimateIracMetadata(ContextState contextState, MetadataState metadataState) throws IOException {
+      SerializationContext ctx = contextState.getCtx();
+      IracMetadata metadata = metadataState.getMetadata();
+      TagWriterImpl sizedWriter = TagWriterImpl.newInstanceSizeEstimate(ctx);
+      WrappedMessage.write(ctx, sizedWriter, metadata);
+      return sizedWriter;
+   }
+
    @Benchmark
    public Object testMarshallIracMetadata(ContextState contextState, MetadataState metadataState) throws IOException {
       RandomAccessOutputStreamImpl raos = contextState.getOutputStream();
@@ -112,6 +139,15 @@ public Object testSizedUUID(ContextState contextState, UUIDState uuidState) thro
       return sizedWriter;
    }
 
+   @Benchmark
+   public Object testSizeEstimateUUID(ContextState contextState, UUIDState uuidState) throws IOException {
+      var ctx = contextState.getCtx();
+      var uuid = uuidState.getUuid();
+      var sizedWriter = TagWriterImpl.newInstanceSizeEstimate(ctx);
+      WrappedMessage.write(ctx, sizedWriter, uuid);
+      return sizedWriter;
+   }
+
    @Benchmark
    public Object testMarshallUUID(ContextState contextState, UUIDState uuidState) throws IOException {
       var raos = contextState.getOutputStream();