From 022e5ef84635b429335220e0c4765fe6225c99e0 Mon Sep 17 00:00:00 2001
From: Emmanuel Bosquet <bjokac@gmail.com>
Date: Mon, 20 Nov 2023 18:21:10 +0100
Subject: [PATCH] add no-clusters option on metrics query

---
 bin/src/cli.rs                 | 6 ++++++
 bin/src/ctl/command.rs         | 2 ++
 bin/src/ctl/mod.rs             | 3 ++-
 command/src/command.proto      | 2 ++
 lib/src/metrics/local_drain.rs | 9 +++++++++
 5 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/bin/src/cli.rs b/bin/src/cli.rs
index 28497a978..6850be8ec 100644
--- a/bin/src/cli.rs
+++ b/bin/src/cli.rs
@@ -220,6 +220,12 @@ pub enum MetricsCmd {
             // parse(try_from_str = split_slash)
         )]
         backends: Vec<String>,
+        #[clap(
+            short,
+            long,
+            help = "get only the metrics of main process and workers only (no cluster metrics)"
+        )]
+        no_clusters: bool,
     },
 }
 
diff --git a/bin/src/ctl/command.rs b/bin/src/ctl/command.rs
index 43822c4b4..518eea513 100644
--- a/bin/src/ctl/command.rs
+++ b/bin/src/ctl/command.rs
@@ -176,12 +176,14 @@ impl CommandManager {
         metric_names: Vec<String>,
         cluster_ids: Vec<String>,
         backend_ids: Vec<String>,
+        no_clusters: bool,
     ) -> Result<(), anyhow::Error> {
         let request: Request = RequestType::QueryMetrics(QueryMetricsOptions {
             list,
             cluster_ids,
             backend_ids,
             metric_names,
+            no_clusters,
         })
         .into();
 
diff --git a/bin/src/ctl/mod.rs b/bin/src/ctl/mod.rs
index e8b23f4f1..94875020a 100644
--- a/bin/src/ctl/mod.rs
+++ b/bin/src/ctl/mod.rs
@@ -78,7 +78,8 @@ impl CommandManager {
                     names,
                     clusters,
                     backends,
-                } => self.get_metrics(list, refresh, names, clusters, backends),
+                    no_clusters,
+                } => self.get_metrics(list, refresh, names, clusters, backends, no_clusters),
                 _ => self.configure_metrics(cmd),
             },
             SubCmd::Logging { level } => self.logging_filter(&level),
diff --git a/command/src/command.proto b/command/src/command.proto
index f69db483d..dcdc048a7 100644
--- a/command/src/command.proto
+++ b/command/src/command.proto
@@ -412,6 +412,8 @@ message QueryMetricsOptions {
     repeated string backend_ids = 3;
     // query only these metrics
     repeated string metric_names = 4;
+    // query only worker and main process metrics (no cluster metrics)
+    required bool no_clusters = 5;
 }
 
 // options to configure metrics collection
diff --git a/lib/src/metrics/local_drain.rs b/lib/src/metrics/local_drain.rs
index faf010915..ea29ef009 100644
--- a/lib/src/metrics/local_drain.rs
+++ b/lib/src/metrics/local_drain.rs
@@ -257,12 +257,21 @@ impl LocalDrain {
             cluster_ids,
             backend_ids,
             list,
+            no_clusters,
         } = options;
 
         if *list {
             return self.list_all_metric_names();
         }
 
+        if *no_clusters {
+            let proxy_metrics = self.dump_proxy_metrics(metric_names);
+            return Ok(ContentType::WorkerMetrics(WorkerMetrics {
+                proxy: proxy_metrics,
+                clusters: BTreeMap::new(),
+            }).into());
+        }
+
         let worker_metrics = match (cluster_ids.is_empty(), backend_ids.is_empty()) {
             (false, _) => self.query_clusters(cluster_ids, metric_names)?,
             (true, false) => self.query_backends(backend_ids, metric_names)?,