Fixing GaugeGuard's API

fulmicoton-dd · fulmicoton-dd · commit ec1add6e8401 · 2025-09-12T12:04:27.000+02:00
diff --git a/quickwit/quickwit-actors/src/mailbox.rs b/quickwit/quickwit-actors/src/mailbox.rs
@@ -395,9 +395,7 @@ fn get_actor_inboxes_count_gauge_guard() -> GaugeGuard<'static> {
             &[],
         )
     });
-    let mut gauge_guard = GaugeGuard::from_gauge(gauge);
-    gauge_guard.add(1);
-    gauge_guard
+    GaugeGuard::from_gauge_with_initial_value(gauge, 1)
 }
 
 pub(crate) fn create_mailbox<A: Actor>(
diff --git a/quickwit/quickwit-common/src/metrics.rs b/quickwit/quickwit-common/src/metrics.rs
@@ -219,8 +219,13 @@ impl std::fmt::Debug for GaugeGuard<'_> {
 }
 
 impl<'a> GaugeGuard<'a> {
-    pub fn from_gauge(gauge: &'a IntGauge) -> Self {
-        Self { gauge, delta: 0i64 }
+    pub fn from_gauge_with_initial_value(gauge: &'a IntGauge, initial_value: i64) -> Self {
+        let mut gauge = Self {
+            gauge,
+            delta: initial_value,
+        };
+        gauge.add(initial_value);
+        gauge
     }
 
     pub fn get(&self) -> i64 {
@@ -256,8 +261,10 @@ impl std::fmt::Debug for OwnedGaugeGuard {
 }
 
 impl OwnedGaugeGuard {
-    pub fn from_gauge(gauge: IntGauge) -> Self {
-        Self { gauge, delta: 0i64 }
+    pub fn from_gauge_with_initial_value(gauge: IntGauge, initial_value: i64) -> Self {
+        let mut gauge = Self { gauge, delta: 0i64 };
+        gauge.add(initial_value);
+        gauge
     }
 
     pub fn get(&self) -> i64 {
diff --git a/quickwit/quickwit-common/src/stream_utils.rs b/quickwit/quickwit-common/src/stream_utils.rs
@@ -240,9 +240,8 @@ where T: fmt::Debug
 
 impl<T> InFlightValue<T> {
     pub fn new(value: T, value_size: ByteSize, gauge: &'static IntGauge) -> Self {
-        let mut gauge_guard = GaugeGuard::from_gauge(gauge);
-        gauge_guard.add(value_size.as_u64() as i64);
-
+        let gauge_guard =
+            GaugeGuard::from_gauge_with_initial_value(gauge, value_size.as_u64() as i64);
         Self(value, gauge_guard)
     }
 
diff --git a/quickwit/quickwit-common/src/thread_pool.rs b/quickwit/quickwit-common/src/thread_pool.rs
@@ -85,18 +85,17 @@ impl ThreadPool {
     {
         let span = tracing::Span::current();
         let ongoing_tasks = self.ongoing_tasks.clone();
-        let mut pending_tasks_guard: OwnedGaugeGuard =
-            OwnedGaugeGuard::from_gauge(self.pending_tasks.clone());
-        pending_tasks_guard.add(1i64);
+        let pending_tasks_guard: OwnedGaugeGuard =
+            OwnedGaugeGuard::from_gauge_with_initial_value(self.pending_tasks.clone(), 1i64);
         let (tx, rx) = oneshot::channel();
         self.thread_pool.spawn(move || {
             drop(pending_tasks_guard);
             if tx.is_closed() {
                 return;
             }
             let _guard = span.enter();
-            let mut ongoing_task_guard = GaugeGuard::from_gauge(&ongoing_tasks);
-            ongoing_task_guard.add(1i64);
+            let _ongoing_task_guard =
+                GaugeGuard::from_gauge_with_initial_value(&ongoing_tasks, 1i64);
             let result = cpu_intensive_fn();
             let _ = tx.send(result);
         });
diff --git a/quickwit/quickwit-indexing/src/actors/indexer.rs b/quickwit/quickwit-indexing/src/actors/indexer.rs
@@ -219,9 +219,10 @@ impl IndexerState {
         let publish_lock = self.publish_lock.clone();
         let publish_token_opt = self.publish_token_opt.clone();
 
-        let mut split_builders_guard =
-            GaugeGuard::from_gauge(&crate::metrics::INDEXER_METRICS.split_builders);
-        split_builders_guard.add(1);
+        let split_builders_guard = GaugeGuard::from_gauge_with_initial_value(
+            &crate::metrics::INDEXER_METRICS.split_builders,
+            1,
+        );
 
         let workbench = IndexingWorkbench {
             workbench_id,
@@ -233,10 +234,11 @@ impl IndexerState {
             publish_lock,
             publish_token_opt,
             last_delete_opstamp,
-            memory_usage: GaugeGuard::from_gauge(
+            memory_usage: GaugeGuard::from_gauge_with_initial_value(
                 &quickwit_common::metrics::MEMORY_METRICS
                     .in_flight
                     .index_writer,
+                0i64,
             ),
             cooperative_indexing_period,
             split_builders_guard,
diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs
@@ -159,9 +159,8 @@ impl IndexingPipeline {
         let indexing_pipelines_gauge = crate::metrics::INDEXER_METRICS
             .indexing_pipelines
             .with_label_values([&params.pipeline_id.index_uid.index_id]);
-        let mut indexing_pipelines_gauge_guard =
-            OwnedGaugeGuard::from_gauge(indexing_pipelines_gauge);
-        indexing_pipelines_gauge_guard.add(1);
+        let indexing_pipelines_gauge_guard =
+            OwnedGaugeGuard::from_gauge_with_initial_value(indexing_pipelines_gauge, 1);
         let params_fingerprint = params.params_fingerprint;
         IndexingPipeline {
             params,
diff --git a/quickwit/quickwit-indexing/src/models/processed_doc.rs b/quickwit/quickwit-indexing/src/models/processed_doc.rs
@@ -51,8 +51,10 @@ impl ProcessedDocBatch {
         force_commit: bool,
     ) -> Self {
         let delta = docs.iter().map(|doc| doc.num_bytes as i64).sum::<i64>();
-        let mut gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.indexer_mailbox);
-        gauge_guard.add(delta);
+        let gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.indexer_mailbox,
+            delta,
+        );
         Self {
             docs,
             checkpoint_delta,
diff --git a/quickwit/quickwit-indexing/src/models/raw_doc_batch.rs b/quickwit/quickwit-indexing/src/models/raw_doc_batch.rs
@@ -34,9 +34,10 @@ impl RawDocBatch {
         force_commit: bool,
     ) -> Self {
         let delta = docs.iter().map(|doc| doc.len() as i64).sum::<i64>();
-        let mut gauge_guard =
-            GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.doc_processor_mailbox);
-        gauge_guard.add(delta);
+        let gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.doc_processor_mailbox,
+            delta,
+        );
 
         Self {
             docs,
@@ -67,7 +68,10 @@ impl fmt::Debug for RawDocBatch {
 
 impl Default for RawDocBatch {
     fn default() -> Self {
-        let _gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.doc_processor_mailbox);
+        let _gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.doc_processor_mailbox,
+            0i64,
+        );
         Self {
             docs: Vec::new(),
             checkpoint_delta: SourceCheckpointDelta::default(),
diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs
@@ -532,7 +532,7 @@ impl BatchBuilder {
             SourceType::Pulsar => MEMORY_METRICS.in_flight.pulsar(),
             _ => MEMORY_METRICS.in_flight.other(),
         };
-        let gauge_guard = GaugeGuard::from_gauge(gauge);
+        let gauge_guard = GaugeGuard::from_gauge_with_initial_value(gauge, 0i64);
 
         Self {
             docs: Vec::with_capacity(capacity),
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/ingester.rs b/quickwit/quickwit-ingest/src/ingest_v2/ingester.rs
@@ -1115,8 +1115,10 @@ impl IngesterService for Ingester {
                 _ => None,
             })
             .sum::<usize>();
-        let mut gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.ingester_persist);
-        gauge_guard.add(request_size_bytes as i64);
+        let _gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.ingester_persist,
+            request_size_bytes as i64,
+        );
 
         self.persist_inner(persist_request).await
     }
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/replication.rs b/quickwit/quickwit-ingest/src/ingest_v2/replication.rs
@@ -504,8 +504,10 @@ impl ReplicationTask {
             )));
         }
         let request_size_bytes = replicate_request.num_bytes();
-        let mut gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.ingester_replicate);
-        gauge_guard.add(request_size_bytes as i64);
+        let _gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.ingester_replicate,
+            request_size_bytes as i64,
+        );
 
         self.current_replication_seqno += 1;
 
diff --git a/quickwit/quickwit-ingest/src/ingest_v2/router.rs b/quickwit/quickwit-ingest/src/ingest_v2/router.rs
@@ -602,8 +602,10 @@ impl IngestRouterService for IngestRouter {
     async fn ingest(&self, ingest_request: IngestRequestV2) -> IngestV2Result<IngestResponseV2> {
         let request_size_bytes = ingest_request.num_bytes();
 
-        let mut gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.ingest_router);
-        gauge_guard.add(request_size_bytes as i64);
+        let _gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.ingest_router,
+            request_size_bytes as i64,
+        );
         let num_subrequests = ingest_request.subrequests.len();
 
         let _permit = self
diff --git a/quickwit/quickwit-search/src/scroll_context.rs b/quickwit/quickwit-search/src/scroll_context.rs
@@ -148,9 +148,10 @@ impl Default for MiniKV {
 
 impl MiniKV {
     pub async fn put(&self, key: Vec<u8>, payload: Vec<u8>, ttl: Duration) {
-        let mut metric_guard =
-            GaugeGuard::from_gauge(&crate::SEARCH_METRICS.searcher_local_kv_store_size_bytes);
-        metric_guard.add(payload.len() as i64);
+        let metric_guard = GaugeGuard::from_gauge_with_initial_value(
+            &crate::SEARCH_METRICS.searcher_local_kv_store_size_bytes,
+            payload.len() as i64,
+        );
         let mut cache_lock = self.ttl_with_cache.write().await;
         cache_lock.insert(
             key,
diff --git a/quickwit/quickwit-search/src/search_permit_provider.rs b/quickwit/quickwit-search/src/search_permit_provider.rs
@@ -212,10 +212,10 @@ impl SearchPermitActor {
         while let Some((permit_requester_tx, next_permit_size)) =
             self.pop_next_request_if_serviceable()
         {
-            let mut ongoing_gauge_guard = GaugeGuard::from_gauge(
+            let ongoing_gauge_guard = GaugeGuard::from_gauge_with_initial_value(
                 &crate::SEARCH_METRICS.leaf_search_single_split_tasks_ongoing,
+                1,
             );
-            ongoing_gauge_guard.add(1);
             self.total_memory_allocated += next_permit_size;
             self.num_warmup_slots_available -= 1;
             permit_requester_tx
diff --git a/quickwit/quickwit-serve/src/decompression.rs b/quickwit/quickwit-serve/src/decompression.rs
@@ -114,8 +114,10 @@ pub(crate) struct Body {
 
 impl Body {
     pub fn new(content: Bytes, load_shield_permit: LoadShieldPermit) -> Body {
-        let mut gauge_guard = GaugeGuard::from_gauge(&MEMORY_METRICS.in_flight.rest_server);
-        gauge_guard.add(content.len() as i64);
+        let gauge_guard = GaugeGuard::from_gauge_with_initial_value(
+            &MEMORY_METRICS.in_flight.rest_server,
+            content.len() as i64,
+        );
         Body {
             content,
             _gauge_guard: gauge_guard,
diff --git a/quickwit/quickwit-serve/src/load_shield.rs b/quickwit/quickwit-serve/src/load_shield.rs
@@ -78,13 +78,13 @@ impl LoadShield {
     }
 
     pub async fn acquire_permit(&'static self) -> Result<LoadShieldPermit, warp::Rejection> {
-        let mut pending_gauge_guard = GaugeGuard::from_gauge(&self.pending_gauge);
-        pending_gauge_guard.add(1);
+        let pending_gauge_guard =
+            GaugeGuard::from_gauge_with_initial_value(&self.pending_gauge, 1i64);
         let in_flight_permit_opt = self.acquire_in_flight_permit().await?;
         let concurrency_permit_opt = self.acquire_concurrency_permit().await;
         drop(pending_gauge_guard);
-        let mut ongoing_gauge_guard = GaugeGuard::from_gauge(&self.ongoing_gauge);
-        ongoing_gauge_guard.add(1);
+        let ongoing_gauge_guard =
+            GaugeGuard::from_gauge_with_initial_value(&self.ongoing_gauge, 1i64);
         Ok(LoadShieldPermit {
             _in_flight_permit_opt: in_flight_permit_opt,
             _concurrency_permit_opt: concurrency_permit_opt,
diff --git a/quickwit/quickwit-storage/src/metrics.rs b/quickwit/quickwit-storage/src/metrics.rs
@@ -232,12 +232,13 @@ pub static CACHE_METRICS_FOR_TESTS: Lazy<CacheMetrics> =
 pub fn object_storage_get_slice_in_flight_guards(
     get_request_size: usize,
 ) -> (GaugeGuard<'static>, GaugeGuard<'static>) {
-    let mut bytes_guard = GaugeGuard::from_gauge(
+    let bytes_guard = GaugeGuard::from_gauge_with_initial_value(
         &crate::STORAGE_METRICS.object_storage_get_slice_in_flight_num_bytes,
+        get_request_size as i64,
+    );
+    let count_guard = GaugeGuard::from_gauge_with_initial_value(
+        &crate::STORAGE_METRICS.object_storage_get_slice_in_flight_count,
+        1,
     );
-    bytes_guard.add(get_request_size as i64);
-    let mut count_guard =
-        GaugeGuard::from_gauge(&crate::STORAGE_METRICS.object_storage_get_slice_in_flight_count);
-    count_guard.add(1);
     (bytes_guard, count_guard)
 }

Original file line number	Diff line number	Diff line change
`@@ -395,9 +395,7 @@ fn get_actor_inboxes_count_gauge_guard() -> GaugeGuard<'static> {`
`395`	`395`	`&[],`
`396`	`396`	`)`
`397`	`397`	`});`
`398`		`- let mut gauge_guard = GaugeGuard::from_gauge(gauge);`
`399`		`- gauge_guard.add(1);`
`400`		`- gauge_guard`
	`398`	`+ GaugeGuard::from_gauge_with_initial_value(gauge, 1)`
`401`	`399`	`}`
`402`	`400`
`403`	`401`	`pub(crate) fn create_mailbox<A: Actor>(`
Original file line number	Diff line number	Diff line change
`@@ -219,8 +219,13 @@ impl std::fmt::Debug for GaugeGuard<'_> {`
`219`	`219`	`}`
`220`	`220`
`221`	`221`	`impl<'a> GaugeGuard<'a> {`
`222`		`- pub fn from_gauge(gauge: &'a IntGauge) -> Self {`
`223`		`- Self { gauge, delta: 0i64 }`
	`222`	`+ pub fn from_gauge_with_initial_value(gauge: &'a IntGauge, initial_value: i64) -> Self {`
	`223`	`+ let mut gauge = Self {`
	`224`	`+ gauge,`
	`225`	`+ delta: initial_value,`
	`226`	`+ };`
	`227`	`+ gauge.add(initial_value);`
	`228`	`+ gauge`
`224`	`229`	`}`
`225`	`230`
`226`	`231`	`pub fn get(&self) -> i64 {`
`@@ -256,8 +261,10 @@ impl std::fmt::Debug for OwnedGaugeGuard {`
`256`	`261`	`}`
`257`	`262`
`258`	`263`	`impl OwnedGaugeGuard {`
`259`		`- pub fn from_gauge(gauge: IntGauge) -> Self {`
`260`		`- Self { gauge, delta: 0i64 }`
	`264`	`+ pub fn from_gauge_with_initial_value(gauge: IntGauge, initial_value: i64) -> Self {`
	`265`	`+ let mut gauge = Self { gauge, delta: 0i64 };`
	`266`	`+ gauge.add(initial_value);`
	`267`	`+ gauge`
`261`	`268`	`}`
`262`	`269`
`263`	`270`	`pub fn get(&self) -> i64 {`