Pruivo feedback

core/src/main/java/org/infinispan/protostream/ProtobufUtil.java

@@ -74,7 +74,8 @@ public static void writeTo(ImmutableSerializationContext ctx, RandomAccessOutput
    }
 
    public static void writeTo(ImmutableSerializationContext ctx, OutputStream out, Object t) throws IOException {
-      write(ctx, TagWriterImpl.newInstance(ctx, out), t);
+      TagWriterImpl writer = out instanceof RandomAccessOutputStream raos ? TagWriterImpl.newInstance(ctx, raos) : TagWriterImpl.newInstance(ctx, out);
+      write(ctx, writer, t);
    }
 
    public static byte[] toByteArray(ImmutableSerializationContext ctx, Object t) throws IOException {
@@ -161,7 +162,7 @@ public static void toWrappedStream(ImmutableSerializationContext ctx, RandomAcce
    }
 
    public static void toWrappedStream(ImmutableSerializationContext ctx, OutputStream out, Object t) throws IOException {
-      TagWriter writer = out instanceof RandomAccessOutputStream raw ? TagWriterImpl.newInstance(ctx, raw) : TagWriterImpl.newInstance(ctx, out);
+      TagWriter writer = out instanceof RandomAccessOutputStream raos ? TagWriterImpl.newInstance(ctx, raos) : TagWriterImpl.newInstance(ctx, out);
       WrappedMessage.write(ctx, writer, t);
    }
 

core/src/main/java/org/infinispan/protostream/RandomAccessOutputStream.java

@@ -1,6 +1,7 @@
 package org.infinispan.protostream;
 
 import java.io.Closeable;
+import java.io.DataOutput;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 
@@ -18,14 +19,14 @@ public interface RandomAccessOutputStream extends Closeable {
    int getPosition();
 
    /**
-    * Set the position in the stream where the next {@link #write(int)} will be written;
+    * Set the position in the stream where the next {@link #write(int)} will be written.
     */
    void setPosition(int position);
 
    /**
     * Ensure that the stream has sufficient capacity to hold at least `capacity` bytes.
     *
-    * @param capacity the minimum number of bytes that should be writeable.
+    * @param capacity the minimum number of bytes that should be writeable
     */
    void ensureCapacity(int capacity);
 
@@ -38,17 +39,17 @@ default void reset() {
    }
 
    /**
-    * @return a {@link ByteBuffer} representation of the stream.
+    * @return a {@link ByteBuffer} representation of the stream
     */
    ByteBuffer getByteBuffer();
 
    /**
-    * @return a trimmed {@link byte[]} instance based upon the current {@link #getPosition()} of the array.
+    * @return a trimmed {@link byte[]} instance based upon the current {@link #getPosition()} of the array
     */
    byte[] toByteArray();
 
    /**
-    * @param position the position in the stream to read the byte from.
+    * @param position the position in the stream to read the byte from
     * @return the byte associated with the specified position
     */
    byte get(int position);
@@ -85,15 +86,15 @@ default void write(byte[] b, int off, int len) throws IOException {
    /**
     * Write a single byte to the specified position in the stream.
     *
-    * @param position the position in the stream to write the byte to.
+    * @param position the position in the stream to write the byte to
     * @param b the byte to be written
     */
    void write(int position, int b) throws IOException;
 
    /**
     * Write all bytes to the specified position in the stream.
     *
-    * @param position the position in the stream to write the bytes to.
+    * @param position the position in the stream to write the bytes to
     * @param b the array of bytes to be written
     */
    default void write(int position, byte[] b) throws IOException {
@@ -103,10 +104,16 @@ default void write(int position, byte[] b) throws IOException {
    /**
     * Write all bytes to the specified position in the stream.
     *
-    * @param position the position in the stream to write the bytes to.
+    * @param position the position in the stream to write the bytes to
     * @param b the array of bytes to be written
     * @param off the offset within the array to be written
     * @param len the number of bytes from the array to be written
     */
    void write(int position, byte[] b, int off, int len) throws IOException;
+
+   /**
+    * Write all bytes in this stream to the provided {@link DataOutput} stream
+    * @param output the stream to write this stream's bytes to
+    */
+   void copyTo(DataOutput output) throws IOException;
 }

core/src/main/java/org/infinispan/protostream/impl/RandomAccessOutputStreamImpl.java

@@ -1,5 +1,7 @@
 package org.infinispan.protostream.impl;
 
+import java.io.DataOutput;
+import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
@@ -72,7 +74,7 @@ public void ensureCapacity(int capacity) {
       }
    }
 
-   private int getNewBufferSize(int curSize, int minNewSize) {
+   private static int getNewBufferSize(int curSize, int minNewSize) {
       if (curSize <= DEFAULT_DOUBLING_SIZE)
          return Math.max(curSize << 1, minNewSize);
       else
@@ -100,4 +102,9 @@ public ByteBuffer getByteBuffer() {
             ByteBuffer.wrap(new byte[0]) :
             ByteBuffer.wrap(buf, 0, pos);
    }
+
+   @Override
+   public void copyTo(DataOutput output) throws IOException {
+      output.write(buf, 0, pos);
+   }
 }

core/src/main/java/org/infinispan/protostream/impl/TagWriterImpl.java

@@ -957,7 +957,7 @@ void writeUInt32Field(int fieldNumber, int value) throws IOException {
       }
 
       @Override
-      void writeUInt64Field(int fieldNumber, long value) throws IOException{
+      void writeUInt64Field(int fieldNumber, long value) throws IOException {
          int pos = out.getPosition();
          int tag = WireType.makeTag(fieldNumber, WireType.WIRETYPE_VARINT);
          int tagLen = varIntBytes(tag);
@@ -1199,4 +1199,5 @@ private static int varIntBytes(long value) {
          }
          return i;
       }
+   }
 }

core/src/main/java/org/infinispan/protostream/impl/UnknownFieldSetImpl.java

@@ -201,12 +201,8 @@ public void writeExternal(ObjectOutput out) throws IOException {
       RandomAccessOutputStream baos = new RandomAccessOutputStreamImpl();
       TagWriter output = TagWriterImpl.newInstance(null, baos);
       writeTo(output);
-      output.flush();
-
-      int pos = baos.getPosition();
-      out.writeInt(pos);
-      for (int i = 0; i < pos; i++)
-         out.write(baos.get(i));
+      out.writeInt(baos.getPosition());
+      baos.copyTo(out);
    }
 
    @Override