[Extensions] Add ClusterStateRequest parameter to cluster state transport request

server/src/main/java/org/opensearch/extensions/ExtensionsManager.java

@@ -32,7 +32,7 @@
 import org.opensearch.OpenSearchException;
 import org.opensearch.Version;
 import org.opensearch.action.ActionModule;
-import org.opensearch.action.admin.cluster.state.ClusterStateResponse;
+import org.opensearch.action.admin.cluster.state.ClusterStateRequest;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.ClusterSettingsResponse;
 import org.opensearch.cluster.node.DiscoveryNode;
@@ -246,8 +246,8 @@ private void registerRequestHandler() {
             ThreadPool.Names.GENERIC,
             false,
             false,
-            ExtensionRequest::new,
-            ((request, channel, task) -> channel.sendResponse(handleExtensionRequest(request)))
+            ClusterStateRequest::new,
+            ((request, channel, task) -> channel.sendResponse(extensionTransportActionsHandler.handleClusterStateRequest(request)))
         );
         transportService.registerRequestHandler(
             REQUEST_EXTENSION_CLUSTER_SETTINGS,
@@ -436,8 +436,6 @@ public String executor() {
      */
     TransportResponse handleExtensionRequest(ExtensionRequest extensionRequest) throws Exception {
         switch (extensionRequest.getRequestType()) {
-            case REQUEST_EXTENSION_CLUSTER_STATE:
-                return new ClusterStateResponse(clusterService.getClusterName(), clusterService.state(), false);
             case REQUEST_EXTENSION_CLUSTER_SETTINGS:
                 return new ClusterSettingsResponse(clusterService);
             case REQUEST_EXTENSION_ENVIRONMENT_SETTINGS:

server/src/main/java/org/opensearch/extensions/action/ExtensionTransportActionsHandler.java

@@ -13,6 +13,9 @@
 import org.opensearch.action.ActionListener;
 import org.opensearch.action.ActionModule;
 import org.opensearch.action.ActionModule.DynamicActionRegistry;
+import org.opensearch.action.admin.cluster.state.ClusterStateAction;
+import org.opensearch.action.admin.cluster.state.ClusterStateRequest;
+import org.opensearch.action.admin.cluster.state.ClusterStateResponse;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.common.io.stream.StreamInput;
@@ -31,6 +34,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
@@ -310,4 +314,42 @@ public String executor() {
         }
         return extensionActionResponse;
     }
+
+    /**
+     * Handles a {@link ClusterStateRequest}.
+     *
+     * @param request The request to handle.
+     * @return A {@link ClusterStateResponse}.
+     */
+    public ClusterStateResponse handleClusterStateRequest(ClusterStateRequest request) {
+        final CompletableFuture<ClusterStateResponse> inProgressFuture = new CompletableFuture<>();
+        client.execute(ClusterStateAction.INSTANCE, request, new ActionListener<ClusterStateResponse>() {
+            @Override
+            public void onResponse(ClusterStateResponse response) {
+                inProgressFuture.complete(response);
+            }
+
+            @Override
+            public void onFailure(Exception exp) {
+                logger.debug("Cluster State Request failed", exp);
+                inProgressFuture.completeExceptionally(exp);
+            }
+        });
+        try {
+            return inProgressFuture.orTimeout(ExtensionsManager.EXTENSION_REQUEST_WAIT_TIMEOUT, TimeUnit.SECONDS).get();
+        } catch (ExecutionException e) {
+            if (e.getCause() instanceof TimeoutException) {
+                logger.debug("No response from extension to request.");
+            }
+            if (e.getCause() instanceof RuntimeException) {
+                throw (RuntimeException) e.getCause();
+            } else if (e.getCause() instanceof Error) {
+                throw (Error) e.getCause();
+            } else {
+                throw new RuntimeException(e.getCause());
+            }
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

server/src/test/java/org/opensearch/extensions/ExtensionsManagerTests.java

@@ -42,7 +42,6 @@
 import org.junit.Before;
 import org.opensearch.Version;
 import org.opensearch.action.ActionModule;
-import org.opensearch.action.admin.cluster.state.ClusterStateResponse;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.ClusterSettingsResponse;
 import org.opensearch.env.EnvironmentSettingsResponse;
@@ -524,9 +523,6 @@ public void testHandleExtensionRequest() throws Exception {
         ExtensionsManager extensionsManager = new ExtensionsManager(settings, extensionDir);
         initialize(extensionsManager);
 
-        ExtensionRequest clusterStateRequest = new ExtensionRequest(ExtensionsManager.RequestType.REQUEST_EXTENSION_CLUSTER_STATE);
-        assertEquals(ClusterStateResponse.class, extensionsManager.handleExtensionRequest(clusterStateRequest).getClass());
-
         ExtensionRequest clusterSettingRequest = new ExtensionRequest(ExtensionsManager.RequestType.REQUEST_EXTENSION_CLUSTER_SETTINGS);
         assertEquals(ClusterSettingsResponse.class, extensionsManager.handleExtensionRequest(clusterSettingRequest).getClass());
 

server/src/test/java/org/opensearch/extensions/action/ExtensionTransportActionsHandlerTests.java

@@ -13,6 +13,8 @@
 import org.opensearch.Version;
 import org.opensearch.action.ActionModule;
 import org.opensearch.action.ActionModule.DynamicActionRegistry;
+import org.opensearch.action.admin.cluster.state.ClusterStateAction;
+import org.opensearch.action.admin.cluster.state.ClusterStateRequest;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.node.DiscoveryNode;
@@ -43,7 +45,12 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import static java.util.Collections.emptyMap;
@@ -95,7 +102,7 @@ public void setup() throws Exception {
             Version.fromString("3.0.0"),
             Collections.emptyList()
         );
-        client = new NoOpNodeClient(this.getTestName());
+        client = spy(new NoOpNodeClient(this.getTestName()));
         ActionModule mockActionModule = mock(ActionModule.class);
         DynamicActionRegistry dynamicActionRegistry = new DynamicActionRegistry();
         dynamicActionRegistry.registerUnmodifiableActionMap(Collections.emptyMap());
@@ -176,4 +183,11 @@ public void testSendTransportRequestToExtension() throws InterruptedException {
 
         expectThrows(NodeNotConnectedException.class, () -> extensionTransportActionsHandler.sendTransportRequestToExtension(request));
     }
+
+    public void testHandleClusterStateRequest() throws Exception {
+        ClusterStateRequest clusterStateRequest = new ClusterStateRequest().all();
+        // The client is no-op so doesn't actually execute this request, but we can verify it's called
+        extensionTransportActionsHandler.handleClusterStateRequest(clusterStateRequest);
+        verify(client, times(1)).execute(eq(ClusterStateAction.INSTANCE), eq(clusterStateRequest), any());
+    }
 }