[Backport 2.x] [Refactor] Stream Reader and Write Generics (#7465)

CHANGELOG.md

@@ -82,6 +82,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Check UTF16 string size before converting to String to avoid OOME ([#7963](https://github.com/opensearch-project/OpenSearch/pull/7963))
 - [Refactor] remaining ImmutableOpenMap usage to j.u.Map and remove class ([#7309](https://github.com/opensearch-project/OpenSearch/pull/7309))
 - [Refactor] Version and LegacyESVersion from server module to core lib ([#7328](https://github.com/opensearch-project/OpenSearch/pull/7328))
+- [Refactor] Stream Reader and Write Generics ([#7465](https://github.com/opensearch-project/OpenSearch/pull/7465))
 - Move ZSTD compression codecs out of the sandbox ([#7908](https://github.com/opensearch-project/OpenSearch/pull/7908))
 
 ### Deprecated

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/InternalGeoGrid.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 InternalGeoGrid(String name, int requiredSize, List<InternalGeoGridBuc
         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/InternalGeoHashGrid.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/InternalGeoTileGrid.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;
     }