ISPN-6056 Improve ReplicableCommand marshalling

commons/src/main/java/org/infinispan/commons/marshall/AbstractDelegatingMarshaller.java

@@ -30,6 +30,11 @@ public ObjectOutput startObjectOutput(OutputStream os, boolean isReentrant, fina
       return marshaller.startObjectOutput(os, isReentrant, estimatedSize);
    }
 
+   @Override
+   public ObjectOutput startObjectOutput(ObjectOutput oo, boolean isReentrant, int estimatedSize) throws IOException {
+      return marshaller.startObjectOutput(oo, isReentrant, estimatedSize);
+   }
+
    @Override
    public void finishObjectOutput(ObjectOutput oo) {
       marshaller.finishObjectOutput(oo);
@@ -45,6 +50,11 @@ public ObjectInput startObjectInput(InputStream is, boolean isReentrant) throws
       return marshaller.startObjectInput(is, isReentrant);
    }
 
+   @Override
+   public ObjectInput startObjectInput(ObjectInput oi, boolean isReentrant) throws IOException {
+      return marshaller.startObjectInput(oi, isReentrant);
+   }
+
    @Override
    public void finishObjectInput(ObjectInput oi) {
       marshaller.finishObjectInput(oi);

commons/src/main/java/org/infinispan/commons/marshall/AbstractDelegatingObjectOutput.java

@@ -0,0 +1,104 @@
+package org.infinispan.commons.marshall;
+
+import java.io.IOException;
+import java.io.ObjectOutput;
+
+/**
+ * // TODO: Document this
+ *
+ * @author Pedro Ruivo
+ * @since 8.0
+ */
+public abstract class AbstractDelegatingObjectOutput implements ObjectOutput {
+
+   protected final ObjectOutput objectOutput;
+
+   protected AbstractDelegatingObjectOutput(ObjectOutput objectOutput) {
+      this.objectOutput = objectOutput;
+   }
+
+   @Override
+   public void writeObject(Object obj) throws IOException {
+      objectOutput.writeObject(obj);
+   }
+
+   @Override
+   public void write(int b) throws IOException {
+      objectOutput.write(b);
+   }
+
+   @Override
+   public void write(byte[] b) throws IOException {
+      objectOutput.write(b);
+   }
+
+   @Override
+   public void write(byte[] b, int off, int len) throws IOException {
+      objectOutput.write(b, off, len);
+   }
+
+   @Override
+   public void flush() throws IOException {
+      objectOutput.flush();
+   }
+
+   @Override
+   public void close() throws IOException {
+      objectOutput.close();
+   }
+
+   @Override
+   public void writeBoolean(boolean v) throws IOException {
+      objectOutput.writeBoolean(v);
+   }
+
+   @Override
+   public void writeByte(int v) throws IOException {
+      objectOutput.writeByte(v);
+   }
+
+   @Override
+   public void writeShort(int v) throws IOException {
+      objectOutput.writeShort(v);
+   }
+
+   @Override
+   public void writeChar(int v) throws IOException {
+      objectOutput.writeChar(v);
+   }
+
+   @Override
+   public void writeInt(int v) throws IOException {
+      objectOutput.writeInt(v);
+   }
+
+   @Override
+   public void writeLong(long v) throws IOException {
+      objectOutput.writeLong(v);
+   }
+
+   @Override
+   public void writeFloat(float v) throws IOException {
+      objectOutput.writeFloat(v);
+   }
+
+   @Override
+   public void writeDouble(double v) throws IOException {
+      objectOutput.writeDouble(v);
+   }
+
+   @Override
+   public void writeBytes(String s) throws IOException {
+      objectOutput.writeBytes(s);
+   }
+
+   @Override
+   public void writeChars(String s) throws IOException {
+      objectOutput.writeChars(s);
+   }
+
+   @Override
+   public void writeUTF(String s) throws IOException {
+      objectOutput.writeUTF(s);
+   }
+}

commons/src/main/java/org/infinispan/commons/marshall/MarshallUtil.java

@@ -1,14 +1,15 @@
 package org.infinispan.commons.marshall;
 
 import net.jcip.annotations.Immutable;
+import org.infinispan.commons.io.UnsignedNumeric;
 
 import java.io.IOException;
 import java.io.ObjectInput;
 import java.io.ObjectOutput;
 import java.util.Collection;
 import java.util.Map;
-
-import org.infinispan.commons.io.UnsignedNumeric;
+import java.util.Objects;
+import java.util.UUID;
 
 /**
  * MarshallUtil.
@@ -19,26 +20,174 @@
 @Immutable
 public class MarshallUtil {
 
-   public static void marshallCollection(Collection<?> c, ObjectOutput out) throws IOException {
-      UnsignedNumeric.writeUnsignedInt(out, c.size());
-      for (Object o : c) {
-         out.writeObject(o);
+   public static <K, V, T extends Map<K, V>> void marshallMap(T map, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(map, out)) {
+         return;
       }
-   }
-
-   public static void marshallMap(Map<?, ?> map, ObjectOutput out) throws IOException {
       int mapSize = map.size();
       UnsignedNumeric.writeUnsignedInt(out, mapSize);
       if (mapSize == 0) return;
 
-      for (Map.Entry<?, ?> me : map.entrySet()) {
+      for (Map.Entry<K, V> me : map.entrySet()) {
          out.writeObject(me.getKey());
          out.writeObject(me.getValue());
       }
    }
 
-   public static void unmarshallMap(Map<Object, Object> map, ObjectInput in) throws IOException, ClassNotFoundException {
+   public static <K, V, T extends Map<K, V>> T unmarshallMap(ObjectInput in, MapBuilder<K, V, T> builder, boolean checkNull) throws IOException, ClassNotFoundException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      T map = Objects.requireNonNull(builder, "MapBuilder must be non-null").build();
+      int size = UnsignedNumeric.readUnsignedInt(in);
+      for (int i = 0; i < size; i++) //noinspection unchecked
+         map.put((K) in.readObject(), (V) in.readObject());
+      return map;
+   }
+
+   public static void marshallUUID(UUID uuid, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(uuid, out)) {
+         return;
+      }
+      out.writeLong(uuid.getMostSignificantBits());
+      out.writeLong(uuid.getLeastSignificantBits());
+   }
+
+   public static UUID unmarshallUUID(ObjectInput in, boolean checkNull) throws IOException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      return new UUID(in.readLong(), in.readLong());
+   }
+
+   public static <E> void marshallArray(E[] array, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(array, out)) {
+         return;
+      }
+      final int size = array.length;
+      UnsignedNumeric.writeUnsignedInt(out, size);
+      for (int i = 0; i < size; ++i) {
+         out.writeObject(array[i]);
+      }
+   }
+
+   public static <E> E[] unmarshallArray(ObjectInput in, ArrayBuilder<E> builder, boolean checkNull) throws IOException, ClassNotFoundException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      final int size = UnsignedNumeric.readUnsignedInt(in);
+      final E[] array = Objects.requireNonNull(builder, "ArrayBuilder must be non-null").build(size);
+      for (int i = 0; i < size; ++i) {
+         //noinspection unchecked
+         array[i] = (E) in.readObject();
+      }
+      return array;
+   }
+
+   public static <E> void marshallCollection(Collection<E> collection, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(collection, out)) {
+         return;
+      }
+      int size = collection.size();
+      UnsignedNumeric.writeUnsignedInt(out, size);
+      if (size == 0) {
+         return;
+      }
+      for (E e : collection) {
+         out.writeObject(e);
+      }
+   }
+
+   public static <E, T extends Collection<E>> T unmarshallCollection(ObjectInput in, CollectionBuilder<E, T> builder, boolean checkNull) throws IOException, ClassNotFoundException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
       int size = UnsignedNumeric.readUnsignedInt(in);
-      for (int i = 0; i < size; i++) map.put(in.readObject(), in.readObject());
+      T collection = Objects.requireNonNull(builder, "CollectionBuilder must be non-null").build(size);
+      for (int i = 0; i < size; ++i) {
+         //noinspection unchecked
+         collection.add((E) in.readObject());
+      }
+      return collection;
+   }
+
+   public static <E, T extends Collection<E>> T unmarshallCollectionUnbounded(ObjectInput in, UnboundedCollectionBuilder<E, T> builder, boolean checkNull) throws IOException, ClassNotFoundException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      T collection = Objects.requireNonNull(builder, "UnboundedCollectionBuilder must be non-null").build();
+      int size = UnsignedNumeric.readUnsignedInt(in);
+      for (int i = 0; i < size; ++i) {
+         //noinspection unchecked
+         collection.add((E) in.readObject());
+      }
+      return collection;
+   }
+
+   public static void marshallString(String string, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(string, out)) {
+         return;
+      }
+      out.writeUTF(string);
+   }
+
+   public static String unmarshallString(ObjectInput in, boolean checkNull) throws IOException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      return in.readUTF();
+   }
+
+   public static <E extends Enum<E>> void marshallEnum(E e, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(e, out)) {
+         return;
+      }
+      out.writeByte(e.ordinal());
+   }
+
+   public static <E extends Enum<E>> E unmarshallEnum(E[] values, ObjectInput in, boolean checkNull) throws IOException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      return Objects.requireNonNull(values, "Enum values must be non-null.")[in.readByte()];
+   }
+
+   public static void marshallByteArray(byte[] array, ObjectOutput out, boolean checkNull) throws IOException {
+      if (checkNull && checkIfNull(array, out)) {
+         return;
+      }
+      UnsignedNumeric.writeUnsignedInt(out, array.length);
+      out.write(array);
+   }
+
+   public static byte[] unmarshallByteArray(ObjectInput in, boolean checkNull) throws IOException {
+      if (checkNull && in.readBoolean()) {
+         return null;
+      }
+      byte[] array = new byte[UnsignedNumeric.readUnsignedInt(in)];
+      in.readFully(array);
+      return array;
+   }
+
+   private static boolean checkIfNull(Object obj, ObjectOutput output) throws IOException {
+      boolean isNull = obj == null;
+      output.writeBoolean(isNull);
+      return isNull;
+   }
+
+   public interface ArrayBuilder<E> {
+      E[] build(int size);
+   }
+
+   public interface CollectionBuilder<E, T extends Collection<E>> {
+      T build(int size);
+   }
+
+   public interface UnboundedCollectionBuilder<E, T extends Collection<E>> {
+      T build();
+   }
+
+   public interface MapBuilder<K, V, T extends Map<K, V>> {
+      T build();
    }
 }

commons/src/main/java/org/infinispan/commons/marshall/StreamingMarshaller.java

@@ -51,6 +51,11 @@ public interface StreamingMarshaller extends Marshaller {
     */
    ObjectOutput startObjectOutput(OutputStream os, boolean isReentrant, final int estimatedSize) throws IOException;
 
+   /**
+    * @see {@link #startObjectOutput(OutputStream, boolean, int)}
+    */
+   ObjectOutput startObjectOutput(ObjectOutput oo, boolean isReentrant, final int estimatedSize) throws IOException;
+
    /**
     * Finish using the given ObjectOutput. After opening a ObjectOutput and calling objectToObjectStream() multiple
     * times, use this method to flush the data and close if necessary
@@ -86,6 +91,11 @@ public interface StreamingMarshaller extends Marshaller {
     */
    ObjectInput startObjectInput(InputStream is, boolean isReentrant) throws IOException;
 
+   /**
+    * @see {@link #startObjectInput(InputStream, boolean)}.
+    */
+   ObjectInput startObjectInput(ObjectInput oi, boolean isReentrant) throws IOException;
+
    /**
     * Finish using the given ObjectInput. After opening a ObjectInput and calling objectFromObjectStream() multiple
     * times, use this method to flush the data and close if necessary