Optimise TransportNodesAction to not send DiscoveryNodes for NodeStat…

CHANGELOG.md

@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add Plugin interface for loading application based configuration templates (([#14659](https://github.com/opensearch-project/OpenSearch/issues/14659)))
 - Refactor remote-routing-table service inline with remote state interfaces([#14668](https://github.com/opensearch-project/OpenSearch/pull/14668))
 - Add prefix mode verification setting for repository verification (([#14790](https://github.com/opensearch-project/OpenSearch/pull/14790)))
+- Optimize TransportNodesAction to not send DiscoveryNodes for NodeStats, NodesInfo and ClusterStats call ([14749](https://github.com/opensearch-project/OpenSearch/pull/14749))
 
 ### Dependencies
 - Bump `org.gradle.test-retry` from 1.5.8 to 1.5.9 ([#13442](https://github.com/opensearch-project/OpenSearch/pull/13442))

server/src/main/java/org/opensearch/action/admin/cluster/node/info/TransportNodesInfoAction.java

@@ -129,7 +129,7 @@ protected NodeInfo nodeOperation(NodeInfoRequest nodeRequest) {
      */
     public static class NodeInfoRequest extends TransportRequest {
 
-        NodesInfoRequest request;
+        protected NodesInfoRequest request;
 
         public NodeInfoRequest(StreamInput in) throws IOException {
             super(in);

server/src/main/java/org/opensearch/action/admin/cluster/node/stats/TransportNodesStatsAction.java

@@ -140,7 +140,7 @@ protected NodeStats nodeOperation(NodeStatsRequest nodeStatsRequest) {
      */
     public static class NodeStatsRequest extends TransportRequest {
 
-        NodesStatsRequest request;
+        protected NodesStatsRequest request;
 
         public NodeStatsRequest(StreamInput in) throws IOException {
             super(in);

server/src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java

@@ -223,7 +223,7 @@ protected ClusterStatsNodeResponse nodeOperation(ClusterStatsNodeRequest nodeReq
      */
     public static class ClusterStatsNodeRequest extends TransportRequest {
 
-        ClusterStatsRequest request;
+        protected ClusterStatsRequest request;
 
         public ClusterStatsNodeRequest(StreamInput in) throws IOException {
             super(in);

server/src/main/java/org/opensearch/action/support/nodes/BaseNodesRequest.java

@@ -65,6 +65,14 @@ public abstract class BaseNodesRequest<Request extends BaseNodesRequest<Request>
      * will be ignored and this will be used.
      * */
     private DiscoveryNode[] concreteNodes;
+
+    /**
+     * Since do not use the discovery nodes coming from the request in all code paths following a request extended off from
+     * BaseNodeRequest, we do not require it to sent around across all nodes.
+     *
+     * Setting default behavior as `true` but can be explicitly changed in requests that do not require.
+     */
+    private boolean sendDiscoveryNodes = true;
     private final TimeValue DEFAULT_TIMEOUT_SECS = TimeValue.timeValueSeconds(30);
 
     private TimeValue timeout;
@@ -119,6 +127,14 @@ public void setConcreteNodes(DiscoveryNode[] concreteNodes) {
         this.concreteNodes = concreteNodes;
     }
 
+    public void sendDiscoveryNodes(boolean value) {
+        sendDiscoveryNodes = value;
+    }
+
+    public boolean sendDiscoveryNodes() {
+        return sendDiscoveryNodes;
+    }
+
     @Override
     public ActionRequestValidationException validate() {
         return null;

server/src/main/java/org/opensearch/action/support/nodes/TransportNodesAction.java

@@ -226,6 +226,7 @@ class AsyncAction {
         private final NodesRequest request;
         private final ActionListener<NodesResponse> listener;
         private final AtomicReferenceArray<Object> responses;
+        private final DiscoveryNode[] concreteNodes;
         private final AtomicInteger counter = new AtomicInteger();
         private final Task task;
 
@@ -238,10 +239,18 @@ class AsyncAction {
                 assert request.concreteNodes() != null;
             }
             this.responses = new AtomicReferenceArray<>(request.concreteNodes().length);
+            this.concreteNodes = request.concreteNodes();
+
+            if (request.sendDiscoveryNodes() == false) {
+                // As we transfer the ownership of discovery nodes to route the request to into the AsyncAction class, we
+                // remove the list of DiscoveryNodes from the request. This reduces the payload of the request and improves
+                // the number of concrete nodes in the memory.
+                request.setConcreteNodes(null);
+            }
         }
 
         void start() {
-            final DiscoveryNode[] nodes = request.concreteNodes();
+            final DiscoveryNode[] nodes = this.concreteNodes;
             if (nodes.length == 0) {
                 // nothing to notify
                 threadPool.generic().execute(() -> listener.onResponse(newResponse(request, responses)));
@@ -260,7 +269,6 @@ void start() {
                     if (task != null) {
                         nodeRequest.setParentTask(clusterService.localNode().getId(), task.getId());
                     }
-
                     transportService.sendRequest(
                         node,
                         getTransportNodeAction(node),

server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterStatsAction.java

@@ -66,6 +66,7 @@ public String getName() {
     public RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException {
         ClusterStatsRequest clusterStatsRequest = new ClusterStatsRequest().nodesIds(request.paramAsStringArray("nodeId", null));
         clusterStatsRequest.timeout(request.param("timeout"));
+        clusterStatsRequest.sendDiscoveryNodes(false);
         return channel -> client.admin().cluster().clusterStats(clusterStatsRequest, new NodesResponseRestListener<>(channel));
     }
 

server/src/main/java/org/opensearch/rest/action/admin/cluster/RestNodesInfoAction.java

@@ -88,7 +88,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         final NodesInfoRequest nodesInfoRequest = prepareRequest(request);
         nodesInfoRequest.timeout(request.param("timeout"));
         settingsFilter.addFilterSettingParams(request);
-
+        nodesInfoRequest.sendDiscoveryNodes(false);
         return channel -> client.admin().cluster().nodesInfo(nodesInfoRequest, new NodesResponseRestListener<>(channel));
     }
 

server/src/main/java/org/opensearch/rest/action/admin/cluster/RestNodesStatsAction.java

@@ -232,6 +232,7 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         // If no levels are passed in this results in an empty array.
         String[] levels = Strings.splitStringByCommaToArray(request.param("level"));
         nodesStatsRequest.indices().setLevels(levels);
+        nodesStatsRequest.sendDiscoveryNodes(false);
 
         return channel -> client.admin().cluster().nodesStats(nodesStatsRequest, new NodesResponseRestListener<>(channel));
     }

server/src/main/java/org/opensearch/rest/action/cat/RestNodesAction.java

@@ -125,6 +125,7 @@
             public void processResponse(final ClusterStateResponse clusterStateResponse) {
                 NodesInfoRequest nodesInfoRequest = new NodesInfoRequest();
                 nodesInfoRequest.timeout(request.param("timeout"));
+                nodesInfoRequest.sendDiscoveryNodes(false);
                 nodesInfoRequest.clear()
                     .addMetrics(
                         NodesInfoRequest.Metric.JVM.metricName(),
@@ -137,6 +138,7 @@
                     public void processResponse(final NodesInfoResponse nodesInfoResponse) {
                         NodesStatsRequest nodesStatsRequest = new NodesStatsRequest();
                         nodesStatsRequest.timeout(request.param("timeout"));
+                        nodesStatsRequest.sendDiscoveryNodes(false);
                         nodesStatsRequest.clear()
                             .indices(true)
                             .addMetrics(

server/src/test/java/org/opensearch/action/RestStatsActionTests.java

@@ -0,0 +1,59 @@
+/*
+ * 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.action;
+
+import org.opensearch.client.node.NodeClient;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.settings.SettingsFilter;
+import org.opensearch.rest.action.admin.cluster.RestClusterStatsAction;
+import org.opensearch.rest.action.admin.cluster.RestNodesInfoAction;
+import org.opensearch.rest.action.admin.cluster.RestNodesStatsAction;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.test.rest.FakeRestRequest;
+import org.opensearch.threadpool.TestThreadPool;
+import org.junit.After;
+
+import java.util.Collections;
+
+public class RestStatsActionTests extends OpenSearchTestCase {
+    private final TestThreadPool threadPool = new TestThreadPool(RestStatsActionTests.class.getName());
+    private final NodeClient client = new NodeClient(Settings.EMPTY, threadPool);
+
+    @After
+    public void terminateThreadPool() {
+        terminate(threadPool);
+    }
+
+    public void testClusterStatsActionPrepareRequestNoError() {
+        RestClusterStatsAction action = new RestClusterStatsAction();
+        try {
+            action.prepareRequest(new FakeRestRequest(), client);
+        } catch (Throwable t) {
+            fail(t.getMessage());
+        }
+    }
+
+    public void testNodesStatsActionPrepareRequestNoError() {
+        RestNodesStatsAction action = new RestNodesStatsAction();
+        try {
+            action.prepareRequest(new FakeRestRequest(), client);
+        } catch (Throwable t) {
+            fail(t.getMessage());
+        }
+    }
+
+    public void testNodesInfoActionPrepareRequestNoError() {
+        RestNodesInfoAction action = new RestNodesInfoAction(new SettingsFilter(Collections.singleton("foo.filtered")));
+        try {
+            action.prepareRequest(new FakeRestRequest(), client);
+        } catch (Throwable t) {
+            fail(t.getMessage());
+        }
+    }
+}

server/src/test/java/org/opensearch/action/support/nodes/TransportClusterStatsActionTests.java

@@ -0,0 +1,165 @@
+/*
+ * 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.action.support.nodes;
+
+import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest;
+import org.opensearch.action.admin.cluster.stats.ClusterStatsRequest;
+import org.opensearch.action.admin.cluster.stats.TransportClusterStatsAction;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.indices.IndicesService;
+import org.opensearch.node.NodeService;
+import org.opensearch.test.transport.CapturingTransport;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class TransportClusterStatsActionTests extends TransportNodesActionTests {
+
+    /**
+     * By default, we send discovery nodes list to each request that is sent across from the coordinator node. This
+     * behavior is asserted in this test.
+     */
+    public void testClusterStatsActionWithRetentionOfDiscoveryNodesList() {
+        ClusterStatsRequest request = new ClusterStatsRequest();
+        request.sendDiscoveryNodes(true);
+        Map<String, List<MockClusterStatsNodeRequest>> combinedSentRequest = performNodesInfoAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> {
+                assertNotNull(sentRequest.getDiscoveryNodes());
+                assertEquals(sentRequest.getDiscoveryNodes().length, clusterService.state().nodes().getSize());
+            });
+        });
+    }
+
+    public void testClusterStatsActionWithPreFilledConcreteNodesAndWithRetentionOfDiscoveryNodesList() {
+        ClusterStatsRequest request = new ClusterStatsRequest();
+        Collection<DiscoveryNode> discoveryNodes = clusterService.state().getNodes().getNodes().values();
+        request.setConcreteNodes(discoveryNodes.toArray(DiscoveryNode[]::new));
+        Map<String, List<MockClusterStatsNodeRequest>> combinedSentRequest = performNodesInfoAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> {
+                assertNotNull(sentRequest.getDiscoveryNodes());
+                assertEquals(sentRequest.getDiscoveryNodes().length, clusterService.state().nodes().getSize());
+            });
+        });
+    }
+
+    /**
+     * In the optimized ClusterStats Request, we do not send the DiscoveryNodes List to each node. This behavior is
+     * asserted in this test.
+     */
+    public void testClusterStatsActionWithoutRetentionOfDiscoveryNodesList() {
+        ClusterStatsRequest request = new ClusterStatsRequest();
+        request.sendDiscoveryNodes(false);
+        Map<String, List<MockClusterStatsNodeRequest>> combinedSentRequest = performNodesInfoAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.getDiscoveryNodes()); });
+        });
+    }
+
+    public void testClusterStatsActionWithPreFilledConcreteNodesAndWithoutRetentionOfDiscoveryNodesList() {
+        ClusterStatsRequest request = new ClusterStatsRequest();
+        Collection<DiscoveryNode> discoveryNodes = clusterService.state().getNodes().getNodes().values();
+        request.setConcreteNodes(discoveryNodes.toArray(DiscoveryNode[]::new));
+        request.sendDiscoveryNodes(false);
+        Map<String, List<MockClusterStatsNodeRequest>> combinedSentRequest = performNodesInfoAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.getDiscoveryNodes()); });
+        });
+    }
+
+    private Map<String, List<MockClusterStatsNodeRequest>> performNodesInfoAction(ClusterStatsRequest request) {
+        TransportNodesAction action = getTestTransportClusterStatsAction();
+        PlainActionFuture<NodesStatsRequest> listener = new PlainActionFuture<>();
+        action.new AsyncAction(null, request, listener).start();
+        Map<String, List<CapturingTransport.CapturedRequest>> capturedRequests = transport.getCapturedRequestsByTargetNodeAndClear();
+        Map<String, List<MockClusterStatsNodeRequest>> combinedSentRequest = new HashMap<>();
+
+        capturedRequests.forEach((node, capturedRequestList) -> {
+            List<MockClusterStatsNodeRequest> sentRequestList = new ArrayList<>();
+
+            capturedRequestList.forEach(preSentRequest -> {
+                BytesStreamOutput out = new BytesStreamOutput();
+                try {
+                    TransportClusterStatsAction.ClusterStatsNodeRequest clusterStatsNodeRequestFromCoordinator =
+                        (TransportClusterStatsAction.ClusterStatsNodeRequest) preSentRequest.request;
+                    clusterStatsNodeRequestFromCoordinator.writeTo(out);
+                    StreamInput in = out.bytes().streamInput();
+                    MockClusterStatsNodeRequest mockClusterStatsNodeRequest = new MockClusterStatsNodeRequest(in);
+                    sentRequestList.add(mockClusterStatsNodeRequest);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+
+            combinedSentRequest.put(node, sentRequestList);
+        });
+
+        return combinedSentRequest;
+    }
+
+    private TestTransportClusterStatsAction getTestTransportClusterStatsAction() {
+        return new TestTransportClusterStatsAction(
+            THREAD_POOL,
+            clusterService,
+            transportService,
+            nodeService,
+            indicesService,
+            new ActionFilters(Collections.emptySet())
+        );
+    }
+
+    private static class TestTransportClusterStatsAction extends TransportClusterStatsAction {
+        public TestTransportClusterStatsAction(
+            ThreadPool threadPool,
+            ClusterService clusterService,
+            TransportService transportService,
+            NodeService nodeService,
+            IndicesService indicesService,
+            ActionFilters actionFilters
+        ) {
+            super(threadPool, clusterService, transportService, nodeService, indicesService, actionFilters);
+        }
+    }
+
+    private static class MockClusterStatsNodeRequest extends TransportClusterStatsAction.ClusterStatsNodeRequest {
+
+        public MockClusterStatsNodeRequest(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        public DiscoveryNode[] getDiscoveryNodes() {
+            return this.request.concreteNodes();
+        }
+    }
+}

server/src/test/java/org/opensearch/action/support/nodes/TransportNodesActionTests.java

@@ -46,6 +46,8 @@
 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.indices.IndicesService;
+import org.opensearch.node.NodeService;
 import org.opensearch.telemetry.tracing.noop.NoopTracer;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.transport.CapturingTransport;
@@ -76,11 +78,12 @@
 
 public class TransportNodesActionTests extends OpenSearchTestCase {
 
-    private static ThreadPool THREAD_POOL;
-
-    private ClusterService clusterService;
-    private CapturingTransport transport;
-    private TransportService transportService;
+    protected static ThreadPool THREAD_POOL;
+    protected ClusterService clusterService;
+    protected CapturingTransport transport;
+    protected TransportService transportService;
+    protected NodeService nodeService;
+    protected IndicesService indicesService;
 
     public void testRequestIsSentToEachNode() throws Exception {
         TransportNodesAction action = getTestTransportNodesAction();

server/src/test/java/org/opensearch/action/support/nodes/TransportNodesInfoActionTests.java

@@ -0,0 +1,131 @@
+/*
+ * 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.action.support.nodes;
+
+import org.opensearch.action.admin.cluster.node.info.NodesInfoRequest;
+import org.opensearch.action.admin.cluster.node.info.TransportNodesInfoAction;
+import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.node.NodeService;
+import org.opensearch.test.transport.CapturingTransport;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class TransportNodesInfoActionTests extends TransportNodesActionTests {
+
+    /**
+     * By default, we send discovery nodes list to each request that is sent across from the coordinator node. This
+     * behavior is asserted in this test.
+     */
+    public void testNodesInfoActionWithRetentionOfDiscoveryNodesList() {
+        NodesInfoRequest request = new NodesInfoRequest();
+        request.sendDiscoveryNodes(true);
+        Map<String, List<MockNodesInfoRequest>> combinedSentRequest = performNodesInfoAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> {
+                assertNotNull(sentRequest.getDiscoveryNodes());
+                assertEquals(sentRequest.getDiscoveryNodes().length, clusterService.state().nodes().getSize());
+            });
+        });
+    }
+
+    /**
+     * In the optimized ClusterStats Request, we do not send the DiscoveryNodes List to each node. This behavior is
+     * asserted in this test.
+     */
+    public void testNodesInfoActionWithoutRetentionOfDiscoveryNodesList() {
+        NodesInfoRequest request = new NodesInfoRequest();
+        request.sendDiscoveryNodes(false);
+        Map<String, List<MockNodesInfoRequest>> combinedSentRequest = performNodesInfoAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.getDiscoveryNodes()); });
+        });
+    }
+
+    private Map<String, List<MockNodesInfoRequest>> performNodesInfoAction(NodesInfoRequest request) {
+        TransportNodesAction action = getTestTransportNodesInfoAction();
+        PlainActionFuture<NodesStatsRequest> listener = new PlainActionFuture<>();
+        action.new AsyncAction(null, request, listener).start();
+        Map<String, List<CapturingTransport.CapturedRequest>> capturedRequests = transport.getCapturedRequestsByTargetNodeAndClear();
+        Map<String, List<MockNodesInfoRequest>> combinedSentRequest = new HashMap<>();
+
+        capturedRequests.forEach((node, capturedRequestList) -> {
+            List<MockNodesInfoRequest> sentRequestList = new ArrayList<>();
+
+            capturedRequestList.forEach(preSentRequest -> {
+                BytesStreamOutput out = new BytesStreamOutput();
+                try {
+                    TransportNodesInfoAction.NodeInfoRequest nodesInfoRequestFromCoordinator =
+                        (TransportNodesInfoAction.NodeInfoRequest) preSentRequest.request;
+                    nodesInfoRequestFromCoordinator.writeTo(out);
+                    StreamInput in = out.bytes().streamInput();
+                    MockNodesInfoRequest nodesStatsRequest = new MockNodesInfoRequest(in);
+                    sentRequestList.add(nodesStatsRequest);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+
+            combinedSentRequest.put(node, sentRequestList);
+        });
+
+        return combinedSentRequest;
+    }
+
+    private TestTransportNodesInfoAction getTestTransportNodesInfoAction() {
+        return new TestTransportNodesInfoAction(
+            THREAD_POOL,
+            clusterService,
+            transportService,
+            nodeService,
+            new ActionFilters(Collections.emptySet())
+        );
+    }
+
+    private static class TestTransportNodesInfoAction extends TransportNodesInfoAction {
+        public TestTransportNodesInfoAction(
+            ThreadPool threadPool,
+            ClusterService clusterService,
+            TransportService transportService,
+            NodeService nodeService,
+            ActionFilters actionFilters
+        ) {
+            super(threadPool, clusterService, transportService, nodeService, actionFilters);
+        }
+    }
+
+    private static class MockNodesInfoRequest extends TransportNodesInfoAction.NodeInfoRequest {
+
+        public MockNodesInfoRequest(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        public DiscoveryNode[] getDiscoveryNodes() {
+            return this.request.concreteNodes();
+        }
+    }
+}

server/src/test/java/org/opensearch/action/support/nodes/TransportNodesStatsActionTests.java

@@ -0,0 +1,130 @@
+/*
+ * 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.action.support.nodes;
+
+import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest;
+import org.opensearch.action.admin.cluster.node.stats.TransportNodesStatsAction;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.PlainActionFuture;
+import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.node.NodeService;
+import org.opensearch.test.transport.CapturingTransport;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportService;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class TransportNodesStatsActionTests extends TransportNodesActionTests {
+
+    /**
+     * By default, we send discovery nodes list to each request that is sent across from the coordinator node. This
+     * behavior is asserted in this test.
+     */
+    public void testNodesStatsActionWithRetentionOfDiscoveryNodesList() {
+        NodesStatsRequest request = new NodesStatsRequest();
+        request.sendDiscoveryNodes(true);
+        Map<String, List<MockNodeStatsRequest>> combinedSentRequest = performNodesStatsAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> {
+                assertNotNull(sentRequest.getDiscoveryNodes());
+                assertEquals(sentRequest.getDiscoveryNodes().length, clusterService.state().nodes().getSize());
+            });
+        });
+    }
+
+    /**
+     * By default, we send discovery nodes list to each request that is sent across from the coordinator node. This
+     * behavior is asserted in this test.
+     */
+    public void testNodesStatsActionWithoutRetentionOfDiscoveryNodesList() {
+        NodesStatsRequest request = new NodesStatsRequest();
+        request.sendDiscoveryNodes(false);
+        Map<String, List<MockNodeStatsRequest>> combinedSentRequest = performNodesStatsAction(request);
+
+        assertNotNull(combinedSentRequest);
+        combinedSentRequest.forEach((node, capturedRequestList) -> {
+            assertNotNull(capturedRequestList);
+            capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.getDiscoveryNodes()); });
+        });
+    }
+
+    private Map<String, List<MockNodeStatsRequest>> performNodesStatsAction(NodesStatsRequest request) {
+        TransportNodesAction action = getTestTransportNodesStatsAction();
+        PlainActionFuture<NodesStatsRequest> listener = new PlainActionFuture<>();
+        action.new AsyncAction(null, request, listener).start();
+        Map<String, List<CapturingTransport.CapturedRequest>> capturedRequests = transport.getCapturedRequestsByTargetNodeAndClear();
+        Map<String, List<MockNodeStatsRequest>> combinedSentRequest = new HashMap<>();
+
+        capturedRequests.forEach((node, capturedRequestList) -> {
+            List<MockNodeStatsRequest> sentRequestList = new ArrayList<>();
+
+            capturedRequestList.forEach(preSentRequest -> {
+                BytesStreamOutput out = new BytesStreamOutput();
+                try {
+                    TransportNodesStatsAction.NodeStatsRequest nodesStatsRequestFromCoordinator =
+                        (TransportNodesStatsAction.NodeStatsRequest) preSentRequest.request;
+                    nodesStatsRequestFromCoordinator.writeTo(out);
+                    StreamInput in = out.bytes().streamInput();
+                    MockNodeStatsRequest nodesStatsRequest = new MockNodeStatsRequest(in);
+                    sentRequestList.add(nodesStatsRequest);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+
+            combinedSentRequest.put(node, sentRequestList);
+        });
+
+        return combinedSentRequest;
+    }
+
+    private TestTransportNodesStatsAction getTestTransportNodesStatsAction() {
+        return new TestTransportNodesStatsAction(
+            THREAD_POOL,
+            clusterService,
+            transportService,
+            nodeService,
+            new ActionFilters(Collections.emptySet())
+        );
+    }
+
+    private static class TestTransportNodesStatsAction extends TransportNodesStatsAction {
+        public TestTransportNodesStatsAction(
+            ThreadPool threadPool,
+            ClusterService clusterService,
+            TransportService transportService,
+            NodeService nodeService,
+            ActionFilters actionFilters
+        ) {
+            super(threadPool, clusterService, transportService, nodeService, actionFilters);
+        }
+    }
+
+    private static class MockNodeStatsRequest extends TransportNodesStatsAction.NodeStatsRequest {
+
+        public MockNodeStatsRequest(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        public DiscoveryNode[] getDiscoveryNodes() {
+            return this.request.concreteNodes();
+        }
+    }
+}