Add StreamInput.readEnum and StreamOutput.writeEnum (#24475)

Implements the common enum serialization/deserialization pattern for enumeration on the StreamInput/StreamOutput.

core/src/main/java/org/elasticsearch/cluster/node/DiscoveryNode.java

@@ -238,11 +238,7 @@ public class DiscoveryNode implements Writeable, ToXContent {
         int rolesSize = in.readVInt();
         this.roles = EnumSet.noneOf(Role.class);
         for (int i = 0; i < rolesSize; i++) {
-            int ordinal = in.readVInt();
+            this.roles.add(in.readEnum(Role.class));
-            if (ordinal < 0 || ordinal >= Role.values().length) {
-                throw new IOException("Unknown Role ordinal [" + ordinal + "]");
-            }
-            this.roles.add(Role.values()[ordinal]);
         }
         this.version = Version.readVersion(in);
     }
@@ -262,7 +258,7 @@ public class DiscoveryNode implements Writeable, ToXContent {
         }
         out.writeVInt(roles.size());
         for (Role role : roles) {
-            out.writeVInt(role.ordinal());
+            out.writeEnum(role);
         }
         Version.writeVersion(version, out);
     }

core/src/main/java/org/elasticsearch/common/geo/ShapeRelation.java

@@ -44,16 +44,12 @@ public enum ShapeRelation implements Writeable {
     }
 
     public static ShapeRelation readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(ShapeRelation.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown ShapeRelation ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(ordinal());
+        out.writeEnum(this);
     }
 
     public static ShapeRelation getRelationByName(String name) {

core/src/main/java/org/elasticsearch/common/geo/SpatialStrategy.java

@@ -40,16 +40,12 @@ public enum SpatialStrategy implements Writeable {
     }
 
     public static SpatialStrategy readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(SpatialStrategy.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown SpatialStrategy ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(ordinal());
+        out.writeEnum(this);
     }
 
     public static SpatialStrategy fromString(String strategyName) {

core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java

@@ -901,6 +901,18 @@ public abstract class StreamInput extends InputStream {
         return builder;
     }
 
+    /**
+     * Reads an enum with type E that was serialized based on the value of it's ordinal
+     */
+    public <E extends Enum<E>> E readEnum(Class<E> enumClass) throws IOException {
+        int ordinal = readVInt();
+        E[] values = enumClass.getEnumConstants();
+        if (ordinal < 0 || ordinal >= values.length) {
+            throw new IOException("Unknown " + enumClass.getSimpleName() + " ordinal [" + ordinal + "]");
+        }
+        return values[ordinal];
+    }
+
     public static StreamInput wrap(byte[] bytes) {
         return wrap(bytes, 0, bytes.length);
     }

core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java

@@ -936,4 +936,11 @@ public abstract class StreamOutput extends OutputStream {
         }
     }
 
+    /**
+     * Writes an enum with type E that by serialized it based on it's ordinal value
+     */
+    public <E extends Enum<E>> void writeEnum(E enumValue) throws IOException {
+        writeVInt(enumValue.ordinal());
+    }
+
 }

core/src/main/java/org/elasticsearch/common/lucene/search/function/CombineFunction.java

@@ -143,15 +143,11 @@ public enum CombineFunction implements Writeable {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(this.ordinal());
+        out.writeEnum(this);
     }
 
     public static CombineFunction readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(CombineFunction.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown CombineFunction ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     public static CombineFunction fromString(String combineFunction) {

core/src/main/java/org/elasticsearch/common/lucene/search/function/FieldValueFactorFunction.java

@@ -191,15 +191,11 @@ public class FieldValueFactorFunction extends ScoreFunction {
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeVInt(this.ordinal());
+            out.writeEnum(this);
         }
 
         public static Modifier readFromStream(StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(Modifier.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown Modifier ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         @Override

core/src/main/java/org/elasticsearch/common/lucene/search/function/FiltersFunctionScoreQuery.java

@@ -81,15 +81,11 @@ public class FiltersFunctionScoreQuery extends Query {
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeVInt(this.ordinal());
+            out.writeEnum(this);
         }
 
         public static ScoreMode readFromStream(StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(ScoreMode.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown ScoreMode ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         public static ScoreMode fromString(String scoreMode) {

core/src/main/java/org/elasticsearch/index/VersionType.java

@@ -364,13 +364,11 @@ public enum VersionType implements Writeable {
     }
 
     public static VersionType readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(VersionType.class);
-        assert (ordinal == 0 || ordinal == 1 || ordinal == 2 || ordinal == 3);
-        return VersionType.values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(ordinal());
+        out.writeEnum(this);
     }
 }

core/src/main/java/org/elasticsearch/index/query/Operator.java

@@ -54,16 +54,12 @@ public enum Operator implements Writeable {
     }
 
     public static Operator readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(Operator.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown Operator ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(this.ordinal());
+        out.writeEnum(this);
     }
 
     public static Operator fromString(String op) {

core/src/main/java/org/elasticsearch/search/MultiValueMode.java

@@ -948,14 +948,10 @@ public enum MultiValueMode implements Writeable {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(this.ordinal());
+        out.writeEnum(this);
     }
 
     public static MultiValueMode readMultiValueModeFrom(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(MultiValueMode.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown MultiValueMode ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 }

core/src/main/java/org/elasticsearch/search/aggregations/Aggregator.java

@@ -139,16 +139,12 @@ public abstract class Aggregator extends BucketCollector implements Releasable {
         }
 
         public static SubAggCollectionMode readFromStream(StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(SubAggCollectionMode.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown SubAggCollectionMode ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeVInt(ordinal());
+            out.writeEnum(this);
         }
     }
 }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/percentiles/PercentilesMethod.java

@@ -53,16 +53,12 @@ public enum PercentilesMethod implements Writeable {
     }
 
     public static PercentilesMethod readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(PercentilesMethod.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown PercentilesMethod ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(ordinal());
+        out.writeEnum(this);
     }
 
     @Override

core/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightBuilder.java

@@ -504,16 +504,12 @@ public class HighlightBuilder extends AbstractHighlighterBuilder<HighlightBuilde
         NONE, SCORE;
 
         public static Order readFromStream(StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(Order.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown Order ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeVInt(this.ordinal());
+            out.writeEnum(this);
         }
 
         public static Order fromString(String order) {
@@ -533,16 +529,12 @@ public class HighlightBuilder extends AbstractHighlighterBuilder<HighlightBuilde
         CHARS, WORD, SENTENCE;
 
         public static BoundaryScannerType readFromStream(StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(BoundaryScannerType.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown BoundaryScannerType ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeVInt(this.ordinal());
+            out.writeEnum(this);
         }
 
         public static BoundaryScannerType fromString(String boundaryScannerType) {

core/src/main/java/org/elasticsearch/search/rescore/QueryRescoreMode.java

@@ -86,16 +86,12 @@ public enum QueryRescoreMode implements Writeable {
     public abstract float combine(float primary, float secondary);
 
     public static QueryRescoreMode readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(QueryRescoreMode.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown ScoreMode ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(this.ordinal());
+        out.writeEnum(this);
     }
 
     public static QueryRescoreMode fromString(String scoreMode) {

core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java

@@ -350,18 +350,14 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
 
         @Override
         public void writeTo(final StreamOutput out) throws IOException {
-            out.writeVInt(ordinal());
+            out.writeEnum(this);
         }
 
         /**
          * Read from a stream.
          */
         static ScriptSortType readFromStream(final StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(ScriptSortType.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown ScriptSortType ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         public static ScriptSortType fromString(final String str) {

core/src/main/java/org/elasticsearch/search/sort/SortMode.java

@@ -52,15 +52,11 @@ public enum SortMode implements Writeable {
 
     @Override
     public void writeTo(final StreamOutput out) throws IOException {
-        out.writeVInt(ordinal());
+        out.writeEnum(this);
     }
 
     public static SortMode readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(SortMode.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown SortMode ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     public static SortMode fromString(final String str) {

core/src/main/java/org/elasticsearch/search/sort/SortOrder.java

@@ -52,16 +52,12 @@ public enum SortOrder implements Writeable {
     };
 
     static SortOrder readFromStream(StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(SortOrder.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown SortOrder ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeVInt(this.ordinal());
+        out.writeEnum(this);
     }
 
     public static SortOrder fromString(String op) {

core/src/main/java/org/elasticsearch/search/suggest/SortBy.java

@@ -38,15 +38,11 @@ public enum SortBy implements Writeable {
 
     @Override
     public void writeTo(final StreamOutput out) throws IOException {
-        out.writeVInt(ordinal());
+        out.writeEnum(this);
     }
 
     public static SortBy readFromStream(final StreamInput in) throws IOException {
-        int ordinal = in.readVInt();
+        return in.readEnum(SortBy.class);
-        if (ordinal < 0 || ordinal >= values().length) {
-            throw new IOException("Unknown SortBy ordinal [" + ordinal + "]");
-        }
-        return values()[ordinal];
     }
 
     public static SortBy resolve(final String str) {

core/src/main/java/org/elasticsearch/search/suggest/term/TermSuggestionBuilder.java

@@ -511,15 +511,11 @@ public class TermSuggestionBuilder extends SuggestionBuilder<TermSuggestionBuild
 
         @Override
         public void writeTo(final StreamOutput out) throws IOException {
-            out.writeVInt(ordinal());
+            out.writeEnum(this);
         }
 
         public static SuggestMode readFromStream(final StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(SuggestMode.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown SuggestMode ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         public static SuggestMode resolve(final String str) {
@@ -571,15 +567,11 @@ public class TermSuggestionBuilder extends SuggestionBuilder<TermSuggestionBuild
 
         @Override
         public void writeTo(final StreamOutput out) throws IOException {
-            out.writeVInt(ordinal());
+            out.writeEnum(this);
         }
 
         public static StringDistanceImpl readFromStream(final StreamInput in) throws IOException {
-            int ordinal = in.readVInt();
+            return in.readEnum(StringDistanceImpl.class);
-            if (ordinal < 0 || ordinal >= values().length) {
-                throw new IOException("Unknown StringDistanceImpl ordinal [" + ordinal + "]");
-            }
-            return values()[ordinal];
         }
 
         public static StringDistanceImpl resolve(final String str) {

core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java

@@ -812,4 +812,34 @@ public class BytesStreamsTests extends ESTestCase {
         StreamInput in = new BytesArray(Base64.getDecoder().decode("////////////AQAAAAAAAA==")).streamInput();
         assertEquals(-1, in.readVLong());
     }
+
+    public enum TestEnum {
+        ONE,
+        TWO,
+        THREE
+    }
+
+    public void testEnum() throws IOException {
+        TestEnum value = randomFrom(TestEnum.values());
+        BytesStreamOutput output = new BytesStreamOutput();
+        output.writeEnum(value);
+        StreamInput input = output.bytes().streamInput();
+        assertEquals(value, input.readEnum(TestEnum.class));
+        assertEquals(0, input.available());
+    }
+
+    public void testInvalidEnum() throws IOException {
+        BytesStreamOutput output = new BytesStreamOutput();
+        int randomNumber = randomInt();
+        boolean validEnum = randomNumber >= 0 && randomNumber < TestEnum.values().length;
+        output.writeVInt(randomNumber);
+        StreamInput input = output.bytes().streamInput();
+        if (validEnum) {
+            assertEquals(TestEnum.values()[randomNumber], input.readEnum(TestEnum.class));
+        } else {
+            IOException ex = expectThrows(IOException.class, () -> input.readEnum(TestEnum.class));
+            assertEquals("Unknown TestEnum ordinal [" + randomNumber + "]", ex.getMessage());
+        }
+        assertEquals(0, input.available());
+    }
 }