[Extensions] Only store names of extension transport actions (#6545) (#…

…6556) * Only store names of extension transport actions (cherry picked from commit 84cec8e) Signed-off-by: Daniel Widdis <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Mar 7, 2023 · a6cff75 · a6cff75
1 parent 39f25dc
commit a6cff75
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 48 deletions.
diff --git a/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java b/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java
@@ -8,73 +8,50 @@
 
 package org.opensearch.extensions;
 
-import org.opensearch.action.ActionRequest;
-import org.opensearch.action.ActionResponse;
-import org.opensearch.action.support.TransportAction;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
 import org.opensearch.transport.TransportRequest;
 
 import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.HashSet;
 import java.util.Objects;
-import java.util.Map.Entry;
+import java.util.Set;
 
 /**
  * Request to register extension Transport actions
  *
  * @opensearch.internal
  */
 public class RegisterTransportActionsRequest extends TransportRequest {
+    // The uniqueId defining the extension which runs this action
     private String uniqueId;
-    private Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> transportActions;
+    // The action names to register
+    private Set<String> transportActions;
 
-    public RegisterTransportActionsRequest(
-        String uniqueId,
-        Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> transportActions
-    ) {
+    public RegisterTransportActionsRequest(String uniqueId, Set<String> transportActions) {
         this.uniqueId = uniqueId;
-        this.transportActions = new HashMap<>(transportActions);
+        this.transportActions = transportActions;
     }
 
     public RegisterTransportActionsRequest(StreamInput in) throws IOException {
         super(in);
         this.uniqueId = in.readString();
-        Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> actions = new HashMap<>();
-        int actionCount = in.readVInt();
-        for (int i = 0; i < actionCount; i++) {
-            try {
-                String actionName = in.readString();
-                @SuppressWarnings("unchecked")
-                Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>> transportAction = (Class<
-                    ? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>) Class.forName(in.readString());
-                actions.put(actionName, transportAction);
-            } catch (ClassNotFoundException e) {
-                throw new IllegalArgumentException("Could not read transport action");
-            }
-        }
-        this.transportActions = actions;
+        this.transportActions = new HashSet<>(in.readStringList());
     }
 
     public String getUniqueId() {
         return uniqueId;
     }
 
-    public Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> getTransportActions() {
+    public Set<String> getTransportActions() {
         return transportActions;
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeString(uniqueId);
-        out.writeVInt(this.transportActions.size());
-        for (Entry<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> action : transportActions
-            .entrySet()) {
-            out.writeString(action.getKey());
-            out.writeString(action.getValue().getName());
-        }
+        out.writeStringCollection(transportActions);
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/extensions/action/ExtensionTransportActionsHandler.java b/server/src/main/java/org/opensearch/extensions/action/ExtensionTransportActionsHandler.java
@@ -93,7 +93,7 @@ public TransportResponse handleRegisterTransportActionsRequest(RegisterTransport
         logger.debug("Register Transport Actions request recieved {}", transportActionsRequest);
         DiscoveryExtensionNode extension = extensionIdMap.get(transportActionsRequest.getUniqueId());
         try {
-            for (String action : transportActionsRequest.getTransportActions().keySet()) {
+            for (String action : transportActionsRequest.getTransportActions()) {
                 registerAction(action, extension);
             }
         } catch (Exception e) {

diff --git a/server/src/test/java/org/opensearch/extensions/RegisterTransportActionsRequestTests.java b/server/src/test/java/org/opensearch/extensions/RegisterTransportActionsRequestTests.java
@@ -9,20 +9,19 @@
 package org.opensearch.extensions;
 
 import org.junit.Before;
-import org.opensearch.action.admin.indices.create.AutoCreateAction.TransportAction;
-import org.opensearch.common.collect.Map;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import java.util.Set;
 
 public class RegisterTransportActionsRequestTests extends OpenSearchTestCase {
     private RegisterTransportActionsRequest originalRequest;
 
     @Before
     public void setup() {
-        this.originalRequest = new RegisterTransportActionsRequest("extension-uniqueId", Map.of("testAction", TransportAction.class));
+        this.originalRequest = new RegisterTransportActionsRequest("extension-uniqueId", Set.of("testAction"));
     }
 
     public void testRegisterTransportActionsRequest() throws IOException {
@@ -31,16 +30,12 @@ public void testRegisterTransportActionsRequest() throws IOException {
         StreamInput input = output.bytes().streamInput();
         RegisterTransportActionsRequest parsedRequest = new RegisterTransportActionsRequest(input);
         assertEquals(parsedRequest.getTransportActions(), originalRequest.getTransportActions());
-        assertEquals(parsedRequest.getTransportActions().get("testAction"), originalRequest.getTransportActions().get("testAction"));
         assertEquals(parsedRequest.getTransportActions().size(), originalRequest.getTransportActions().size());
         assertEquals(parsedRequest.hashCode(), originalRequest.hashCode());
         assertTrue(originalRequest.equals(parsedRequest));
     }
 
     public void testToString() {
-        assertEquals(
-            originalRequest.toString(),
-            "TransportActionsRequest{uniqueId=extension-uniqueId, actions={testAction=class org.opensearch.action.admin.indices.create.AutoCreateAction$TransportAction}}"
-        );
+        assertEquals(originalRequest.toString(), "TransportActionsRequest{uniqueId=extension-uniqueId, actions=[testAction]}");
     }
 }
diff --git a/...src/test/java/org/opensearch/extensions/action/ExtensionTransportActionsHandlerTests.java b/...src/test/java/org/opensearch/extensions/action/ExtensionTransportActionsHandlerTests.java
@@ -11,7 +11,6 @@
 import org.junit.After;
 import org.junit.Before;
 import org.opensearch.Version;
-import org.opensearch.action.admin.indices.create.AutoCreateAction.TransportAction;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.common.io.stream.NamedWriteableRegistry;
@@ -39,6 +38,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 import static java.util.Collections.emptyMap;
@@ -118,7 +118,7 @@ public void testRegisterAction() {
 
     public void testRegisterTransportActionsRequest() {
         String action = "test-action";
-        RegisterTransportActionsRequest request = new RegisterTransportActionsRequest("uniqueid1", Map.of(action, TransportAction.class));
+        RegisterTransportActionsRequest request = new RegisterTransportActionsRequest("uniqueid1", Set.of(action));
         AcknowledgedResponse response = (AcknowledgedResponse) extensionTransportActionsHandler.handleRegisterTransportActionsRequest(
             request
         );
@@ -150,10 +150,7 @@ public void testSendTransportRequestToExtension() throws InterruptedException {
         );
 
         // Register Action
-        RegisterTransportActionsRequest registerRequest = new RegisterTransportActionsRequest(
-            "uniqueid1",
-            Map.of(action, TransportAction.class)
-        );
+        RegisterTransportActionsRequest registerRequest = new RegisterTransportActionsRequest("uniqueid1", Set.of(action));
         AcknowledgedResponse response = (AcknowledgedResponse) extensionTransportActionsHandler.handleRegisterTransportActionsRequest(
             registerRequest
         );