From 0f7bd508cafed1dbd746364778ef3ed105c167ae Mon Sep 17 00:00:00 2001
From: Remi Dettai <rdettai@gmail.com>
Date: Tue, 3 Dec 2024 14:14:00 +0100
Subject: [PATCH] Add warmup cache metric

---
 quickwit/quickwit-search/src/metrics.rs       | 20 +++++++++++++++++++
 .../src/search_permit_provider.rs             |  3 +++
 2 files changed, 23 insertions(+)

diff --git a/quickwit/quickwit-search/src/metrics.rs b/quickwit/quickwit-search/src/metrics.rs
index 35b7d3115c5..0a628240af8 100644
--- a/quickwit/quickwit-search/src/metrics.rs
+++ b/quickwit/quickwit-search/src/metrics.rs
@@ -19,6 +19,7 @@
 
 // See https://prometheus.io/docs/practices/naming/
 
+use bytesize::ByteSize;
 use once_cell::sync::Lazy;
 use quickwit_common::metrics::{
     exponential_buckets, linear_buckets, new_counter, new_counter_vec, new_gauge_vec,
@@ -37,6 +38,7 @@ pub struct SearchMetrics {
     pub job_assigned_total: IntCounterVec<1>,
     pub leaf_search_single_split_tasks_pending: IntGauge,
     pub leaf_search_single_split_tasks_ongoing: IntGauge,
+    pub leaf_search_single_split_warmup_num_bytes: Histogram,
 }
 
 impl Default for SearchMetrics {
@@ -52,6 +54,18 @@ impl Default for SearchMetrics {
         .copied()
         .collect();
 
+        let pseudo_exponential_bytes_buckets = vec![
+            ByteSize::mb(10).as_u64() as f64,
+            ByteSize::mb(20).as_u64() as f64,
+            ByteSize::mb(50).as_u64() as f64,
+            ByteSize::mb(100).as_u64() as f64,
+            ByteSize::mb(200).as_u64() as f64,
+            ByteSize::mb(500).as_u64() as f64,
+            ByteSize::gb(1).as_u64() as f64,
+            ByteSize::mb(2).as_u64() as f64,
+            ByteSize::mb(5).as_u64() as f64,
+        ];
+
         let leaf_search_single_split_tasks = new_gauge_vec::<1>(
             "leaf_search_single_split_tasks",
             "Number of single split search tasks pending or ongoing",
@@ -124,6 +138,12 @@ impl Default for SearchMetrics {
                 .with_label_values(["ongoing"]),
             leaf_search_single_split_tasks_pending: leaf_search_single_split_tasks
                 .with_label_values(["pending"]),
+            leaf_search_single_split_warmup_num_bytes: new_histogram(
+                "leaf_search_single_split_warmup_num_bytes",
+                "Size of the short lived cache for a single split once the warmup is done.",
+                "search",
+                pseudo_exponential_bytes_buckets,
+            ),
             job_assigned_total: new_counter_vec(
                 "job_assigned_total",
                 "Number of job assigned to searchers, per affinity rank.",
diff --git a/quickwit/quickwit-search/src/search_permit_provider.rs b/quickwit/quickwit-search/src/search_permit_provider.rs
index b3a61f0f68f..bc22c1a8336 100644
--- a/quickwit/quickwit-search/src/search_permit_provider.rs
+++ b/quickwit/quickwit-search/src/search_permit_provider.rs
@@ -219,6 +219,9 @@ impl SearchPermit {
     /// the warmup slot.
     pub fn warmup_completed(&mut self, new_memory_usage: ByteSize) {
         let new_usage_bytes = new_memory_usage.as_u64();
+        crate::SEARCH_METRICS
+            .leaf_search_single_split_warmup_num_bytes
+            .observe(new_usage_bytes as f64);
         if new_usage_bytes > self.memory_allocation {
             warn!(
                 memory_usage = new_usage_bytes,