Add general read/write optional support (elastic#112276)

Today `StreamOutput#writeOptionalWriteable` allows to write a
possibly-null value that implements `Writeable` and therefore carries
its own serialization, but sometimes we want to write an optional value
and provide a custom `Writer` too. This commit adds
`StreamOutput#writeOptional` and a corresponding
`StreamInput#readOptional` to support this.

Author: DaveCTurner

server/src/main/java/org/elasticsearch/action/bulk/BulkItemRequest.java

@@ -101,11 +101,11 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalWriteable(primaryResponse);
     }
 
-    public void writeThin(StreamOutput out) throws IOException {
-        out.writeVInt(id);
-        DocWriteRequest.writeDocumentRequestThin(out, request);
-        out.writeOptionalWriteable(primaryResponse == null ? null : primaryResponse::writeThin);
-    }
+    public static final Writer<BulkItemRequest> THIN_WRITER = (out, item) -> {
+        out.writeVInt(item.id);
+        DocWriteRequest.writeDocumentRequestThin(out, item.request);
+        out.writeOptional(BulkItemResponse.THIN_WRITER, item.primaryResponse);
+    };
 
     @Override
     public long ramBytesUsed() {

server/src/main/java/org/elasticsearch/action/bulk/BulkItemResponse.java

@@ -264,15 +264,15 @@ public String toString() {
         id = in.readVInt();
         opType = OpType.fromId(in.readByte());
         response = readResponse(shardId, in);
-        failure = in.readBoolean() ? new Failure(in) : null;
+        failure = in.readOptionalWriteable(Failure::new);
         assertConsistent();
     }
 
     BulkItemResponse(StreamInput in) throws IOException {
         id = in.readVInt();
         opType = OpType.fromId(in.readByte());
         response = readResponse(in);
-        failure = in.readBoolean() ? new Failure(in) : null;
+        failure = in.readOptionalWriteable(Failure::new);
         assertConsistent();
     }
 
@@ -384,31 +384,21 @@ public void writeTo(StreamOutput out) throws IOException {
             writeResponseType(out);
             response.writeTo(out);
         }
-        if (failure == null) {
-            out.writeBoolean(false);
-        } else {
-            out.writeBoolean(true);
-            failure.writeTo(out);
-        }
+        out.writeOptionalWriteable(failure);
     }
 
-    public void writeThin(StreamOutput out) throws IOException {
-        out.writeVInt(id);
-        out.writeByte(opType.getId());
+    public static final Writer<BulkItemResponse> THIN_WRITER = (out, item) -> {
+        out.writeVInt(item.id);
+        out.writeByte(item.opType.getId());
 
-        if (response == null) {
+        if (item.response == null) {
             out.writeByte((byte) 2);
         } else {
-            writeResponseType(out);
-            response.writeThin(out);
+            item.writeResponseType(out);
+            item.response.writeThin(out);
         }
-        if (failure == null) {
-            out.writeBoolean(false);
-        } else {
-            out.writeBoolean(true);
-            failure.writeTo(out);
-        }
-    }
+        out.writeOptionalWriteable(item.failure);
+    };
 
     private void writeResponseType(StreamOutput out) throws IOException {
         if (response instanceof SimulateIndexResponse) {

server/src/main/java/org/elasticsearch/action/bulk/BulkShardRequest.java

@@ -130,14 +130,7 @@ public void writeTo(StreamOutput out) throws IOException {
             throw new IllegalStateException("Inference metadata should have been consumed before writing to the stream");
         }
         super.writeTo(out);
-        out.writeArray((o, item) -> {
-            if (item != null) {
-                o.writeBoolean(true);
-                item.writeThin(o);
-            } else {
-                o.writeBoolean(false);
-            }
-        }, items);
+        out.writeArray((o, item) -> o.writeOptional(BulkItemRequest.THIN_WRITER, item), items);
         if (out.getTransportVersion().onOrAfter(TransportVersions.SIMULATE_VALIDATES_MAPPINGS)) {
             out.writeBoolean(isSimulated);
         }

server/src/main/java/org/elasticsearch/action/bulk/BulkShardResponse.java

@@ -56,6 +56,6 @@ public void setForcedRefresh(boolean forcedRefresh) {
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         shardId.writeTo(out);
-        out.writeArray((o, item) -> item.writeThin(o), responses);
+        out.writeArray(BulkItemResponse.THIN_WRITER, responses);
     }
 }

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

@@ -1095,8 +1095,23 @@ public <T> T[] readOptionalArray(Writeable.Reader<T> reader, IntFunction<T[]> ar
         return readBoolean() ? readArray(reader, arraySupplier) : null;
     }
 
+    /**
+     * Reads a possibly-null value using the given {@link org.elasticsearch.common.io.stream.Writeable.Reader}.
+     *
+     * @see StreamOutput#writeOptionalWriteable
+     */
+    // just an alias for readOptional() since we don't actually care whether T extends Writeable
     @Nullable
     public <T extends Writeable> T readOptionalWriteable(Writeable.Reader<T> reader) throws IOException {
+        return readOptional(reader);
+    }
+
+    /**
+     * Reads a possibly-null value using the given {@link org.elasticsearch.common.io.stream.Writeable.Reader}.
+     *
+     * @see StreamOutput#writeOptional
+     */
+    public <T> T readOptional(Writeable.Reader<T> reader) throws IOException {
         if (readBoolean()) {
             T t = reader.read(this);
             if (t == null) {

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

@@ -1015,6 +1015,12 @@ public <T extends Writeable> void writeOptionalArray(@Nullable T[] array) throws
         writeOptionalArray(StreamOutput::writeWriteable, array);
     }
 
+    /**
+     * Writes a boolean value indicating whether the given object is {@code null}, followed by the object's serialization if it is not
+     * {@code null}.
+     *
+     * @see StreamInput#readOptionalWriteable
+     */
     public void writeOptionalWriteable(@Nullable Writeable writeable) throws IOException {
         if (writeable != null) {
             writeBoolean(true);
@@ -1024,6 +1030,21 @@ public void writeOptionalWriteable(@Nullable Writeable writeable) throws IOExcep
         }
     }
 
+    /**
+     * Writes a boolean value indicating whether the given object is {@code null}, followed by the object's serialization if it is not
+     * {@code null}.
+     *
+     * @see StreamInput#readOptional
+     */
+    public <T> void writeOptional(Writer<T> writer, @Nullable T maybeItem) throws IOException {
+        if (maybeItem != null) {
+            writeBoolean(true);
+            writer.write(this, maybeItem);
+        } else {
+            writeBoolean(false);
+        }
+    }
+
     /**
      * This method allow to use a method reference when writing collection elements such as
      * {@code out.writeMap(map, StreamOutput::writeString, StreamOutput::writeWriteable)}

server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java

@@ -72,8 +72,8 @@ private static Bucket createFromStream(StreamInput in, DocValueFormat format, bo
             String key = in.getTransportVersion().equals(TransportVersions.V_8_0_0) ? in.readString()
                 : in.getTransportVersion().onOrAfter(TransportVersions.V_7_17_1) ? in.readOptionalString()
                 : in.readString();
-            BytesRef from = in.readBoolean() ? in.readBytesRef() : null;
-            BytesRef to = in.readBoolean() ? in.readBytesRef() : null;
+            BytesRef from = in.readOptional(StreamInput::readBytesRef);
+            BytesRef to = in.readOptional(StreamInput::readBytesRef);
             long docCount = in.readLong();
             InternalAggregations aggregations = InternalAggregations.readFrom(in);
 
@@ -89,14 +89,8 @@ public void writeTo(StreamOutput out) throws IOException {
             } else {
                 out.writeString(key == null ? generateKey(from, to, format) : key);
             }
-            out.writeBoolean(from != null);
-            if (from != null) {
-                out.writeBytesRef(from);
-            }
-            out.writeBoolean(to != null);
-            if (to != null) {
-                out.writeBytesRef(to);
-            }
+            out.writeOptional(StreamOutput::writeBytesRef, from);
+            out.writeOptional(StreamOutput::writeBytesRef, to);
             out.writeLong(docCount);
             aggregations.writeTo(out);
         }

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

@@ -761,6 +761,17 @@ public void checkZonedDateTimeSerialization(TransportVersion tv) throws IOExcept
         }
     }
 
+    public void testOptional() throws IOException {
+        try (var output = new BytesStreamOutput()) {
+            output.writeOptional(StreamOutput::writeString, "not-null");
+            output.writeOptional(StreamOutput::writeString, null);
+
+            final var input = getStreamInput(output.bytes());
+            assertEquals("not-null", input.readOptional(StreamInput::readString));
+            assertNull(input.readOptional(StreamInput::readString));
+        }
+    }
+
     private void assertSerialization(
         CheckedConsumer<StreamOutput, IOException> outputAssertions,
         CheckedConsumer<StreamInput, IOException> inputAssertions,

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/rollup/job/RollupJobStatus.java

@@ -74,7 +74,7 @@ public RollupJobStatus(IndexerState state, @Nullable Map<String, Object> positio
 
     public RollupJobStatus(StreamInput in) throws IOException {
         state = IndexerState.fromStream(in);
-        currentPosition = in.readBoolean() ? new TreeMap<>(in.readGenericMap()) : null;
+        currentPosition = in.readOptional(CURRENT_POSITION_READER);
         if (in.getTransportVersion().before(TransportVersions.V_8_0_0)) {
             // 7.x nodes serialize `upgradedDocumentID` flag. We don't need it anymore, but
             // we need to pull it off the stream
@@ -83,6 +83,8 @@ public RollupJobStatus(StreamInput in) throws IOException {
         }
     }
 
+    private static final Reader<TreeMap<String, Object>> CURRENT_POSITION_READER = in -> new TreeMap<>(in.readGenericMap());
+
     public IndexerState getIndexerState() {
         return state;
     }
@@ -118,10 +120,7 @@ public String getWriteableName() {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         state.writeTo(out);
-        out.writeBoolean(currentPosition != null);
-        if (currentPosition != null) {
-            out.writeGenericMap(currentPosition);
-        }
+        out.writeOptional(StreamOutput::writeGenericMap, currentPosition);
         if (out.getTransportVersion().before(TransportVersions.V_8_0_0)) {
             // 7.x nodes expect a boolean `upgradedDocumentID` flag. We don't have it anymore,
             // but we need to tell them we are upgraded in case there is a mixed cluster

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/watcher/transport/actions/execute/ExecuteWatchRequest.java

@@ -59,12 +59,8 @@ public ExecuteWatchRequest(StreamInput in) throws IOException {
         id = in.readOptionalString();
         ignoreCondition = in.readBoolean();
         recordExecution = in.readBoolean();
-        if (in.readBoolean()) {
-            alternativeInput = in.readGenericMap();
-        }
-        if (in.readBoolean()) {
-            triggerData = in.readGenericMap();
-        }
+        alternativeInput = in.readOptional(StreamInput::readGenericMap);
+        triggerData = in.readOptional(StreamInput::readGenericMap);
         long actionModesCount = in.readLong();
         actionModes = new HashMap<>();
         for (int i = 0; i < actionModesCount; i++) {
@@ -83,14 +79,8 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(id);
         out.writeBoolean(ignoreCondition);
         out.writeBoolean(recordExecution);
-        out.writeBoolean(alternativeInput != null);
-        if (alternativeInput != null) {
-            out.writeGenericMap(alternativeInput);
-        }
-        out.writeBoolean(triggerData != null);
-        if (triggerData != null) {
-            out.writeGenericMap(triggerData);
-        }
+        out.writeOptional(StreamOutput::writeGenericMap, alternativeInput);
+        out.writeOptional(StreamOutput::writeGenericMap, triggerData);
         out.writeLong(actionModes.size());
         for (Map.Entry<String, ActionExecutionMode> entry : actionModes.entrySet()) {
             out.writeString(entry.getKey());