From b87e1971e7cfa9646f30d552a6fa57132bbf2992 Mon Sep 17 00:00:00 2001 From: Ryan Ernst Date: Sat, 4 Apr 2020 13:28:30 -0700 Subject: [PATCH] Add generic Set support to streams This commit adds support for reading and writing sets as generic values in stream input and output. closes #54708 --- .../common/io/stream/StreamInput.java | 18 +++++++++++ .../common/io/stream/StreamOutput.java | 17 +++++++++- .../common/io/stream/StreamTests.java | 31 +++++++++++++++++++ 3 files changed, 65 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java b/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java index 7de344d95dc46..7356880e12012 100644 --- a/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java +++ b/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java @@ -70,11 +70,13 @@ 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; import java.util.Set; import java.util.concurrent.TimeUnit; +import java.util.function.Function; import java.util.function.IntFunction; import static org.elasticsearch.ElasticsearchException.readStackTrace; @@ -735,6 +737,10 @@ public Object readGenericValue() throws IOException { 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 + "]"); } @@ -820,6 +826,18 @@ private Map readHashMap() throws IOException { return map10; } + private Set readSet(Function> ctor) throws IOException { + int size = readArraySize(); + if (size == 0) { + return Collections.emptySet(); + } + Set set = ctor.apply(size); + for (int i = 0; i < size; i++) { + set.add(readGenericValue()); + } + return set; + } + private Date readDate() throws IOException { return new Date(readLong()); } diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java b/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java index 7fa616e1c6340..bc59959509b02 100644 --- a/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java +++ b/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java @@ -66,8 +66,10 @@ import java.util.EnumSet; import java.util.Iterator; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.function.IntFunction; @@ -790,7 +792,18 @@ public final void writeOptionalInstant(@Nullable Instant instant) throws IOExcep // joda does not understand "Z" for utc, so we must special case o.writeString(zoneId.equals("Z") ? DateTimeZone.UTC.getID() : zoneId); o.writeLong(zonedDateTime.toInstant().toEpochMilli()); - })); + }), + entry( + Set.class, + (o, v) -> { + if (v instanceof LinkedHashSet) { + o.writeByte((byte) 24); + } else { + o.writeByte((byte) 25); + } + o.writeCollection((Set) v, StreamOutput::writeGenericValue); + } + )); /** * Notice: when serialization a map, the stream out map with the stream in map maybe have the @@ -810,6 +823,8 @@ public void writeGenericValue(@Nullable Object value) throws IOException { 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) { diff --git a/server/src/test/java/org/elasticsearch/common/io/stream/StreamTests.java b/server/src/test/java/org/elasticsearch/common/io/stream/StreamTests.java index ce0bce03b0335..f14c99c4a62c1 100644 --- a/server/src/test/java/org/elasticsearch/common/io/stream/StreamTests.java +++ b/server/src/test/java/org/elasticsearch/common/io/stream/StreamTests.java @@ -21,6 +21,7 @@ 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,11 +39,13 @@ 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; import java.util.Objects; import java.util.Set; +import java.util.function.Consumer; import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -434,4 +437,32 @@ public void testSecureStringSerialization() throws IOException { } } + public void testGenericSet() throws IOException { + Set set = Set.of("a", "b", "c", "d", "e"); + assertGenericRoundtrip(set); + // reverse order in normal set so linked hashset does not match the order + var list = new ArrayList<>(set); + Collections.reverse(list); + assertGenericRoundtrip(new LinkedHashSet<>(list)); + } + + private void assertSerialization(CheckedConsumer outputAssertions, + CheckedConsumer 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)); + }); + } + }