Limiting node attribute serialisation in DiscoveryNode (#15341)

RS146BIJAY · web-flow · commit 451606535752 · 2024-09-03T20:11:02.000+05:30
Signed-off-by: RS146BIJAY &lt;rishavsagar4b1@gmail.com&gt;
diff --git a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java
@@ -969,9 +969,13 @@ public <T extends Writeable> void writeOptionalArray(@Nullable T[] array) throws
     }
 
     public void writeOptionalWriteable(@Nullable Writeable writeable) throws IOException {
+        writeOptionalWriteable((out, writable) -> writable.writeTo(out), writeable);
+    }
+
+    public <T extends Writeable> void writeOptionalWriteable(final Writer<T> writer, @Nullable T writeable) throws IOException {
         if (writeable != null) {
             writeBoolean(true);
-            writeable.writeTo(this);
+            writer.write(this, writeable);
         } else {
             writeBoolean(false);
         }
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/allocation/ClusterAllocationExplanation.java b/server/src/main/java/org/opensearch/action/admin/cluster/allocation/ClusterAllocationExplanation.java
@@ -95,7 +95,7 @@ public ClusterAllocationExplanation(StreamInput in) throws IOException {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         shardRouting.writeTo(out);
-        out.writeOptionalWriteable(currentNode);
+        out.writeOptionalWriteable((stream, node) -> node.writeToWithAttribute(stream), currentNode);
         out.writeOptionalWriteable(relocationTargetNode);
         out.writeOptionalWriteable(clusterInfo);
         shardAllocationDecision.writeTo(out);
diff --git a/server/src/main/java/org/opensearch/action/support/nodes/BaseNodeResponse.java b/server/src/main/java/org/opensearch/action/support/nodes/BaseNodeResponse.java
@@ -67,6 +67,6 @@ public DiscoveryNode getNode() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        node.writeTo(out);
+        node.writeToWithAttribute(out);
     }
 }
diff --git a/server/src/main/java/org/opensearch/cluster/ClusterState.java b/server/src/main/java/org/opensearch/cluster/ClusterState.java
@@ -781,7 +781,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(stateUUID);
         metadata.writeTo(out);
         routingTable.writeTo(out);
-        nodes.writeTo(out);
+        nodes.writeToWithAttribute(out);
         blocks.writeTo(out);
         // filter out custom states not supported by the other node
         int numberOfCustoms = 0;
@@ -887,13 +887,23 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeString(toUuid);
             out.writeLong(toVersion);
             routingTable.writeTo(out);
-            nodes.writeTo(out);
+            nodesWriteToWithAttributes(nodes, out);
             metadata.writeTo(out);
             blocks.writeTo(out);
             customs.writeTo(out);
             out.writeVInt(minimumClusterManagerNodesOnPublishingClusterManager);
         }
 
+        private void nodesWriteToWithAttributes(Diff<DiscoveryNodes> nodes, StreamOutput out) throws IOException {
+            DiscoveryNodes part = nodes.apply(null);
+            if (part != null) {
+                out.writeBoolean(true);
+                part.writeToWithAttribute(out);
+            } else {
+                out.writeBoolean(false);
+            }
+        }
+
         @Override
         public ClusterState apply(ClusterState state) {
             Builder builder = new Builder(clusterName);
diff --git a/server/src/main/java/org/opensearch/cluster/coordination/Join.java b/server/src/main/java/org/opensearch/cluster/coordination/Join.java
@@ -78,8 +78,8 @@ public Join(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        sourceNode.writeTo(out);
-        targetNode.writeTo(out);
+        sourceNode.writeToWithAttribute(out);
+        targetNode.writeToWithAttribute(out);
         out.writeLong(term);
         out.writeLong(lastAcceptedTerm);
         out.writeLong(lastAcceptedVersion);
diff --git a/server/src/main/java/org/opensearch/cluster/coordination/JoinRequest.java b/server/src/main/java/org/opensearch/cluster/coordination/JoinRequest.java
@@ -84,7 +84,7 @@ public JoinRequest(StreamInput in) throws IOException {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
-        sourceNode.writeTo(out);
+        sourceNode.writeToWithAttribute(out);
         out.writeLong(minimumTerm);
         out.writeOptionalWriteable(optionalJoin.orElse(null));
     }
diff --git a/server/src/main/java/org/opensearch/cluster/coordination/StartJoinRequest.java b/server/src/main/java/org/opensearch/cluster/coordination/StartJoinRequest.java
@@ -64,7 +64,7 @@ public StartJoinRequest(StreamInput input) throws IOException {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
-        sourceNode.writeTo(out);
+        sourceNode.writeToWithAttribute(out);
         out.writeLong(term);
     }
 
diff --git a/server/src/main/java/org/opensearch/cluster/node/DiscoveryNode.java b/server/src/main/java/org/opensearch/cluster/node/DiscoveryNode.java
@@ -330,6 +330,7 @@ public DiscoveryNode(StreamInput in) throws IOException {
         for (int i = 0; i < size; i++) {
             this.attributes.put(in.readString(), in.readString());
         }
+
         int rolesSize = in.readVInt();
         final Set<DiscoveryNodeRole> roles = new HashSet<>(rolesSize);
         for (int i = 0; i < rolesSize; i++) {
@@ -359,13 +360,31 @@ public DiscoveryNode(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
+        if (out.getVersion().onOrAfter(Version.V_3_0_0)) {
+            writeToUtil(out, false);
+        } else {
+            writeToUtil(out, true);
+        }
+
+    }
+
+    public void writeToWithAttribute(StreamOutput out) throws IOException {
+        writeToUtil(out, true);
+    }
+
+    public void writeToUtil(StreamOutput out, boolean includeAllAttributes) throws IOException {
         writeNodeDetails(out);
 
-        out.writeVInt(attributes.size());
-        for (Map.Entry<String, String> entry : attributes.entrySet()) {
-            out.writeString(entry.getKey());
-            out.writeString(entry.getValue());
+        if (includeAllAttributes) {
+            out.writeVInt(attributes.size());
+            for (Map.Entry<String, String> entry : attributes.entrySet()) {
+                out.writeString(entry.getKey());
+                out.writeString(entry.getValue());
+            }
+        } else {
+            out.writeVInt(0);
         }
+
         writeRolesAndVersion(out);
     }
 
diff --git a/server/src/main/java/org/opensearch/cluster/node/DiscoveryNodes.java b/server/src/main/java/org/opensearch/cluster/node/DiscoveryNodes.java
@@ -690,10 +690,18 @@ public String shortSummary() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
+        writeToUtil((output, value) -> value.writeTo(output), out);
+    }
+
+    public void writeToWithAttribute(StreamOutput out) throws IOException {
+        writeToUtil((output, value) -> value.writeToWithAttribute(output), out);
+    }
+
+    public void writeToUtil(final Writer<DiscoveryNode> writer, StreamOutput out) throws IOException {
         writeClusterManager(out);
         out.writeVInt(nodes.size());
         for (DiscoveryNode node : this) {
-            node.writeTo(out);
+            writer.write(out, node);
         }
     }
 
diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/AbstractAllocationDecision.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/AbstractAllocationDecision.java
@@ -107,7 +107,7 @@ public List<NodeAllocationResult> getNodeDecisions() {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeOptionalWriteable(targetNode);
+        out.writeOptionalWriteable((stream, node) -> node.writeToWithAttribute(stream), targetNode);
         if (nodeDecisions != null) {
             out.writeBoolean(true);
             out.writeList(nodeDecisions);
diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/NodeAllocationResult.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/NodeAllocationResult.java
@@ -104,7 +104,7 @@ public NodeAllocationResult(StreamInput in) throws IOException {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        node.writeTo(out);
+        node.writeToWithAttribute(out);
         out.writeOptionalWriteable(shardStoreInfo);
         out.writeOptionalWriteable(canAllocateDecision);
         nodeDecision.writeTo(out);
diff --git a/server/src/main/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodes.java b/server/src/main/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodes.java
@@ -88,7 +88,12 @@ public UploadedMetadata getUploadedMetadata() {
 
     @Override
     public InputStream serialize() throws IOException {
-        return DISCOVERY_NODES_FORMAT.serialize(discoveryNodes, generateBlobFileName(), getCompressor()).streamInput();
+        return DISCOVERY_NODES_FORMAT.serialize(
+            (out, discoveryNode) -> discoveryNode.writeToWithAttribute(out),
+            discoveryNodes,
+            generateBlobFileName(),
+            getCompressor()
+        ).streamInput();
     }
 
     @Override
diff --git a/server/src/main/java/org/opensearch/indices/recovery/PeerRecoveryTargetService.java b/server/src/main/java/org/opensearch/indices/recovery/PeerRecoveryTargetService.java
@@ -294,6 +294,7 @@ private void doRecovery(final long recoveryId, final StartRecoveryRequest preExi
                 logger.debug("{} reestablishing recovery from {}", startRequest.shardId(), startRequest.sourceNode());
             }
         }
+
         transportService.sendRequest(
             startRequest.sourceNode(),
             actionName,
diff --git a/server/src/main/java/org/opensearch/indices/recovery/StartRecoveryRequest.java b/server/src/main/java/org/opensearch/indices/recovery/StartRecoveryRequest.java
@@ -144,8 +144,8 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeLong(recoveryId);
         shardId.writeTo(out);
         out.writeString(targetAllocationId);
-        sourceNode.writeTo(out);
-        targetNode.writeTo(out);
+        sourceNode.writeToWithAttribute(out);
+        targetNode.writeToWithAttribute(out);
         metadataSnapshot.writeTo(out);
         out.writeBoolean(primaryRelocation);
         out.writeLong(startingSeqNo);
diff --git a/server/src/main/java/org/opensearch/repositories/blobstore/ChecksumWritableBlobStoreFormat.java b/server/src/main/java/org/opensearch/repositories/blobstore/ChecksumWritableBlobStoreFormat.java
@@ -28,6 +28,7 @@
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.common.io.stream.Writeable.Writer;
 import org.opensearch.core.compress.Compressor;
 import org.opensearch.core.compress.CompressorRegistry;
 import org.opensearch.gateway.CorruptStateException;
@@ -56,6 +57,10 @@ public ChecksumWritableBlobStoreFormat(String codec, CheckedFunction<StreamInput
     }
 
     public BytesReference serialize(final T obj, final String blobName, final Compressor compressor) throws IOException {
+        return serialize((out, unSerializedObj) -> unSerializedObj.writeTo(out), obj, blobName, compressor);
+    }
+
+    public BytesReference serialize(final Writer<T> writer, T obj, final String blobName, final Compressor compressor) throws IOException {
         try (BytesStreamOutput outputStream = new BytesStreamOutput()) {
             try (
                 OutputStreamIndexOutput indexOutput = new OutputStreamIndexOutput(
@@ -76,7 +81,7 @@ public void close() throws IOException {
                 }; StreamOutput stream = new OutputStreamStreamOutput(compressor.threadLocalOutputStream(indexOutputOutputStream));) {
                     // TODO The stream version should be configurable
                     stream.setVersion(Version.CURRENT);
-                    obj.writeTo(stream);
+                    writer.write(stream, obj);
                 }
                 CodecUtil.writeFooter(indexOutput);
             }
diff --git a/server/src/main/java/org/opensearch/transport/TransportService.java b/server/src/main/java/org/opensearch/transport/TransportService.java
@@ -752,7 +752,7 @@ public HandshakeResponse(StreamInput in) throws IOException {
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeOptionalWriteable(discoveryNode);
+            out.writeOptionalWriteable((stream, node) -> node.writeToWithAttribute(stream), discoveryNode);
             clusterName.writeTo(out);
             out.writeVersion(version);
         }
diff --git a/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManagerTests.java b/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManagerTests.java
@@ -140,7 +140,12 @@ public void testGetAsyncReadRunnable_DiscoveryNodes() throws IOException, Interr
         DiscoveryNodes discoveryNodes = getDiscoveryNodes();
         String fileName = randomAlphaOfLength(10);
         when(blobStoreTransferService.downloadBlob(anyIterable(), anyString())).thenReturn(
-            DISCOVERY_NODES_FORMAT.serialize(discoveryNodes, fileName, compressor).streamInput()
+            DISCOVERY_NODES_FORMAT.serialize(
+                (out, discoveryNode) -> discoveryNode.writeToWithAttribute(out),
+                discoveryNodes,
+                fileName,
+                compressor
+            ).streamInput()
         );
         RemoteDiscoveryNodes remoteObjForDownload = new RemoteDiscoveryNodes(fileName, "cluster-uuid", compressor);
         CountDownLatch latch = new CountDownLatch(1);
diff --git a/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java b/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java
@@ -1220,7 +1220,12 @@ public void testGetClusterStateUsingDiff() throws IOException {
             diffManifestBuilder.discoveryNodesUpdated(true);
             manifestBuilder.discoveryNodesMetadata(new UploadedMetadataAttribute(DISCOVERY_NODES, DISCOVERY_NODES_FILENAME));
             when(blobContainer.readBlob(DISCOVERY_NODES_FILENAME)).thenAnswer(invocationOnMock -> {
-                BytesReference bytes = DISCOVERY_NODES_FORMAT.serialize(nodesBuilder.build(), DISCOVERY_NODES_FILENAME, compressor);
+                BytesReference bytes = DISCOVERY_NODES_FORMAT.serialize(
+                    (out, nodes) -> nodes.writeToWithAttribute(out),
+                    nodesBuilder.build(),
+                    DISCOVERY_NODES_FILENAME,
+                    compressor
+                );
                 return new ByteArrayInputStream(bytes.streamInput().readAllBytes());
             });
         }
diff --git a/server/src/test/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodesTests.java b/server/src/test/java/org/opensearch/gateway/remote/model/RemoteDiscoveryNodesTests.java
@@ -143,7 +143,7 @@ public void testSerDe() throws IOException {
     public void testExceptionDuringSerialization() throws IOException {
         DiscoveryNodes nodes = mock(DiscoveryNodes.class);
         RemoteDiscoveryNodes remoteObjectForUpload = new RemoteDiscoveryNodes(nodes, METADATA_VERSION, clusterUUID, compressor);
-        doThrow(new IOException("mock-exception")).when(nodes).writeTo(any());
+        doThrow(new IOException("mock-exception")).when(nodes).writeToWithAttribute(any());
         IOException iea = assertThrows(IOException.class, remoteObjectForUpload::serialize);
     }
 
diff --git a/server/src/test/java/org/opensearch/repositories/blobstore/ChecksumWritableBlobStoreFormatTests.java b/server/src/test/java/org/opensearch/repositories/blobstore/ChecksumWritableBlobStoreFormatTests.java
@@ -35,14 +35,20 @@ public class ChecksumWritableBlobStoreFormatTests extends OpenSearchTestCase {
 
     public void testSerDe() throws IOException {
         IndexMetadata indexMetadata = getIndexMetadata();
-        BytesReference bytesReference = clusterBlocksFormat.serialize(indexMetadata, TEST_BLOB_FILE_NAME, CompressorRegistry.none());
+        BytesReference bytesReference = clusterBlocksFormat.serialize(
+            (out, metadata) -> metadata.writeTo(out),
+            indexMetadata,
+            TEST_BLOB_FILE_NAME,
+            CompressorRegistry.none()
+        );
         IndexMetadata readIndexMetadata = clusterBlocksFormat.deserialize(TEST_BLOB_FILE_NAME, bytesReference);
         assertThat(readIndexMetadata, is(indexMetadata));
     }
 
     public void testSerDeForCompressed() throws IOException {
         IndexMetadata indexMetadata = getIndexMetadata();
         BytesReference bytesReference = clusterBlocksFormat.serialize(
+            (out, metadata) -> metadata.writeTo(out),
             indexMetadata,
             TEST_BLOB_FILE_NAME,
             CompressorRegistry.getCompressor(DeflateCompressor.NAME)

Original file line number	Diff line number	Diff line change
`@@ -969,9 +969,13 @@ public <T extends Writeable> void writeOptionalArray(@Nullable T[] array) throws`
`969`	`969`	`}`
`970`	`970`
`971`	`971`	`public void writeOptionalWriteable(@Nullable Writeable writeable) throws IOException {`
	`972`	`+ writeOptionalWriteable((out, writable) -> writable.writeTo(out), writeable);`
	`973`	`+ }`
	`974`	`+`
	`975`	`+ public <T extends Writeable> void writeOptionalWriteable(final Writer<T> writer, @Nullable T writeable) throws IOException {`
`972`	`976`	`if (writeable != null) {`
`973`	`977`	`writeBoolean(true);`
`974`		`- writeable.writeTo(this);`
	`978`	`+ writer.write(this, writeable);`
`975`	`979`	`} else {`
`976`	`980`	`writeBoolean(false);`
`977`	`981`	`}`
Original file line number	Diff line number	Diff line change
`@@ -67,6 +67,6 @@ public DiscoveryNode getNode() {`
`67`	`67`
`68`	`68`	`@Override`
`69`	`69`	`public void writeTo(StreamOutput out) throws IOException {`
`70`		`- node.writeTo(out);`
	`70`	`+ node.writeToWithAttribute(out);`
`71`	`71`	`}`
`72`	`72`	`}`
Original file line number	Diff line number	Diff line change
`@@ -84,7 +84,7 @@ public JoinRequest(StreamInput in) throws IOException {`
`84`	`84`	`@Override`
`85`	`85`	`public void writeTo(StreamOutput out) throws IOException {`
`86`	`86`	`super.writeTo(out);`
`87`		`- sourceNode.writeTo(out);`
	`87`	`+ sourceNode.writeToWithAttribute(out);`
`88`	`88`	`out.writeLong(minimumTerm);`
`89`	`89`	`out.writeOptionalWriteable(optionalJoin.orElse(null));`
`90`	`90`	`}`
Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ public StartJoinRequest(StreamInput input) throws IOException {`
`64`	`64`	`@Override`
`65`	`65`	`public void writeTo(StreamOutput out) throws IOException {`
`66`	`66`	`super.writeTo(out);`
`67`		`- sourceNode.writeTo(out);`
	`67`	`+ sourceNode.writeToWithAttribute(out);`
`68`	`68`	`out.writeLong(term);`
`69`	`69`	`}`
`70`	`70`