Enable feature using flags

Pranshu-S · Jun 13, 2024 · 14e3db5 · 14e3db5
1 parent d6b78a4
commit 14e3db5
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 8 deletions.
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/stats/ClusterStatsRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/stats/ClusterStatsRequest.java
@@ -47,9 +47,28 @@
 @PublicApi(since = "1.0.0")
 public class ClusterStatsRequest extends BaseNodesRequest<ClusterStatsRequest> {
 
+    private boolean includeMappingStats;
+    private boolean includeAnalysisStats;
+
     public ClusterStatsRequest(StreamInput in) throws IOException {
         super(in);
     }
+    public void setIncludeMappingStats(boolean bool) {
+        this.includeMappingStats = bool;
+    }
+
+    public boolean isIncludeMappingStats() {
+        return includeMappingStats;
+    }
+
+    public void setIncludeAnalysisStats(boolean bool) {
+        this.includeAnalysisStats = bool;
+    }
+
+    public boolean isIncludeAnalysisStats() {
+        return includeAnalysisStats;
+    }
+
 
     /**
      * Get stats from nodes based on the nodes ids specified. If none are passed, stats

diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/stats/ClusterStatsResponse.java b/server/src/main/java/org/opensearch/action/admin/cluster/stats/ClusterStatsResponse.java
@@ -105,6 +105,33 @@ public ClusterStatsResponse(
         this.status = status;
     }
 
+    public ClusterStatsResponse(
+        long timestamp,
+        String clusterUUID,
+        ClusterName clusterName,
+        List<ClusterStatsNodeResponse> nodes,
+        List<FailedNodeException> failures,
+        ClusterState state,
+        ClusterStatsRequest request
+    ) {
+        super(clusterName, nodes, failures);
+        this.clusterUUID = clusterUUID;
+        this.timestamp = timestamp;
+        nodesStats = new ClusterStatsNodes(nodes);
+        indicesStats = new ClusterStatsIndices(nodes,
+            request.isIncludeMappingStats() ? MappingStats.of(state) : null,
+            request.isIncludeAnalysisStats() ? AnalysisStats.of(state) : null);
+        ClusterHealthStatus status = null;
+        for (ClusterStatsNodeResponse response : nodes) {
+            // only the cluster-manager node populates the status
+            if (response.clusterStatus() != null) {
+                status = response.clusterStatus();
+                break;
+            }
+        }
+        this.status = status;
+    }
+
     public String getClusterUUID() {
         return this.clusterUUID;
     }

diff --git a/.../src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java b/.../src/main/java/org/opensearch/action/admin/cluster/stats/TransportClusterStatsAction.java
@@ -46,6 +46,7 @@
 import org.opensearch.cluster.health.ClusterStateHealth;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.index.IndexService;
@@ -123,14 +124,28 @@ protected ClusterStatsResponse newResponse(
                 + " the cluster state that are too slow for a transport thread"
         );
         ClusterState state = clusterService.state();
-        return new ClusterStatsResponse(
-            System.currentTimeMillis(),
-            state.metadata().clusterUUID(),
-            clusterService.getClusterName(),
-            responses,
-            failures,
-            state
-        );
+
+        if (FeatureFlags.isEnabled(FeatureFlags.OPTIMIZED_CLUSTER_STATS_SETTING)) {
+            return new ClusterStatsResponse(
+                System.currentTimeMillis(),
+                state.metadata().clusterUUID(),
+                clusterService.getClusterName(),
+                responses,
+                failures,
+                state,
+                request
+            );
+        } else {
+            return new ClusterStatsResponse(
+                System.currentTimeMillis(),
+                state.metadata().clusterUUID(),
+                clusterService.getClusterName(),
+                responses,
+                failures,
+                state
+            );
+        }
+
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterStatsAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterStatsAction.java
@@ -66,6 +66,10 @@ 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.setIncludeMappingStats(request.paramAsBoolean("include_mapping_stats", false));
+        clusterStatsRequest.setIncludeAnalysisStats(request.paramAsBoolean("include_analysis_stats", false));
+
         return channel -> client.admin().cluster().clusterStats(clusterStatsRequest, new NodesResponseRestListener<>(channel));
     }