Compress and cache cluster state during validate join request

CHANGELOG.md

@@ -49,6 +49,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Change http code on create index API with bad input raising NotXContentException from 500 to 400 ([#4773](https://github.com/opensearch-project/OpenSearch/pull/4773))
 - Change http code for DecommissioningFailedException from 500 to 400 ([#5283](https://github.com/opensearch-project/OpenSearch/pull/5283))
 - Improve summary error message for invalid setting updates ([#4792](https://github.com/opensearch-project/OpenSearch/pull/4792))
+- Compress and cache cluster state during validate join request ([#7321](https://github.com/opensearch-project/OpenSearch/pull/7321))
 
 ### Deprecated
 
@@ -154,4 +155,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Security
 
 [Unreleased 3.0]: https://github.com/opensearch-project/OpenSearch/compare/2.x...HEAD
-[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.7...2.x
+[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.7...2.x

server/src/main/java/org/opensearch/cluster/coordination/ClusterStateUtils.java

@@ -0,0 +1,98 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.cluster.coordination;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.Diff;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.common.bytes.BytesReference;
+import org.opensearch.common.compress.Compressor;
+import org.opensearch.common.compress.CompressorFactory;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.io.stream.InputStreamStreamInput;
+import org.opensearch.common.io.stream.NamedWriteableAwareStreamInput;
+import org.opensearch.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.common.io.stream.OutputStreamStreamOutput;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.io.stream.Writeable;
+import org.opensearch.transport.BytesTransportRequest;
+
+import java.io.IOException;
+
+/**
+ * A helper class to utilize the compressed stream.
+ *
+ * @opensearch.internal
+ */
+public final class ClusterStateUtils {
+    private static final Logger logger = LogManager.getLogger(ClusterStateUtils.class);
+
+    /**
+     * Serialize the given cluster state or diff. It'll always use compression before writing on a newly created output
+     * stream.
+     *
+     * @param writer             Object which is going to write the content
+     * @param node               version of cluster node
+     * @param isFullClusterState flag used at receiver end to make intelligent decisions. For example, ClusterState
+     *                           assumes full state of diff of the states based on this flag.
+     * @return reference to serialized bytes
+     * @throws IOException if writing on the compressed stream is failed.
+     */
+    public static BytesReference serializeClusterState(Writeable writer, DiscoveryNode node, boolean isFullClusterState)
+        throws IOException {
+        final BytesStreamOutput bStream = new BytesStreamOutput();
+        try (StreamOutput stream = new OutputStreamStreamOutput(CompressorFactory.COMPRESSOR.threadLocalOutputStream(bStream))) {
+            stream.setVersion(node.getVersion());
+            stream.writeBoolean(isFullClusterState);
+            writer.writeTo(stream);
+        }
+        final BytesReference serializedByteRef = bStream.bytes();
+        logger.trace("serialized writable object for node version [{}] with size [{}]", node.getVersion(), serializedByteRef.length());
+        return serializedByteRef;
+    }
+
+    /**
+     * Decompress the incoming compressed BytesTransportRequest into StreamInput which can be deserialized.
+     * @param request incoming compressed request in bytes form
+     * @param namedWriteableRegistry existing registry of readers which contains ClusterState writable
+     * @return StreamInput object containing uncompressed request sent by sender
+     * @throws IOException if creating StreamInput object fails due to EOF
+     */
+    public static StreamInput decompressClusterState(BytesTransportRequest request, NamedWriteableRegistry namedWriteableRegistry)
+        throws IOException {
+        final Compressor compressor = CompressorFactory.compressor(request.bytes());
+        StreamInput in = request.bytes().streamInput();
+        if (compressor != null) {
+            in = new InputStreamStreamInput(compressor.threadLocalInputStream(in));
+        }
+        in = new NamedWriteableAwareStreamInput(in, namedWriteableRegistry);
+        in.setVersion(request.version());
+        return in;
+    }
+
+    public static ClusterState deserializeFullClusterState(StreamInput in, DiscoveryNode localNode) throws IOException {
+        final ClusterState incomingState;
+        try (StreamInput input = in) {
+            incomingState = ClusterState.readFrom(input, localNode);
+        }
+        return incomingState;
+    }
+
+    public static Diff<ClusterState> deserializeClusterStateDiff(StreamInput in, DiscoveryNode localNode) throws IOException {
+        final Diff<ClusterState> incomingStateDiff;
+        // Close stream early to release resources used by the de-compression as early as possible
+        try (StreamInput input = in) {
+            incomingStateDiff = ClusterState.readDiffFrom(input, localNode);
+        }
+        return incomingStateDiff;
+    }
+}

server/src/main/java/org/opensearch/cluster/coordination/Coordinator.java

@@ -222,7 +222,8 @@ public Coordinator(
             this.onJoinValidators,
             rerouteService,
             nodeHealthService,
-            this::onNodeCommissionStatusChange
+            this::onNodeCommissionStatusChange,
+            namedWriteableRegistry
         );
         this.persistedStateSupplier = persistedStateSupplier;
         this.noClusterManagerBlockService = new NoClusterManagerBlockService(settings, clusterSettings);

server/src/main/java/org/opensearch/cluster/coordination/JoinHelper.java

@@ -35,6 +35,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.opensearch.Version;
 import org.opensearch.action.ActionListener;
 import org.opensearch.action.ActionListenerResponseHandler;
 import org.opensearch.cluster.ClusterState;
@@ -49,15 +50,19 @@
 import org.opensearch.cluster.routing.allocation.AllocationService;
 import org.opensearch.cluster.service.ClusterManagerService;
 import org.opensearch.common.Priority;
+import org.opensearch.common.bytes.BytesReference;
 import org.opensearch.common.collect.Tuple;
+import org.opensearch.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.monitor.NodeHealthService;
 import org.opensearch.monitor.StatusInfo;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.threadpool.ThreadPool.Names;
+import org.opensearch.transport.BytesTransportRequest;
 import org.opensearch.transport.RemoteTransportException;
 import org.opensearch.transport.TransportChannel;
 import org.opensearch.transport.TransportException;
@@ -69,6 +74,7 @@
 import org.opensearch.transport.TransportService;
 
 import java.io.IOException;
+import java.io.InvalidObjectException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -78,6 +84,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -98,6 +105,7 @@ public class JoinHelper {
 
     public static final String JOIN_ACTION_NAME = "internal:cluster/coordination/join";
     public static final String VALIDATE_JOIN_ACTION_NAME = "internal:cluster/coordination/join/validate";
+    public static final String VALIDATE_COMPRESSED_JOIN_ACTION_NAME = JOIN_ACTION_NAME + "/validate_compressed";
     public static final String START_JOIN_ACTION_NAME = "internal:cluster/coordination/start_join";
 
     // the timeout for Zen1 join attempts
@@ -122,6 +130,8 @@ public class JoinHelper {
 
     private final Supplier<JoinTaskExecutor> joinTaskExecutorGenerator;
     private final Consumer<Boolean> nodeCommissioned;
+    private final NamedWriteableRegistry namedWriteableRegistry;
+    private final ConcurrentHashMap<Long, BytesReference> serializedStates = new ConcurrentHashMap<>();
 
     JoinHelper(
         Settings settings,
@@ -135,13 +145,16 @@ public class JoinHelper {
         Collection<BiConsumer<DiscoveryNode, ClusterState>> joinValidators,
         RerouteService rerouteService,
         NodeHealthService nodeHealthService,
-        Consumer<Boolean> nodeCommissioned
+        Consumer<Boolean> nodeCommissioned,
+        NamedWriteableRegistry namedWriteableRegistry
     ) {
         this.clusterManagerService = clusterManagerService;
         this.transportService = transportService;
         this.nodeHealthService = nodeHealthService;
         this.joinTimeout = JOIN_TIMEOUT_SETTING.get(settings);
         this.nodeCommissioned = nodeCommissioned;
+        this.namedWriteableRegistry = namedWriteableRegistry;
+
         this.joinTaskExecutorGenerator = () -> new JoinTaskExecutor(settings, allocationService, logger, rerouteService) {
 
             private final long term = currentTermSupplier.getAsLong();
@@ -208,22 +221,61 @@ public ClusterTasksResult<JoinTaskExecutor.Task> execute(ClusterState currentSta
             ThreadPool.Names.GENERIC,
             ValidateJoinRequest::new,
             (request, channel, task) -> {
-                final ClusterState localState = currentStateSupplier.get();
-                if (localState.metadata().clusterUUIDCommitted()
-                    && localState.metadata().clusterUUID().equals(request.getState().metadata().clusterUUID()) == false) {
-                    throw new CoordinationStateRejectedException(
-                        "join validation on cluster state"
-                            + " with a different cluster uuid "
-                            + request.getState().metadata().clusterUUID()
-                            + " than local cluster uuid "
-                            + localState.metadata().clusterUUID()
-                            + ", rejecting"
-                    );
-                }
-                joinValidators.forEach(action -> action.accept(transportService.getLocalNode(), request.getState()));
+                runJoinValidators(currentStateSupplier, request.getState(), joinValidators);
+                channel.sendResponse(Empty.INSTANCE);
+            }
+        );
+
+        transportService.registerRequestHandler(
+            VALIDATE_COMPRESSED_JOIN_ACTION_NAME,
+            ThreadPool.Names.GENERIC,
+            BytesTransportRequest::new,
+            (request, channel, task) -> {
+                handleCompressedValidateJoinRequest(currentStateSupplier, joinValidators, request);
                 channel.sendResponse(Empty.INSTANCE);
             }
         );
+
+    }
+
+    private void runJoinValidators(
+        Supplier<ClusterState> currentStateSupplier,
+        ClusterState incomingState,
+        Collection<BiConsumer<DiscoveryNode, ClusterState>> joinValidators
+    ) {
+        final ClusterState localState = currentStateSupplier.get();
+        if (localState.metadata().clusterUUIDCommitted()
+            && localState.metadata().clusterUUID().equals(incomingState.metadata().clusterUUID()) == false) {
+            throw new CoordinationStateRejectedException(
+                "join validation on cluster state"
+                    + " with a different cluster uuid "
+                    + incomingState.metadata().clusterUUID()
+                    + " than local cluster uuid "
+                    + localState.metadata().clusterUUID()
+                    + ", rejecting"
+            );
+        }
+        joinValidators.forEach(action -> action.accept(transportService.getLocalNode(), incomingState));
+    }
+
+    private void handleCompressedValidateJoinRequest(
+        Supplier<ClusterState> currentStateSupplier,
+        Collection<BiConsumer<DiscoveryNode, ClusterState>> joinValidators,
+        BytesTransportRequest request
+    ) throws IOException {
+        StreamInput in = null;
+        try {
+            in = ClusterStateUtils.decompressClusterState(request, namedWriteableRegistry);
+            if (in.readBoolean()) {
+                ClusterState incomingState = ClusterStateUtils.deserializeFullClusterState(in, transportService.getLocalNode());
+                runJoinValidators(currentStateSupplier, incomingState, joinValidators);
+            } else {
+                logger.error("validate new node join request requires full cluster state");
+                throw new InvalidObjectException("validate new node join request requires full cluster state");
+            }
+        } finally {
+            IOUtils.close(in);
+        }
     }
 
     private JoinCallback transportJoinCallback(TransportRequest request, TransportChannel channel) {
@@ -407,12 +459,38 @@ public String executor() {
     }
 
     public void sendValidateJoinRequest(DiscoveryNode node, ClusterState state, ActionListener<TransportResponse.Empty> listener) {
-        transportService.sendRequest(
-            node,
-            VALIDATE_JOIN_ACTION_NAME,
-            new ValidateJoinRequest(state),
-            new ActionListenerResponseHandler<>(listener, i -> Empty.INSTANCE, ThreadPool.Names.GENERIC)
-        );
+        if (node.getVersion().before(Version.V_3_0_0)) {
+            transportService.sendRequest(
+                node,
+                VALIDATE_JOIN_ACTION_NAME,
+                new ValidateJoinRequest(state),
+                new ActionListenerResponseHandler<>(listener, i -> Empty.INSTANCE, ThreadPool.Names.GENERIC)
+            );
+        } else {
+            try {
+                if (!serializedStates.containsKey(state.version())) {
+                    serializedStates.clear();
+                }
+                BytesReference bytes = serializedStates.computeIfAbsent(state.version(), version -> {
+                    try {
+                        return ClusterStateUtils.serializeClusterState(state, node, true);
+                    } catch (IOException e) {
+                        // mandatory as ConcurrentHashMap doesn't rethrow IOException.
+                        throw new RuntimeException(e);
+                    }
+                });
+                final BytesTransportRequest request = new BytesTransportRequest(bytes, node.getVersion());
+                transportService.sendRequest(
+                    node,
+                    VALIDATE_COMPRESSED_JOIN_ACTION_NAME,
+                    request,
+                    new ActionListenerResponseHandler<>(listener, i -> Empty.INSTANCE, ThreadPool.Names.GENERIC)
+                );
+            } catch (Exception e) {
+                logger.warn("error sending cluster state to {}", node);
+                listener.onFailure(e);
+            }
+        }
     }
 
     /**