From e721d40846588da279cc319845f303af935298a8 Mon Sep 17 00:00:00 2001
From: William Burns <mudokonman@gmail.com>
Date: Tue, 9 May 2023 14:45:51 -0700
Subject: [PATCH] Add in new tests for ByteBuf based Encoder/Decoder instances

---
 protostream/pom.xml                           |   7 +-
 .../protostream/ByteBufDecoder.java           | 165 ++++++++++++++++++
 .../protostream/ByteBufEncoder.java           | 116 ++++++++++++
 .../protostream/ProtostreamBenchmark.java     |  94 ++++++----
 .../protostream/SerializationType.java        |   7 +
 .../protostream/state/AddressState.java       |   9 +
 .../protostream/state/MetadataState.java      |   9 +
 .../protostream/state/UserState.java          |   9 +
 8 files changed, 380 insertions(+), 36 deletions(-)
 create mode 100644 protostream/src/main/java/org/infinispan/protostream/ByteBufDecoder.java
 create mode 100644 protostream/src/main/java/org/infinispan/protostream/ByteBufEncoder.java
 create mode 100644 protostream/src/main/java/org/infinispan/protostream/SerializationType.java

diff --git a/protostream/pom.xml b/protostream/pom.xml
index e936f6b..ca1eed8 100644
--- a/protostream/pom.xml
+++ b/protostream/pom.xml
@@ -67,6 +67,11 @@ THE POSSIBILITY OF SUCH DAMAGE.
             <artifactId>protostream-processor</artifactId>
             <version>${protostream.version}</version>
         </dependency>
+        <dependency>
+            <groupId>io.netty</groupId>
+            <artifactId>netty-buffer</artifactId>
+            <version>4.1.92.Final</version>
+        </dependency>
 
     </dependencies>
 
@@ -81,7 +86,7 @@ THE POSSIBILITY OF SUCH DAMAGE.
         <!--
             Java source/target to use for compilation.
           -->
-        <javac.target>11</javac.target>
+        <javac.target>17</javac.target>
 
         <!--
             Name of the benchmark Uber-JAR to generate.
diff --git a/protostream/src/main/java/org/infinispan/protostream/ByteBufDecoder.java b/protostream/src/main/java/org/infinispan/protostream/ByteBufDecoder.java
new file mode 100644
index 0000000..be03a6f
--- /dev/null
+++ b/protostream/src/main/java/org/infinispan/protostream/ByteBufDecoder.java
@@ -0,0 +1,165 @@
+package org.infinispan.protostream;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+import org.infinispan.protostream.impl.TagReaderImpl.Decoder;
+
+import io.netty.buffer.ByteBuf;
+
+public class ByteBufDecoder extends Decoder {
+   private final ByteBuf buf;
+   private final int initialOffset;
+
+   public ByteBufDecoder(ByteBuf buf) {
+      this.buf = buf;
+      this.initialOffset = buf.readerIndex();
+   }
+
+   MalformedProtobufException truncated() {
+      // TODO: add this from log eventually
+      return new MalformedProtobufException("Input data ended unexpectedly in the middle of a field. The message is corrupt.");
+   }
+
+   MalformedProtobufException malformedVarint() {
+      // TODO: add this from log eventually
+      return new MalformedProtobufException("Encountered a malformed varint.");
+   }
+
+   @Override
+   protected int getEnd() {
+      return buf.writerIndex() - initialOffset;
+   }
+
+   @Override
+   protected int getPos() {
+      return buf.readableBytes() - initialOffset;
+   }
+
+   @Override
+   protected byte[] getBufferArray() {
+      byte[] bytes = new byte[buf.readableBytes()];
+      buf.readBytes(bytes);
+      return bytes;
+   }
+
+   @Override
+   protected boolean isAtEnd() {
+      return buf.readableBytes() == 0;
+   }
+
+   @Override
+   protected void skipVarint() throws IOException {
+      for (int i = 0; i < 5; ++i) {
+         if (readRawByte() >= 0) {
+            return;
+         }
+      }
+      throw malformedVarint();
+   }
+
+   @Override
+   protected void skipRawBytes(int length) throws IOException {
+      try {
+         buf.skipBytes(length);
+      } catch (IndexOutOfBoundsException e) {
+         throw truncated();
+      }
+   }
+
+   @Override
+   protected String readString() throws IOException {
+      int length = readVarint32();
+      if (length > buf.readableBytes()) {
+         throw truncated();
+      }
+      String string = buf.toString(buf.readerIndex(), length, StandardCharsets.UTF_8);
+      buf.skipBytes(length);
+      return string;
+   }
+
+   @Override
+   protected byte readRawByte() throws IOException {
+      try {
+         return buf.readByte();
+      } catch (IndexOutOfBoundsException e) {
+         throw truncated();
+      }
+   }
+
+   @Override
+   protected byte[] readRawByteArray(int length) throws IOException {
+      if (length > buf.readableBytes()) {
+         throw truncated();
+      }
+      byte[] bytes = new byte[length];
+      buf.readBytes(bytes);
+      return bytes;
+   }
+
+   @Override
+   protected ByteBuffer readRawByteBuffer(int length) throws IOException {
+      if (length > buf.readableBytes()) {
+         throw truncated();
+      }
+      return buf.nioBuffer(buf.readerIndex(), length);
+   }
+
+   @Override
+   protected long readVarint64() throws IOException {
+      long value = 0;
+      for (int i = 0; i < 64; i += 7) {
+         byte b = readRawByte();
+         value |= (long) (b & 0x7F) << i;
+         if (b >= 0) {
+            return value;
+         }
+      }
+      throw malformedVarint();
+   }
+
+   @Override
+   protected int readFixed32() throws IOException {
+      try {
+         return buf.readInt();
+      } catch (IndexOutOfBoundsException e) {
+         throw truncated();
+      }
+   }
+
+   @Override
+   protected long readFixed64() throws IOException {
+      try {
+         return buf.readLong();
+      } catch (IndexOutOfBoundsException e) {
+         throw truncated();
+      }
+   }
+
+   @Override
+   protected int pushLimit(int newLimit) throws IOException {
+      int prev = buf.writerIndex() - initialOffset;
+      if (prev < newLimit) {
+         throw truncated();
+      }
+      buf.writerIndex(buf.readerIndex() + newLimit);
+      return prev;
+   }
+
+   @Override
+   protected void popLimit(int oldLimit) {
+      buf.writerIndex(oldLimit + initialOffset);
+   }
+
+   @Override
+   protected Decoder decoderFromLength(int length) {
+      ByteBuf slicedBuf = buf.readSlice(length);
+      return new ByteBufDecoder(slicedBuf);
+   }
+
+   @Override
+   protected int setGlobalLimit(int globalLimit) {
+      return Integer.MAX_VALUE;
+   }
+}
diff --git a/protostream/src/main/java/org/infinispan/protostream/ByteBufEncoder.java b/protostream/src/main/java/org/infinispan/protostream/ByteBufEncoder.java
new file mode 100644
index 0000000..93cf507
--- /dev/null
+++ b/protostream/src/main/java/org/infinispan/protostream/ByteBufEncoder.java
@@ -0,0 +1,116 @@
+package org.infinispan.protostream;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import org.infinispan.protostream.impl.TagWriterImpl;
+
+import io.netty.buffer.ByteBuf;
+
+public class ByteBufEncoder extends TagWriterImpl.FixedVarintEncoder {
+   private final ByteBuf buf;
+
+   public ByteBufEncoder(ByteBuf buf) {
+      this.buf = buf;
+   }
+
+   private IOException outOfWriteBufferSpace(Throwable t) {
+      return new IOException("Ran out of buffer space", t);
+   }
+
+   @Override
+   protected void writeVarint32(int value) throws IOException {
+      try {
+         while (true) {
+            if ((value & 0xFFFFFF80) == 0) {
+               buf.writeByte((byte) value);
+               break;
+            } else {
+               buf.writeByte((byte) (value & 0x7F | 0x80));
+               value >>>= 7;
+            }
+         }
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected void writeVarint64(long value) throws IOException {
+      try {
+         while (true) {
+            if ((value & 0xFFFFFFFFFFFFFF80L) == 0) {
+               buf.writeByte((byte) value);
+               break;
+            } else {
+               buf.writeByte((byte) ((int) value & 0x7F | 0x80));
+               value >>>= 7;
+            }
+         }
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected void writeFixed32(int value) throws IOException {
+      try {
+         buf.writeInt(value);
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected void writeFixed64(long value) throws IOException {
+      try {
+         buf.writeLong(value);
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected void writeByte(byte value) throws IOException {
+      try {
+         buf.writeByte(value);
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected void writeBytes(byte[] value, int offset, int length) throws IOException {
+      try {
+         buf.writeBytes(value, offset, length);
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected void writeBytes(ByteBuffer value) throws IOException {
+      try {
+         buf.writeBytes(value);
+      } catch (IndexOutOfBoundsException e) {
+         throw outOfWriteBufferSpace(e);
+      }
+   }
+
+   @Override
+   protected int skipFixedVarint() {
+      int prev = buf.writerIndex();
+      buf.skipBytes(5);
+      return prev;
+   }
+
+   @Override
+   protected void writePositiveFixedVarint(int pos) {
+      int length = buf.writerIndex() - pos - 5;
+      buf.setByte(pos++, (byte) (length & 0x7F | 0x80));
+      buf.setByte(pos++, (byte) ((length >>> 7) & 0x7F | 0x80));
+      buf.setByte(pos++, (byte) ((length >>> 14) & 0x7F | 0x80));
+      buf.setByte(pos++, (byte) ((length >>> 21) & 0x7F | 0x80));
+      buf.setByte(pos, (byte) ((length >>> 28) & 0x7F));
+   }
+}
diff --git a/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java b/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java
index f7fbff7..6d257a2 100644
--- a/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java
+++ b/protostream/src/main/java/org/infinispan/protostream/ProtostreamBenchmark.java
@@ -23,6 +23,9 @@
 import org.openjdk.jmh.annotations.State;
 import org.openjdk.jmh.annotations.Warmup;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.NANOSECONDS)
 @Warmup(iterations = 3, time = 10)
@@ -31,75 +34,96 @@
 @State(Scope.Benchmark)
 public class ProtostreamBenchmark {
 
-   @Param({"true", "false"})
-   boolean byteArrayOrStream;
+   @Param
+   SerializationType serializationType;
 
    @Benchmark
    public Object testMarshallAddress(ContextState contextState, AddressState addressState) throws IOException {
       SerializationContext ctx = contextState.getCtx();
       Address address = addressState.getAddress();
-      if (byteArrayOrStream)
-         return ProtobufUtil.toWrappedByteArray(ctx, address);
-      else {
-         ByteArrayOutputStreamEx os = new ByteArrayOutputStreamEx(addressState.getAddressBytes().length);
-         ProtobufUtil.toWrappedStream(ctx, os, address);
-         return os.getByteBuffer();
-      }
+      return switch (serializationType) {
+         case ARRAY -> ProtobufUtil.toWrappedByteArray(ctx, address);
+         case INPUT_STREAM -> {
+            ByteArrayOutputStreamEx os = new ByteArrayOutputStreamEx(addressState.getAddressBytes().length);
+            ProtobufUtil.toWrappedStream(ctx, os, address);
+            yield os.getByteBuffer();
+         }
+         case BYTE_BUF -> {
+            ByteBuf buf = ByteBufAllocator.DEFAULT.buffer(addressState.getAddressBytes().length);
+            ProtobufUtil.toWrappedEncoder(ctx, new ByteBufEncoder(buf), address);
+            yield buf.release();
+         }
+      };
    }
 
    @Benchmark
    public Object testMarshallUser(ContextState contextState, UserState userState) throws IOException {
       SerializationContext ctx = contextState.getCtx();
       User user = userState.getUser();
-      if (byteArrayOrStream)
-         return ProtobufUtil.toWrappedByteArray(ctx, user);
-      else {
-         ByteArrayOutputStreamEx os = new ByteArrayOutputStreamEx(userState.getUserBytes().length);
-         ProtobufUtil.toWrappedStream(ctx, os, user);
-         return os.getByteBuffer();
-      }
+      return switch (serializationType) {
+         case ARRAY -> ProtobufUtil.toWrappedByteArray(ctx, user);
+         case INPUT_STREAM -> {
+            ByteArrayOutputStreamEx os = new ByteArrayOutputStreamEx(userState.getUserBytes().length);
+            ProtobufUtil.toWrappedStream(ctx, os, user);
+            yield os.getByteBuffer();
+         }
+         case BYTE_BUF -> {
+            ByteBuf buf = ByteBufAllocator.DEFAULT.buffer(userState.getUserBytes().length);
+            ProtobufUtil.toWrappedEncoder(ctx, new ByteBufEncoder(buf), user);
+            yield buf.release();
+         }
+      };
    }
 
    @Benchmark
    public Object testMarshallIracMetadata(ContextState contextState, MetadataState metadataState) throws IOException {
       SerializationContext ctx = contextState.getCtx();
       IracMetadata metadata = metadataState.getMetadata();
-      if (byteArrayOrStream)
-         return ProtobufUtil.toWrappedByteArray(ctx, metadata);
-      else {
-         ByteArrayOutputStreamEx os = new ByteArrayOutputStreamEx(metadataState.getMetadataBytes().length);
-         ProtobufUtil.toWrappedStream(ctx, os, metadata);
-         return os.getByteBuffer();
-      }
+      return switch (serializationType) {
+         case ARRAY -> ProtobufUtil.toWrappedByteArray(ctx, metadata);
+         case INPUT_STREAM -> {
+            ByteArrayOutputStreamEx os = new ByteArrayOutputStreamEx(metadataState.getMetadataBytes().length);
+            ProtobufUtil.toWrappedStream(ctx, os, metadata);
+            yield os.getByteBuffer();
+         }
+         case BYTE_BUF -> {
+            ByteBuf buf = ByteBufAllocator.DEFAULT.buffer(metadataState.getMetadataBytes().length);
+            ProtobufUtil.toWrappedEncoder(ctx, new ByteBufEncoder(buf), metadata);
+            yield buf.release();
+         }
+      };
    }
 
    @Benchmark
    public Address testUnmarshallAddress(ContextState contextState, AddressState addressState) throws IOException {
       SerializationContext ctx = contextState.getCtx();
       byte[] addressBytes = addressState.getAddressBytes();
-      if (byteArrayOrStream)
-         return ProtobufUtil.fromWrappedByteArray(ctx, addressBytes);
-      else
-         return ProtobufUtil.fromWrappedStream(ctx, new ByteArrayInputStream(addressBytes));
+      return switch (serializationType) {
+         case ARRAY -> ProtobufUtil.fromWrappedByteArray(ctx, addressBytes);
+         case INPUT_STREAM -> ProtobufUtil.fromWrappedStream(ctx, new ByteArrayInputStream(addressBytes));
+         case BYTE_BUF -> ProtobufUtil.fromWrappedDecoder(ctx, new ByteBufDecoder(addressState.getAddressByteBuf()));
+      };
    }
 
    @Benchmark
    public User testUnmarshallUser(ContextState contextState, UserState userState) throws IOException {
       SerializationContext ctx = contextState.getCtx();
       byte[] userBytes = userState.getUserBytes();
-      if (byteArrayOrStream)
-         return ProtobufUtil.fromWrappedByteArray(ctx, userBytes);
-      else
-         return ProtobufUtil.fromWrappedStream(ctx, new ByteArrayInputStream(userBytes));
+      return switch (serializationType) {
+         case ARRAY -> ProtobufUtil.fromWrappedByteArray(ctx, userBytes);
+         case INPUT_STREAM -> ProtobufUtil.fromWrappedStream(ctx, new ByteArrayInputStream(userBytes));
+         case BYTE_BUF -> ProtobufUtil.fromWrappedDecoder(ctx, new ByteBufDecoder(userState.getUserByteBuf()));
+      };
    }
 
    @Benchmark
    public IracMetadata testUnmarshallMetadata(ContextState contextState, MetadataState metadataState) throws IOException {
       SerializationContext ctx = contextState.getCtx();
       byte[] metadataBytes = metadataState.getMetadataBytes();
-      if (byteArrayOrStream)
-         return ProtobufUtil.fromWrappedByteArray(ctx, metadataBytes);
-      else
-         return ProtobufUtil.fromWrappedStream(ctx, new ByteArrayInputStream(metadataBytes));
+      return switch (serializationType) {
+         case ARRAY -> ProtobufUtil.fromWrappedByteArray(ctx, metadataBytes);
+         case INPUT_STREAM -> ProtobufUtil.fromWrappedStream(ctx, new ByteArrayInputStream(metadataBytes));
+         case BYTE_BUF -> ProtobufUtil.fromWrappedDecoder(ctx, new ByteBufDecoder(metadataState.getMetadataByteBuf()));
+      };
    }
 }
diff --git a/protostream/src/main/java/org/infinispan/protostream/SerializationType.java b/protostream/src/main/java/org/infinispan/protostream/SerializationType.java
new file mode 100644
index 0000000..640a655
--- /dev/null
+++ b/protostream/src/main/java/org/infinispan/protostream/SerializationType.java
@@ -0,0 +1,7 @@
+package org.infinispan.protostream;
+
+public enum SerializationType {
+   ARRAY,
+   INPUT_STREAM,
+   BYTE_BUF
+}
diff --git a/protostream/src/main/java/org/infinispan/protostream/state/AddressState.java b/protostream/src/main/java/org/infinispan/protostream/state/AddressState.java
index 78f7b36..08e4a43 100644
--- a/protostream/src/main/java/org/infinispan/protostream/state/AddressState.java
+++ b/protostream/src/main/java/org/infinispan/protostream/state/AddressState.java
@@ -10,12 +10,16 @@
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
 @State(Scope.Benchmark)
 public class AddressState {
    private Address address;
    private byte[] addressBytes;
 
    private List<Address> addressesForUser;
+   private ByteBuf addressByteBuf;
 
    @Setup
    public void setup(ContextState contextState) throws IOException {
@@ -28,6 +32,7 @@ public void setup(ContextState contextState) throws IOException {
       addressesForUser.add(new Address("Out of street", "3333-111", 1, true));
 
       addressBytes = ProtobufUtil.toWrappedByteArray(contextState.getCtx(), address);
+      addressByteBuf = Unpooled.wrappedBuffer(addressBytes);
    }
 
    public Address getAddress() {
@@ -41,4 +46,8 @@ public byte[] getAddressBytes() {
    public List<Address> getAddressesForUser() {
       return addressesForUser;
    }
+
+   public ByteBuf getAddressByteBuf() {
+      return addressByteBuf.slice();
+   }
 }
diff --git a/protostream/src/main/java/org/infinispan/protostream/state/MetadataState.java b/protostream/src/main/java/org/infinispan/protostream/state/MetadataState.java
index 27612f4..dca0069 100644
--- a/protostream/src/main/java/org/infinispan/protostream/state/MetadataState.java
+++ b/protostream/src/main/java/org/infinispan/protostream/state/MetadataState.java
@@ -12,11 +12,15 @@
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
 @State(Scope.Benchmark)
 public class MetadataState {
    private IracMetadata metadata;
 
    byte[] metadataBytes;
+   private ByteBuf metadataByteBuf;
    @Setup
    public void setup(ContextState contextState) throws IOException {
       Map<String, TopologyIracVersion> versions = new HashMap<>();
@@ -26,6 +30,7 @@ public void setup(ContextState contextState) throws IOException {
       metadata = new IracMetadata("site_1", version);
 
       metadataBytes = ProtobufUtil.toWrappedByteArray(contextState.getCtx(), metadata);
+      metadataByteBuf = Unpooled.wrappedBuffer(metadataBytes);
    }
 
    public IracMetadata getMetadata() {
@@ -35,4 +40,8 @@ public IracMetadata getMetadata() {
    public byte[] getMetadataBytes() {
       return metadataBytes;
    }
+
+   public ByteBuf getMetadataByteBuf() {
+      return metadataByteBuf.slice();
+   }
 }
diff --git a/protostream/src/main/java/org/infinispan/protostream/state/UserState.java b/protostream/src/main/java/org/infinispan/protostream/state/UserState.java
index aed63ee..247208e 100644
--- a/protostream/src/main/java/org/infinispan/protostream/state/UserState.java
+++ b/protostream/src/main/java/org/infinispan/protostream/state/UserState.java
@@ -15,12 +15,16 @@
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
 @State(Scope.Benchmark)
 public class UserState {
    @Param({"10", "8096"})
    int userByteArraySize;
    private User user;
    private byte[] userBytes;
+   private ByteBuf userByteBuf;
    @Setup
    public void setup(ContextState contextState, AddressState addressState) throws IOException {
       byte[] data = new byte[userByteArraySize];
@@ -32,6 +36,7 @@ public void setup(ContextState contextState, AddressState addressState) throws I
             User.Gender.MALE, null, Instant.now(), Instant.now().plusMillis(TimeUnit.DAYS.toMillis(30)), data);
 
       userBytes = ProtobufUtil.toWrappedByteArray(contextState.getCtx(), user);
+      userByteBuf = Unpooled.wrappedBuffer(userBytes);
    }
 
    public User getUser() {
@@ -41,4 +46,8 @@ public User getUser() {
    public byte[] getUserBytes() {
       return userBytes;
    }
+
+   public ByteBuf getUserByteBuf() {
+      return userByteBuf.slice();
+   }
 }