Serializing node attribute in discoveryNode only in scenarioes where …

…it is required Signed-off-by: RS146BIJAY <[email protected]>
opensearch-project · Aug 29, 2024 · 6c0e2c4 · 6c0e2c4
1 parent b830d68
commit 6c0e2c4
Show file tree

Hide file tree

Showing 19 changed files with 119 additions and 24 deletions.
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/...ain/java/org/opensearch/action/admin/cluster/allocation/ClusterAllocationExplanation.java b/...ain/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/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;
@@ -859,13 +859,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
@@ -141,8 +141,10 @@ public static boolean isDedicatedSearchNode(Settings settings) {
     private final String hostAddress;
     private final TransportAddress address;
     private final Map<String, String> attributes;
+    private final Map<String, String> filteredAttributes;
     private final Version version;
     private final SortedSet<DiscoveryNodeRole> roles;
+    private static final String ZONE = "zone";
 
     /**
      * Creates a new {@link DiscoveryNode}
@@ -268,6 +270,7 @@ public DiscoveryNode(
             this.version = version;
         }
         this.attributes = Collections.unmodifiableMap(attributes);
+        this.filteredAttributes = Collections.unmodifiableMap(getFilteredAttributes(attributes));
         // verify that no node roles are being provided as attributes
         Predicate<Map<String, String>> predicate = (attrs) -> {
             boolean success = true;
@@ -281,6 +284,23 @@ public DiscoveryNode(
         this.roles = Collections.unmodifiableSortedSet(new TreeSet<>(roles));
     }
 
+    private Map<String, String> getFilteredAttributes(Map<String, String> attributes) {
+        final Map<String, String> filteredAttributes = new HashMap<>();
+        final Set<String> filteredKeySet = attributes.keySet()
+            .stream()
+            .filter(key -> key.startsWith(REMOTE_STORE_NODE_ATTRIBUTE_KEY_PREFIX))
+            .collect(Collectors.toSet());
+        for (String key : filteredKeySet) {
+            filteredAttributes.put(key, attributes.get(key));
+        }
+
+        if (attributes.containsKey(ZONE)) {
+            filteredAttributes.put(ZONE, attributes.get(ZONE));
+        }
+
+        return filteredAttributes;
+    }
+
     /** Creates a DiscoveryNode representing the local node. */
     public static DiscoveryNode createLocal(Settings settings, TransportAddress publishAddress, String nodeId) {
         Map<String, String> attributes = Node.NODE_ATTRIBUTES.getAsMap(settings);
@@ -329,6 +349,8 @@ public DiscoveryNode(StreamInput in) throws IOException {
         for (int i = 0; i < size; i++) {
             this.attributes.put(in.readString(), in.readString());
         }
+
+        this.filteredAttributes = Collections.unmodifiableMap(getFilteredAttributes(attributes));
         int rolesSize = in.readVInt();
         final Set<DiscoveryNodeRole> roles = new HashSet<>(rolesSize);
         for (int i = 0; i < rolesSize; i++) {
@@ -358,17 +380,33 @@ 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);
+    }
+
+    private void writeToUtil(StreamOutput out, boolean includeAllAttributes) throws IOException {
         out.writeString(nodeName);
         out.writeString(nodeId);
         out.writeString(ephemeralId);
         out.writeString(hostName);
         out.writeString(hostAddress);
         address.writeTo(out);
-        out.writeVInt(attributes.size());
-        for (Map.Entry<String, String> entry : attributes.entrySet()) {
-            out.writeString(entry.getKey());
-            out.writeString(entry.getValue());
+        if (includeAllAttributes) {
+            serializeAttributes(attributes, out);
+        } else {
+            // Serialize only remote store and zone (needed for SearchWeightedRouting) attributes if present.
+
+            serializeAttributes(filteredAttributes, out);
         }
+
         out.writeVInt(roles.size());
         for (final DiscoveryNodeRole role : roles) {
             final DiscoveryNodeRole compatibleRole = role.getCompatibilityRole(out.getVersion());
@@ -379,6 +417,14 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeVersion(version);
     }
 
+    private void serializeAttributes(final Map<String, String> attributes, final StreamOutput out) throws IOException {
+        out.writeVInt(attributes.size());
+        for (Map.Entry<String, String> entry : attributes.entrySet()) {
+            out.writeString(entry.getKey());
+            out.writeString(entry.getValue());
+        }
+    }
+
     /**
      * The address that the node can be communicated with.
      */

diff --git a/server/src/main/java/org/opensearch/cluster/node/DiscoveryNodes.java b/server/src/main/java/org/opensearch/cluster/node/DiscoveryNodes.java
@@ -688,6 +688,14 @@ 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);
+    }
+
+    private void writeToUtil(final Writer<DiscoveryNode> writer, StreamOutput out) throws IOException {
         if (clusterManagerNodeId == null) {
             out.writeBoolean(false);
         } else {
@@ -696,7 +704,7 @@ public void writeTo(StreamOutput out) throws IOException {
         }
         out.writeVInt(nodes.size());
         for (DiscoveryNode node : this) {
-            node.writeTo(out);
+            writer.write(out, node);
         }
     }
 

diff --git a/...r/src/main/java/org/opensearch/cluster/routing/allocation/AbstractAllocationDecision.java b/...r/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/.../src/main/java/org/opensearch/repositories/blobstore/ChecksumWritableBlobStoreFormat.java b/.../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/...src/test/java/org/opensearch/gateway/remote/RemoteClusterStateAttributesManagerTests.java b/...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
@@ -1172,7 +1172,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/...test/java/org/opensearch/repositories/blobstore/ChecksumWritableBlobStoreFormatTests.java b/...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)