Add generic Set support to streams (#54769) (#55123)

This commit adds support for reading and writing sets as generic values
in stream input and output.

closes #54708

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

@@ -70,6 +70,7 @@ import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -735,6 +736,10 @@ public abstract class StreamInput extends InputStream {
                 return readGeoPoint();
             case 23:
                 return readZonedDateTime();
+            case 24:
+                return readCollection(StreamInput::readGenericValue, LinkedHashSet::new, Collections.emptySet());
+            case 25:
+                return readCollection(StreamInput::readGenericValue, HashSet::new, Collections.emptySet());
             default:
                 throw new IOException("Can't read unknown type [" + type + "]");
         }

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

@@ -70,6 +70,7 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -776,6 +777,14 @@ public abstract class StreamOutput extends OutputStream {
             o.writeString(zoneId.equals("Z") ? DateTimeZone.UTC.getID() : zoneId);
             o.writeLong(zonedDateTime.toInstant().toEpochMilli());
         });
+        writers.put(Set.class, (o, v) -> {
+            if (v instanceof LinkedHashSet) {
+                o.writeByte((byte) 24);
+            } else {
+                o.writeByte((byte) 25);
+            }
+            o.writeCollection((Set<?>) v, StreamOutput::writeGenericValue);
+        });
         WRITERS = Collections.unmodifiableMap(writers);
     }
 
@@ -797,6 +806,8 @@ public abstract class StreamOutput extends OutputStream {
             type = Object[].class;
         } else if (value instanceof Map) {
             type = Map.class;
+        } else if (value instanceof Set) {
+            type = Set.class;
         } else if (value instanceof ReadableInstant) {
             type = ReadableInstant.class;
         } else if (value instanceof BytesReference) {

server/src/test/java/org/elasticsearch/common/io/stream/StreamTests.java

@@ -21,6 +21,7 @@ package org.elasticsearch.common.io.stream;
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.CheckedBiConsumer;
+import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.common.CheckedFunction;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
@@ -38,6 +39,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -434,4 +436,32 @@ public class StreamTests extends ESTestCase {
         }
     }
 
+    public void testGenericSet() throws IOException {
+        Set<String> set = new HashSet<>(Arrays.asList("a", "b", "c", "d", "e"));
+        assertGenericRoundtrip(set);
+        // reverse order in normal set so linked hashset does not match the order
+        List<String> list = new ArrayList<>(set);
+        Collections.reverse(list);
+        assertGenericRoundtrip(new LinkedHashSet<>(list));
+    }
+
+    private void assertSerialization(CheckedConsumer<StreamOutput, IOException> outputAssertions,
+                                     CheckedConsumer<StreamInput, IOException> inputAssertions) throws IOException {
+        try (BytesStreamOutput output = new BytesStreamOutput()) {
+            outputAssertions.accept(output);
+            final BytesReference bytesReference = output.bytes();
+            final StreamInput input = bytesReference.streamInput();
+            inputAssertions.accept(input);
+        }
+    }
+
+    private void assertGenericRoundtrip(Object original) throws IOException {
+        assertSerialization(output -> {
+            output.writeGenericValue(original);
+        }, input -> {
+            Object read = input.readGenericValue();
+            assertThat(read, equalTo(original));
+        });
+    }
+
 }