[Refactor] Stream Reader and Write Generics

libs/core/src/main/java/org/opensearch/core/common/io/stream/BaseStreamInput.java

@@ -0,0 +1,19 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.core.common.io.stream;
+
+import java.io.InputStream;
+
+/**
+ * Foundation class for reading core types off the transport stream
+ *
+ * todo: refactor {@code StreamInput} primitive readers to this class
+ *
+ * @opensearch.internal
+ */
+public abstract class BaseStreamInput extends InputStream {}

libs/core/src/main/java/org/opensearch/core/common/io/stream/BaseStreamOutput.java

@@ -0,0 +1,19 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.core.common.io.stream;
+
+import java.io.OutputStream;
+
+/**
+ * Foundation class for writing core types over the transport stream
+ *
+ * todo: refactor {@code StreamOutput} primitive writers to this class
+ *
+ * @opensearch.internal
+ */
+public abstract class BaseStreamOutput extends OutputStream {}

libs/core/src/main/java/org/opensearch/core/common/io/stream/BaseWriteable.java

@@ -0,0 +1,130 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.core.common.io.stream;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Implementers can be written to a {@code StreamOutput} and read from a {@code StreamInput}. This allows them to be "thrown
+ * across the wire" using OpenSearch's internal protocol. If the implementer also implements equals and hashCode then a copy made by
+ * serializing and deserializing must be equal and have the same hashCode. It isn't required that such a copy be entirely unchanged.
+ *
+ * @opensearch.internal
+ */
+public interface BaseWriteable<S extends BaseStreamOutput> {
+    /**
+     * A WriteableRegistry registers {@link Writer} methods for writing data types over a
+     * {@link BaseStreamOutput} channel and {@link Reader} methods for reading data from a
+     * {@link BaseStreamInput} channel.
+     *
+     * @opensearch.internal
+     */
+    class WriteableRegistry {
+        private static final Map<Class<?>, Writer<? extends BaseStreamOutput, ?>> WRITER_REGISTRY = new ConcurrentHashMap<>();
+        private static final Map<Byte, Reader<? extends BaseStreamInput, ?>> READER_REGISTRY = new ConcurrentHashMap<>();
+
+        /**
+         * registers a streamable writer
+         *
+         * @opensearch.internal
+         */
+        public static <W extends Writer<? extends BaseStreamOutput, ?>> void registerWriter(final Class<?> clazz, final W writer) {
+            if (WRITER_REGISTRY.containsKey(clazz)) {
+                throw new IllegalArgumentException("Streamable writer already registered for type [" + clazz.getName() + "]");
+            }
+            WRITER_REGISTRY.put(clazz, writer);
+        }
+
+        /**
+         * registers a streamable reader
+         *
+         * @opensearch.internal
+         */
+        public static <R extends Reader<? extends BaseStreamInput, ?>> void registerReader(final byte ordinal, final R reader) {
+            if (READER_REGISTRY.containsKey(ordinal)) {
+                throw new IllegalArgumentException("Streamable reader already registered for ordinal [" + (int) ordinal + "]");
+            }
+            READER_REGISTRY.put(ordinal, reader);
+        }
+
+        /**
+         * Returns the registered writer keyed by the class type
+         */
+        @SuppressWarnings("unchecked")
+        public static <W extends Writer<? extends BaseStreamOutput, ?>> W getWriter(final Class<?> clazz) {
+            return (W) WRITER_REGISTRY.get(clazz);
+        }
+
+        /**
+         * Returns the ristered reader keyed by the unique ordinal
+         */
+        @SuppressWarnings("unchecked")
+        public static <R extends Reader<? extends BaseStreamInput, ?>> R getReader(final byte b) {
+            return (R) READER_REGISTRY.get(b);
+        }
+    }
+
+    /**
+     * Write this into the {@linkplain BaseStreamOutput}.
+     */
+    void writeTo(final S out) throws IOException;
+
+    /**
+     * Reference to a method that can write some object to a {@link BaseStreamOutput}.
+     * <p>
+     * By convention this is a method from {@link BaseStreamOutput} itself (e.g., {@code StreamOutput#writeString}). If the value can be
+     * {@code null}, then the "optional" variant of methods should be used!
+     * <p>
+     * Most classes should implement {@code Writeable} and the {@code Writeable#writeTo(BaseStreamOutput)} method should <em>use</em>
+     * {@link BaseStreamOutput} methods directly or this indirectly:
+     * <pre><code>
+     * public void writeTo(StreamOutput out) throws IOException {
+     *     out.writeVInt(someValue);
+     *     out.writeMapOfLists(someMap, StreamOutput::writeString, StreamOutput::writeString);
+     * }
+     * </code></pre>
+     */
+    @FunctionalInterface
+    interface Writer<S extends BaseStreamOutput, V> {
+
+        /**
+         * Write {@code V}-type {@code value} to the {@code out}put stream.
+         *
+         * @param out Output to write the {@code value} too
+         * @param value The value to add
+         */
+        void write(final S out, V value) throws IOException;
+    }
+
+    /**
+     * Reference to a method that can read some object from a stream. By convention this is a constructor that takes
+     * {@linkplain BaseStreamInput} as an argument for most classes and a static method for things like enums. Returning null from one of these
+     * is always wrong - for that we use methods like {@code StreamInput#readOptionalWriteable(Reader)}.
+     * <p>
+     * As most classes will implement this via a constructor (or a static method in the case of enumerations), it's something that should
+     * look like:
+     * <pre><code>
+     * public MyClass(final StreamInput in) throws IOException {
+     *     this.someValue = in.readVInt();
+     *     this.someMap = in.readMapOfLists(StreamInput::readString, StreamInput::readString);
+     * }
+     * </code></pre>
+     */
+    @FunctionalInterface
+    interface Reader<S extends BaseStreamInput, V> {
+
+        /**
+         * Read {@code V}-type value from a stream.
+         *
+         * @param in Input to read the value from
+         */
+        V read(final S in) throws IOException;
+    }
+}

libs/core/src/main/java/org/opensearch/core/common/io/stream/package-info.java

@@ -0,0 +1,9 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+/** Core transport stream classes */
+package org.opensearch.core.common.io.stream;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/BaseGeoGrid.java

@@ -33,7 +33,6 @@
 
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
-import org.opensearch.common.io.stream.Writeable;
 import org.opensearch.common.util.LongObjectPagedHashMap;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.search.aggregations.InternalAggregation;
@@ -69,7 +68,7 @@ protected BaseGeoGrid(String name, int requiredSize, List<BaseGeoGridBucket> buc
         this.buckets = buckets;
     }
 
-    protected abstract Writeable.Reader<B> getBucketReader();
+    protected abstract Reader<B> getBucketReader();
 
     /**
      * Read from a stream.

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoHashGrid.java

@@ -76,7 +76,7 @@ protected InternalGeoHashGridBucket createBucket(long hashAsLong, long docCount,
     }
 
     @Override
-    protected Reader getBucketReader() {
+    protected Reader<InternalGeoHashGridBucket> getBucketReader() {
         return InternalGeoHashGridBucket::new;
     }
 

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoTileGrid.java

@@ -76,7 +76,7 @@ protected InternalGeoTileGridBucket createBucket(long hashAsLong, long docCount,
     }
 
     @Override
-    protected Reader getBucketReader() {
+    protected Reader<InternalGeoTileGridBucket> getBucketReader() {
         return InternalGeoTileGridBucket::new;
     }
 

server/src/main/java/org/opensearch/common/geo/GeoBoundingBox.java

@@ -79,8 +79,8 @@ public GeoBoundingBox(GeoPoint topLeft, GeoPoint bottomRight) {
     }
 
     public GeoBoundingBox(StreamInput input) throws IOException {
-        this.topLeft = input.readGeoPoint();
-        this.bottomRight = input.readGeoPoint();
+        this.topLeft = new GeoPoint(input);
+        this.bottomRight = new GeoPoint(input);
     }
 
     public boolean isUnbounded() {
@@ -164,8 +164,8 @@ public boolean pointInBounds(double lon, double lat) {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeGeoPoint(topLeft);
-        out.writeGeoPoint(bottomRight);
+        topLeft.writeTo(out);
+        bottomRight.writeTo(out);
     }
 
     @Override

server/src/main/java/org/opensearch/common/geo/GeoPoint.java

@@ -40,6 +40,11 @@
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.OpenSearchParseException;
 import org.opensearch.common.geo.GeoUtils.EffectivePoint;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.BaseWriteable.Reader;
+import org.opensearch.core.common.io.stream.BaseWriteable.Writer;
+import org.opensearch.core.common.io.stream.BaseWriteable.WriteableRegistry;
 import org.opensearch.core.xcontent.ToXContentFragment;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.geometry.Geometry;
@@ -87,6 +92,22 @@ public GeoPoint(GeoPoint template) {
         this(template.getLat(), template.getLon());
     }
 
+    public GeoPoint(final StreamInput in) throws IOException {
+        this.lat = in.readDouble();
+        this.lon = in.readDouble();
+    }
+
+    /**
+     * Register this type as a streamable so it can be serialized over the wire
+     */
+    public static void registerStreamables() {
+        WriteableRegistry.<Writer<StreamOutput, ?>>registerWriter(GeoPoint.class, (o, v) -> {
+            o.writeByte((byte) 22);
+            ((GeoPoint) v).writeTo(o);
+        });
+        WriteableRegistry.<Reader<StreamInput, ?>>registerReader(Byte.valueOf((byte) 22), GeoPoint::new);
+    }
+
     public GeoPoint reset(double lat, double lon) {
         this.lat = lat;
         this.lon = lon;
@@ -210,6 +231,11 @@ public GeoPoint resetFromGeoHash(long geohashLong) {
         return this.resetFromIndexHash(BitUtil.flipFlop((geohashLong >>> 4) << ((level * 5) + 2)));
     }
 
+    public void writeTo(final StreamOutput out) throws IOException {
+        out.writeDouble(this.lat);
+        out.writeDouble(this.lon);
+    }
+
     public double lat() {
         return this.lat;
     }

server/src/main/java/org/opensearch/common/io/stream/StreamInput.java

@@ -50,20 +50,18 @@
 import org.opensearch.common.Strings;
 import org.opensearch.common.bytes.BytesArray;
 import org.opensearch.common.bytes.BytesReference;
-import org.opensearch.common.geo.GeoPoint;
 import org.opensearch.common.settings.SecureString;
 import org.opensearch.common.text.Text;
-import org.opensearch.common.time.DateUtils;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.common.io.stream.BaseStreamInput;
+import org.opensearch.core.common.io.stream.BaseWriteable;
 import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
-import org.opensearch.script.JodaCompatibleZonedDateTime;
 
 import java.io.ByteArrayInputStream;
 import java.io.EOFException;
 import java.io.FileNotFoundException;
 import java.io.FilterInputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.math.BigInteger;
 import java.nio.file.AccessDeniedException;
 import java.nio.file.AtomicMoveNotSupportedException;
@@ -108,7 +106,7 @@
  *
  * @opensearch.internal
  */
-public abstract class StreamInput extends InputStream {
+public abstract class StreamInput extends BaseStreamInput {
 
     private Version version = Version.CURRENT;
 
@@ -686,6 +684,11 @@ public Map<String, Object> readMap() throws IOException {
     @Nullable
     public Object readGenericValue() throws IOException {
         byte type = readByte();
+        BaseWriteable.Reader<StreamInput, ?> r = BaseWriteable.WriteableRegistry.getReader(type);
+        if (r != null) {
+            return r.read(this);
+        }
+
         switch (type) {
             case -1:
                 return null;
@@ -715,8 +718,6 @@ public Object readGenericValue() throws IOException {
                 return readByte();
             case 12:
                 return readDate();
-            case 13:
-                return readDateTime();
             case 14:
                 return readBytesReference();
             case 15:
@@ -733,8 +734,6 @@ public Object readGenericValue() throws IOException {
                 return readDoubleArray();
             case 21:
                 return readBytesRef();
-            case 22:
-                return readGeoPoint();
             case 23:
                 return readZonedDateTime();
             case 24:
@@ -778,14 +777,6 @@ private List readArrayList() throws IOException {
         return list;
     }
 
-    private JodaCompatibleZonedDateTime readDateTime() throws IOException {
-        // we reuse DateTime to communicate with older nodes that don't know about the joda compat layer, but
-        // here we are on a new node so we always want a compat datetime
-        final ZoneId zoneId = DateUtils.dateTimeZoneToZoneId(DateTimeZone.forID(readString()));
-        long millis = readLong();
-        return new JodaCompatibleZonedDateTime(Instant.ofEpochMilli(millis), zoneId);
-    }
-
     private ZonedDateTime readZonedDateTime() throws IOException {
         final String timeZoneId = readString();
         return ZonedDateTime.ofInstant(Instant.ofEpochMilli(readLong()), ZoneId.of(timeZoneId));
@@ -833,13 +824,6 @@ private Date readDate() throws IOException {
         return new Date(readLong());
     }
 
-    /**
-     * Reads a {@link GeoPoint} from this stream input
-     */
-    public GeoPoint readGeoPoint() throws IOException {
-        return new GeoPoint(readDouble(), readDouble());
-    }
-
     /**
      * Read a {@linkplain DateTimeZone}.
      */
@@ -1181,7 +1165,7 @@ public <C extends NamedWriteable> C readOptionalNamedWriteable(Class<C> category
      * @return the list of objects
      * @throws IOException if an I/O exception occurs reading the list
      */
-    public <T> List<T> readList(final Writeable.Reader<T> reader) throws IOException {
+    public <T> List<T> readList(final BaseWriteable.Reader<StreamInput, T> reader) throws IOException {
         return readCollection(reader, ArrayList::new, Collections.emptyList());
     }
 
@@ -1223,8 +1207,11 @@ public <T> Set<T> readSet(Writeable.Reader<T> reader) throws IOException {
     /**
      * Reads a collection of objects
      */
-    private <T, C extends Collection<? super T>> C readCollection(Writeable.Reader<T> reader, IntFunction<C> constructor, C empty)
-        throws IOException {
+    private <T, C extends Collection<? super T>> C readCollection(
+        BaseWriteable.Reader<StreamInput, T> reader,
+        IntFunction<C> constructor,
+        C empty
+    ) throws IOException {
         int count = readArraySize();
         if (count == 0) {
             return empty;