From e77eafc12e1f43086a5ba9bbdd7e8ad241330af5 Mon Sep 17 00:00:00 2001
From: Stas Malyshev <stas.malyshev@elastic.co>
Date: Wed, 11 Sep 2024 15:54:04 -0600
Subject: [PATCH] Add capabilities for stats depending on the flag for now

---
 .../cluster/stats/TransportClusterStatsAction.java    | 11 +++++++++--
 .../action/admin/cluster/RestClusterStatsAction.java  | 11 ++++++++++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java
index 267199761d52e..41c06fad5eb93 100644
--- a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java
+++ b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java
@@ -23,6 +23,7 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.CancellableSingleObjectCache;
+import org.elasticsearch.common.util.FeatureFlag;
 import org.elasticsearch.common.util.concurrent.ListenableFuture;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.indices.IndicesService;
@@ -52,6 +53,8 @@ public class TransportClusterStatsAction extends TransportClusterStatsBaseAction
 
     public static final ActionType<ClusterStatsResponse> TYPE = new ActionType<>("cluster:monitor/stats");
 
+    public static final FeatureFlag CCS_TELEMETRY_FEATURE_FLAG = new FeatureFlag("ccs_telemetry");
+
     private final MetadataStatsCache<MappingStats> mappingStatsCache;
     private final MetadataStatsCache<AnalysisStats> analysisStatsCache;
     private final RemoteClusterService remoteClusterService;
@@ -178,8 +181,12 @@ protected boolean isFresh(Long currentKey, Long newKey) {
         }
     }
 
+    private static boolean doRemotes(ClusterStatsRequest request) {
+        return CCS_TELEMETRY_FEATURE_FLAG.isEnabled() && request.doRemotes();
+    }
+
     private Map<String, ClusterStatsResponse.RemoteClusterStats> getRemoteClusterStats(ClusterStatsRequest request) {
-        if (request.doRemotes() == false) {
+        if (doRemotes(request) == false) {
             return null;
         }
         Map<String, ClusterStatsResponse.RemoteClusterStats> remoteClustersStats = new HashMap<>();
@@ -211,7 +218,7 @@ private Map<String, RemoteClusterStatsResponse> resolveRemoteClusterStats() {
     }
 
     private ActionFuture<Map<String, RemoteClusterStatsResponse>> getStatsFromRemotes(ClusterStatsRequest request) {
-        if (request.doRemotes() == false) {
+        if (doRemotes(request) == false) {
             // this will never be used since getRemoteClusterStats has the same check
             return null;
         }
diff --git a/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java b/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java
index 908d04995ff8d..1651e0abc3b09 100644
--- a/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java
+++ b/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java
@@ -10,6 +10,8 @@
 
 import org.elasticsearch.action.admin.cluster.stats.ClusterStatsRequest;
 import org.elasticsearch.client.internal.node.NodeClient;
+import org.elasticsearch.common.util.FeatureFlag;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.Scope;
@@ -28,6 +30,8 @@
 public class RestClusterStatsAction extends BaseRestHandler {
 
     private static final Set<String> SUPPORTED_CAPABILITIES = Set.of("human-readable-total-docs-size");
+    private static final Set<String> SUPPORTED_CAPABILITIES_CCS_STATS = Sets.union(SUPPORTED_CAPABILITIES, Set.of("ccs-stats"));
+    public static final FeatureFlag CCS_TELEMETRY_FEATURE_FLAG = new FeatureFlag("ccs_telemetry");
 
     @Override
     public List<Route> routes() {
@@ -39,6 +43,11 @@ public String getName() {
         return "cluster_stats_action";
     }
 
+    @Override
+    public Set<String> supportedQueryParameters() {
+        return Set.of("include_remotes", "nodeId");
+    }
+
     @Override
     public RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException {
         ClusterStatsRequest clusterStatsRequest = new ClusterStatsRequest(
@@ -58,6 +67,6 @@ public boolean canTripCircuitBreaker() {
 
     @Override
     public Set<String> supportedCapabilities() {
-        return SUPPORTED_CAPABILITIES;
+        return CCS_TELEMETRY_FEATURE_FLAG.isEnabled() ? SUPPORTED_CAPABILITIES_CCS_STATS : SUPPORTED_CAPABILITIES;
     }
 }