Optimize checksum creation for remote cluster state

server/src/main/java/org/opensearch/gateway/remote/ClusterStateChecksum.java

@@ -14,6 +14,7 @@
 import org.opensearch.cluster.metadata.DiffableStringMap;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.io.stream.BufferedChecksumStreamOutput;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -22,11 +23,16 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParseException;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
+import java.util.function.Function;
 
 import com.jcraft.jzlib.JZlib;
 
@@ -50,6 +56,7 @@
     static final String INDICES_CS = "indices_md";
     private static final String CLUSTER_STATE_CS = "cluster_state";
     private static final int CHECKSUM_SIZE = 8;
+    private static final int COMPONENT_SIZE = 11;
     private static final Logger logger = LogManager.getLogger(ClusterStateChecksum.class);
 
     long routingTableChecksum;
@@ -65,62 +72,147 @@
     long indicesChecksum;
     long clusterStateChecksum;
 
-    public ClusterStateChecksum(ClusterState clusterState) {
-        try (
-            BytesStreamOutput out = new BytesStreamOutput();
-            BufferedChecksumStreamOutput checksumOut = new BufferedChecksumStreamOutput(out)
-        ) {
-            clusterState.routingTable().writeVerifiableTo(checksumOut);
-            routingTableChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            clusterState.nodes().writeVerifiableTo(checksumOut);
-            nodesChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            clusterState.coordinationMetadata().writeVerifiableTo(checksumOut);
-            coordinationMetadataChecksum = checksumOut.getChecksum();
-
-            // Settings create sortedMap by default, so no explicit sorting required here.
-            checksumOut.reset();
-            Settings.writeSettingsToStream(clusterState.metadata().persistentSettings(), checksumOut);
-            settingMetadataChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            Settings.writeSettingsToStream(clusterState.metadata().transientSettings(), checksumOut);
-            transientSettingsMetadataChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            clusterState.metadata().templatesMetadata().writeVerifiableTo(checksumOut);
-            templatesMetadataChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            checksumOut.writeStringCollection(clusterState.metadata().customs().keySet());
-            customMetadataMapChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            ((DiffableStringMap) clusterState.metadata().hashesOfConsistentSettings()).writeTo(checksumOut);
-            hashesOfConsistentSettingsChecksum = checksumOut.getChecksum();
-
-            checksumOut.reset();
-            checksumOut.writeMapValues(
-                clusterState.metadata().indices(),
-                (stream, value) -> value.writeVerifiableTo((BufferedChecksumStreamOutput) stream)
-            );
-            indicesChecksum = checksumOut.getChecksum();
+    public ClusterStateChecksum(ClusterState clusterState, ThreadPool threadpool) {
+        long start = threadpool.relativeTimeInNanos();
+        ExecutorService executorService = threadpool.executor(ThreadPool.Names.REMOTE_STATE_CHECKSUM);
+        CountDownLatch latch = new CountDownLatch(COMPONENT_SIZE);
 
-            checksumOut.reset();
-            clusterState.blocks().writeVerifiableTo(checksumOut);
-            blocksChecksum = checksumOut.getChecksum();
+        executeChecksumTask((stream) -> {
+            try {
+                clusterState.routingTable().writeVerifiableTo(stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for routing table", e);
+            }
+            return null;
+        }, checksum -> routingTableChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                clusterState.nodes().writeVerifiableTo(stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for discovery nodes", e);
+            }
+            return null;
+        }, checksum -> nodesChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                clusterState.coordinationMetadata().writeVerifiableTo(stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for coordination metadata", e);
+            }
+            return null;
+        }, checksum -> coordinationMetadataChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                Settings.writeSettingsToStream(clusterState.metadata().persistentSettings(), stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for settings metadata", e);
+            }
+            return null;
+        }, checksum -> settingMetadataChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                Settings.writeSettingsToStream(clusterState.metadata().transientSettings(), stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for transient settings metadata", e);
+            }
+            return null;
+        }, checksum -> transientSettingsMetadataChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                clusterState.metadata().templatesMetadata().writeVerifiableTo(stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for templates metadata", e);
+            }
+            return null;
+        }, checksum -> templatesMetadataChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                stream.writeStringCollection(clusterState.metadata().customs().keySet());
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for customs metadata", e);
+            }
+            return null;
+        }, checksum -> customMetadataMapChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                ((DiffableStringMap) clusterState.metadata().hashesOfConsistentSettings()).writeTo(stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for hashesOfConsistentSettings", e);
+            }
+            return null;
+        }, checksum -> hashesOfConsistentSettingsChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                stream.writeMapValues(
+                    clusterState.metadata().indices(),
+                    (checksumStream, value) -> value.writeVerifiableTo((BufferedChecksumStreamOutput) checksumStream)
+                );
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for indices metadata", e);
+            }
+            return null;
+        }, checksum -> indicesChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                clusterState.blocks().writeVerifiableTo(stream);
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for cluster state blocks", e);
+            }
+            return null;
+        }, checksum -> blocksChecksum = checksum, executorService, latch);
+
+        executeChecksumTask((stream) -> {
+            try {
+                stream.writeStringCollection(clusterState.customs().keySet());
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to create checksum for cluster state customs", e);
+            }
+            return null;
+        }, checksum -> clusterStateCustomsChecksum = checksum, executorService, latch);
 
-            checksumOut.reset();
-            checksumOut.writeStringCollection(clusterState.customs().keySet());
-            clusterStateCustomsChecksum = checksumOut.getChecksum();
-        } catch (IOException e) {
-            logger.error("Failed to create checksum for cluster state.", e);
+        try {
+            latch.await();
+        } catch (InterruptedException e) {
             throw new RemoteStateTransferException("Failed to create checksum for cluster state.", e);
         }
         createClusterStateChecksum();
+        logger.debug("Checksum execution time {}", TimeValue.nsecToMSec(threadpool.relativeTimeInNanos() - start));
+    }
+
+    private void executeChecksumTask(
+        Function<BufferedChecksumStreamOutput, Void> checksumTask,
+        Consumer<Long> checksumConsumer,
+        ExecutorService executorService,
+        CountDownLatch latch
+    ) {
+        executorService.execute(() -> {
+            try {
+                long checksum = createChecksum(checksumTask);
+                checksumConsumer.accept(checksum);
+                latch.countDown();
+            } catch (IOException e) {
+                throw new RemoteStateTransferException("Failed to execute checksum task", e);
+            }
+        });
+    }
+
+    private long createChecksum(Function<BufferedChecksumStreamOutput, Void> o) throws IOException {
+        try (
+            BytesStreamOutput out = new BytesStreamOutput();
+            BufferedChecksumStreamOutput checksumOut = new BufferedChecksumStreamOutput(out)
+        ) {
+            o.apply(checksumOut);
+            return checksumOut.getChecksum();
+        }
     }
 
     private void createClusterStateChecksum() {

server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java

@@ -332,7 +332,9 @@
             uploadedMetadataResults,
             previousClusterUUID,
             clusterStateDiffManifest,
-            !remoteClusterStateValidationMode.equals(RemoteClusterStateValidationMode.NONE) ? new ClusterStateChecksum(clusterState) : null,
+            !remoteClusterStateValidationMode.equals(RemoteClusterStateValidationMode.NONE)
+                ? new ClusterStateChecksum(clusterState, threadpool)
+                : null,
             false,
             codecVersion
         );
@@ -539,7 +541,9 @@
             uploadedMetadataResults,
             previousManifest.getPreviousClusterUUID(),
             clusterStateDiffManifest,
-            !remoteClusterStateValidationMode.equals(RemoteClusterStateValidationMode.NONE) ? new ClusterStateChecksum(clusterState) : null,
+            !remoteClusterStateValidationMode.equals(RemoteClusterStateValidationMode.NONE)
+                ? new ClusterStateChecksum(clusterState, threadpool)
+                : null,
             false,
             previousManifest.getCodecVersion()
         );
@@ -1010,7 +1014,9 @@
             uploadedMetadataResults,
             previousManifest.getPreviousClusterUUID(),
             previousManifest.getDiffManifest(),
-            !remoteClusterStateValidationMode.equals(RemoteClusterStateValidationMode.NONE) ? new ClusterStateChecksum(clusterState) : null,
+            !remoteClusterStateValidationMode.equals(RemoteClusterStateValidationMode.NONE)
+                ? new ClusterStateChecksum(clusterState, threadpool)
+                : null,
             true,
             previousManifest.getCodecVersion()
         );
@@ -1631,7 +1637,7 @@
         String localNodeId,
         boolean isFullStateDownload
     ) {
-        ClusterStateChecksum newClusterStateChecksum = new ClusterStateChecksum(clusterState);
+        ClusterStateChecksum newClusterStateChecksum = new ClusterStateChecksum(clusterState, threadpool);
         List<String> failedValidation = newClusterStateChecksum.getMismatchEntities(manifest.getClusterStateChecksum());
         if (failedValidation.isEmpty()) {
             return;

server/src/main/java/org/opensearch/threadpool/ThreadPool.java

@@ -118,6 +118,7 @@ public static class Names {
         public static final String REMOTE_RECOVERY = "remote_recovery";
         public static final String REMOTE_STATE_READ = "remote_state_read";
         public static final String INDEX_SEARCHER = "index_searcher";
+        public static final String REMOTE_STATE_CHECKSUM = "remote_state_checksum";
     }
 
     /**
@@ -191,6 +192,7 @@ public static ThreadPoolType fromType(String type) {
         map.put(Names.REMOTE_RECOVERY, ThreadPoolType.SCALING);
         map.put(Names.REMOTE_STATE_READ, ThreadPoolType.SCALING);
         map.put(Names.INDEX_SEARCHER, ThreadPoolType.RESIZABLE);
+        map.put(Names.REMOTE_STATE_CHECKSUM, ThreadPoolType.SCALING);
         THREAD_POOL_TYPES = Collections.unmodifiableMap(map);
     }
 
@@ -307,6 +309,15 @@ public ThreadPool(
                 runnableTaskListener
             )
         );
+        builders.put(
+            Names.REMOTE_STATE_CHECKSUM,
+            new ScalingExecutorBuilder(
+                Names.REMOTE_STATE_CHECKSUM,
+                1,
+                twiceAllocatedProcessors(allocatedProcessors),
+                TimeValue.timeValueMinutes(5)
+            )
+        );
 
         for (final ExecutorBuilder<?> builder : customBuilders) {
             if (builders.containsKey(builder.name())) {

server/src/test/java/org/opensearch/gateway/remote/ClusterMetadataManifestTests.java

@@ -34,6 +34,9 @@
 import org.opensearch.gateway.remote.ClusterMetadataManifest.UploadedMetadataAttribute;
 import org.opensearch.test.EqualsHashCodeTestUtils;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.TestThreadPool;
+import org.opensearch.threadpool.ThreadPool;
+import org.junit.After;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -64,6 +67,14 @@
 
 public class ClusterMetadataManifestTests extends OpenSearchTestCase {
 
+    private final ThreadPool threadPool = new TestThreadPool(getClass().getName());
+
+    @After
+    public void teardown() throws Exception {
+        super.tearDown();
+        threadPool.shutdown();
+    }
+
     public void testClusterMetadataManifestXContentV0() throws IOException {
         UploadedIndexMetadata uploadedIndexMetadata = new UploadedIndexMetadata("test-index", "test-uuid", "/test/upload/path", CODEC_V0);
         ClusterMetadataManifest originalManifest = ClusterMetadataManifest.builder()
@@ -214,7 +225,7 @@ public void testClusterMetadataManifestSerializationEqualsHashCode() {
                     "indicesRoutingDiffPath"
                 )
             )
-            .checksum(new ClusterStateChecksum(createClusterState()))
+            .checksum(new ClusterStateChecksum(createClusterState(), threadPool))
             .build();
         {  // Mutate Cluster Term
             EqualsHashCodeTestUtils.checkEqualsAndHashCode(
@@ -647,7 +658,7 @@ public void testClusterMetadataManifestXContentV4() throws IOException {
         UploadedIndexMetadata uploadedIndexMetadata = new UploadedIndexMetadata("test-index", "test-uuid", "/test/upload/path");
         UploadedMetadataAttribute uploadedMetadataAttribute = new UploadedMetadataAttribute("attribute_name", "testing_attribute");
         final StringKeyDiffProvider<IndexRoutingTable> routingTableIncrementalDiff = Mockito.mock(StringKeyDiffProvider.class);
-        ClusterStateChecksum checksum = new ClusterStateChecksum(createClusterState());
+        ClusterStateChecksum checksum = new ClusterStateChecksum(createClusterState(), threadPool);
         ClusterMetadataManifest originalManifest = ClusterMetadataManifest.builder()
             .clusterTerm(1L)
             .stateVersion(1L)