Skip to content

Commit

Permalink
[Refactor] Stream Reader and Write Generics (opensearch-project#7465)
Browse files Browse the repository at this point in the history
StreamInput and StreamOutput provide the core functionality for
marshalling / unmarshalling objects over the transport wire. The class
is intended to be generic but it is tightly coupled to the types defined
in the server module. Because of this tight coupling, the classes cannot
be refactored to the core library, thus all new types are required to be
hard coded in the server module.

To decouple this logic and make it more generic across opensearch
modules and plugins, this commit introduces a reader and writer registry
in a new BaseWriteable interface. The StreamInput and StreamOutput also
now inherits from new BaseStreamInput and BaseStreamOutput classes,
respectively, located in the core library. This will be the new home for
streaming primitives in a follow up commit.

Signed-off-by: Nicholas Walter Knize <[email protected]>
  • Loading branch information
nknize committed Jul 5, 2023
1 parent d150117 commit 578db4c
Show file tree
Hide file tree
Showing 19 changed files with 310 additions and 94 deletions.
Original file line number Diff line number Diff line change
@@ -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 {}
Original file line number Diff line number Diff line change
@@ -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 {}
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
@@ -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;
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ protected InternalGeoHashGridBucket createBucket(long hashAsLong, long docCount,
}

@Override
protected Reader getBucketReader() {
protected Reader<InternalGeoHashGridBucket> getBucketReader() {
return InternalGeoHashGridBucket::new;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ protected InternalGeoTileGridBucket createBucket(long hashAsLong, long docCount,
}

@Override
protected Reader getBucketReader() {
protected Reader<InternalGeoTileGridBucket> getBucketReader() {
return InternalGeoTileGridBucket::new;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -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
Expand Down
26 changes: 26 additions & 0 deletions server/src/main/java/org/opensearch/common/geo/GeoPoint.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -108,7 +106,7 @@
*
* @opensearch.internal
*/
public abstract class StreamInput extends InputStream {
public abstract class StreamInput extends BaseStreamInput {

private Version version = Version.CURRENT;

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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}.
*/
Expand Down Expand Up @@ -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());
}

Expand Down Expand Up @@ -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;
Expand Down
Loading

0 comments on commit 578db4c

Please sign in to comment.