From fc7638bc4d0463cd33a6ddf703671d8b11805c24 Mon Sep 17 00:00:00 2001
From: trinity-1686a <trinity@quickwit.io>
Date: Mon, 3 Jun 2024 09:53:48 +0200
Subject: [PATCH] add metric for search split affinity (#4998)

* add metric for search split affinity
---
 quickwit/quickwit-search/src/metrics.rs           | 11 ++++++++++-
 quickwit/quickwit-search/src/search_job_placer.rs | 12 +++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/quickwit/quickwit-search/src/metrics.rs b/quickwit/quickwit-search/src/metrics.rs
index 718d7b307b..449b747f16 100644
--- a/quickwit/quickwit-search/src/metrics.rs
+++ b/quickwit/quickwit-search/src/metrics.rs
@@ -21,12 +21,14 @@
 
 use once_cell::sync::Lazy;
 use quickwit_common::metrics::{
-    exponential_buckets, new_counter, new_histogram, Histogram, IntCounter,
+    exponential_buckets, new_counter, new_counter_vec, new_histogram, Histogram, IntCounter,
+    IntCounterVec,
 };
 
 pub struct SearchMetrics {
     pub leaf_searches_splits_total: IntCounter,
     pub leaf_search_split_duration_secs: Histogram,
+    pub job_assigned_total: IntCounterVec<1>,
 }
 
 impl Default for SearchMetrics {
@@ -45,6 +47,13 @@ impl Default for SearchMetrics {
                 "search",
                 exponential_buckets(0.005, 2.0, 10).unwrap(),
             ),
+            job_assigned_total: new_counter_vec(
+                "job_assigned_total",
+                "Number of job assigned to searchers, per affinity rank.",
+                "search",
+                &[],
+                ["affinity"],
+            ),
         }
     }
 }
diff --git a/quickwit/quickwit-search/src/search_job_placer.rs b/quickwit/quickwit-search/src/search_job_placer.rs
index 24c6677d0e..dc708763ff 100644
--- a/quickwit/quickwit-search/src/search_job_placer.rs
+++ b/quickwit/quickwit-search/src/search_job_placer.rs
@@ -29,7 +29,7 @@ use quickwit_common::pubsub::EventSubscriber;
 use quickwit_common::rendezvous_hasher::{node_affinity, sort_by_rendez_vous_hash};
 use quickwit_proto::search::{ReportSplit, ReportSplitsRequest};
 
-use crate::{SearchJob, SearchServiceClient, SearcherPool};
+use crate::{SearchJob, SearchServiceClient, SearcherPool, SEARCH_METRICS};
 
 /// Job.
 /// The unit in which distributed search is performed.
@@ -185,6 +185,16 @@ impl SearchJobPlacer {
             } else {
                 0
             };
+            let metric_node_idx = match chosen_node_idx {
+                0 => "0",
+                1 => "1",
+                _ => "> 1",
+            };
+            SEARCH_METRICS
+                .job_assigned_total
+                .with_label_values([metric_node_idx])
+                .inc();
+
             let chosen_node = &mut candidate_nodes[chosen_node_idx];
             chosen_node.load += job.cost();