From 75a0e81dfd1fc6aa743bc32e372b2a84bea25342 Mon Sep 17 00:00:00 2001
From: Joshua Palis <jpalis@amazon.com>
Date: Fri, 9 Jun 2023 10:24:02 -0700
Subject: [PATCH] [Extensions] Improves performance by replacing cluster state
 calls to retrieve cluster name (#924)

* replacing cluster state calls to retrieve cluster name

Signed-off-by: Joshua Palis <jpalis@amazon.com>

* moving all getEnvironmentSettings calls out of constructor since settings are populated after initialization

Signed-off-by: Joshua Palis <jpalis@amazon.com>

* Removing unnecessary extensionsRunners, retrieving cluster name from SDKClusterService directly

Signed-off-by: Joshua Palis <jpalis@amazon.com>

---------

Signed-off-by: Joshua Palis <jpalis@amazon.com>
---
 .../ad/transport/ADCancelTaskTransportAction.java        | 2 +-
 .../ad/transport/ADStatsNodesTransportAction.java        | 2 +-
 .../ad/transport/ADTaskProfileTransportAction.java       | 2 +-
 .../ad/transport/DeleteModelTransportAction.java         | 2 +-
 .../opensearch/ad/transport/ProfileTransportAction.java  | 9 +++++----
 .../ad/transport/DeleteModelTransportActionTests.java    | 2 ++
 6 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/main/java/org/opensearch/ad/transport/ADCancelTaskTransportAction.java b/src/main/java/org/opensearch/ad/transport/ADCancelTaskTransportAction.java
index d228ddeb3..c8c4cd6c4 100644
--- a/src/main/java/org/opensearch/ad/transport/ADCancelTaskTransportAction.java
+++ b/src/main/java/org/opensearch/ad/transport/ADCancelTaskTransportAction.java
@@ -53,7 +53,7 @@ protected ADCancelTaskResponse newResponse(
         List<ADCancelTaskNodeResponse> responses,
         List<FailedNodeException> failures
     ) {
-        return new ADCancelTaskResponse(clusterService.state().getClusterName(), responses, failures);
+        return new ADCancelTaskResponse(clusterService.getClusterName(), responses, failures);
     }
 
     @Override
diff --git a/src/main/java/org/opensearch/ad/transport/ADStatsNodesTransportAction.java b/src/main/java/org/opensearch/ad/transport/ADStatsNodesTransportAction.java
index f37522e23..d5eaaa373 100644
--- a/src/main/java/org/opensearch/ad/transport/ADStatsNodesTransportAction.java
+++ b/src/main/java/org/opensearch/ad/transport/ADStatsNodesTransportAction.java
@@ -77,7 +77,7 @@ protected ADStatsNodesResponse newResponse(
         List<ADStatsNodeResponse> responses,
         List<FailedNodeException> failures
     ) {
-        return new ADStatsNodesResponse(sdkClusterService.state().getClusterName(), responses, failures);
+        return new ADStatsNodesResponse(sdkClusterService.getClusterName(), responses, failures);
     }
 
     protected ADStatsNodeRequest newNodeRequest(ADStatsRequest request) {
diff --git a/src/main/java/org/opensearch/ad/transport/ADTaskProfileTransportAction.java b/src/main/java/org/opensearch/ad/transport/ADTaskProfileTransportAction.java
index c828dd760..5b54b4164 100644
--- a/src/main/java/org/opensearch/ad/transport/ADTaskProfileTransportAction.java
+++ b/src/main/java/org/opensearch/ad/transport/ADTaskProfileTransportAction.java
@@ -63,7 +63,7 @@ protected ADTaskProfileResponse newResponse(
         List<ADTaskProfileNodeResponse> responses,
         List<FailedNodeException> failures
     ) {
-        return new ADTaskProfileResponse(sdkClusterService.state().getClusterName(), responses, failures);
+        return new ADTaskProfileResponse(sdkClusterService.getClusterName(), responses, failures);
     }
 
     protected ADTaskProfileNodeRequest newNodeRequest(ADTaskProfileRequest request) {
diff --git a/src/main/java/org/opensearch/ad/transport/DeleteModelTransportAction.java b/src/main/java/org/opensearch/ad/transport/DeleteModelTransportAction.java
index e8022b67b..0aa5a6fef 100644
--- a/src/main/java/org/opensearch/ad/transport/DeleteModelTransportAction.java
+++ b/src/main/java/org/opensearch/ad/transport/DeleteModelTransportAction.java
@@ -106,6 +106,6 @@ protected DeleteModelResponse newResponse(
         List<DeleteModelNodeResponse> responses,
         List<FailedNodeException> failures
     ) {
-        return new DeleteModelResponse(clusterService.state().getClusterName(), responses, failures);
+        return new DeleteModelResponse(clusterService.getClusterName(), responses, failures);
     }
 }
diff --git a/src/main/java/org/opensearch/ad/transport/ProfileTransportAction.java b/src/main/java/org/opensearch/ad/transport/ProfileTransportAction.java
index 68d5475da..47deed33f 100644
--- a/src/main/java/org/opensearch/ad/transport/ProfileTransportAction.java
+++ b/src/main/java/org/opensearch/ad/transport/ProfileTransportAction.java
@@ -29,7 +29,6 @@
 import org.opensearch.ad.ml.ModelManager;
 import org.opensearch.ad.model.DetectorProfileName;
 import org.opensearch.ad.model.ModelProfile;
-import org.opensearch.common.settings.Settings;
 import org.opensearch.sdk.ExtensionsRunner;
 import org.opensearch.sdk.SDKClusterService;
 import org.opensearch.tasks.Task;
@@ -52,8 +51,11 @@ public class ProfileTransportAction extends TransportAction<ProfileRequest, Prof
     /**
      * Constructor
      *
+     * @param extensionsRunner Extensions Runner
      * @param actionFilters Action Filters
+     * @param taskManager Task Manager
      * @param modelManager model manager object
+     * @param sdkClusterService extension cluster service
      * @param featureManager feature manager object
      * @param cacheProvider cache provider
      */
@@ -72,13 +74,12 @@ public ProfileTransportAction(
         this.featureManager = featureManager;
         this.cacheProvider = cacheProvider;
         this.sdkClusterService = sdkClusterService;
-        Settings settings = extensionsRunner.getEnvironmentSettings();
-        this.numModelsToReturn = MAX_MODEL_SIZE_PER_NODE.get(settings);
+        this.numModelsToReturn = MAX_MODEL_SIZE_PER_NODE.get(extensionsRunner.getEnvironmentSettings());
         this.sdkClusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_MODEL_SIZE_PER_NODE, it -> this.numModelsToReturn = it);
     }
 
     private ProfileResponse newResponse(ProfileRequest request, List<ProfileNodeResponse> responses, List<FailedNodeException> failures) {
-        return new ProfileResponse(sdkClusterService.state().getClusterName(), responses, failures);
+        return new ProfileResponse(sdkClusterService.getClusterName(), responses, failures);
     }
 
     @Override
diff --git a/src/test/java/org/opensearch/ad/transport/DeleteModelTransportActionTests.java b/src/test/java/org/opensearch/ad/transport/DeleteModelTransportActionTests.java
index 6a1438436..494ead078 100644
--- a/src/test/java/org/opensearch/ad/transport/DeleteModelTransportActionTests.java
+++ b/src/test/java/org/opensearch/ad/transport/DeleteModelTransportActionTests.java
@@ -68,7 +68,9 @@ public void setUp() throws Exception {
         ThreadPool threadPool = mock(ThreadPool.class);
 
         SDKClusterService clusterService = mock(SDKClusterService.class);
+        ClusterName clusterName = mock(ClusterName.class);
         ClusterState clusterState = mock(ClusterState.class);
+        when(clusterService.getClusterName()).thenReturn(clusterName);
         when(clusterService.state()).thenReturn(clusterState);
         when(clusterState.getClusterName()).thenReturn(new ClusterName("clustername"));
         localNodeID = "foo";