Added metrics tracking the memory usage in the index writer. (#4724)

* Added metrics tracking the memory usage in the index writer.

* Fix fmt

---------

Co-authored-by: Adrien Guillo <adrien@quickwit.io>

Author: fulmicoton

quickwit/quickwit-common/src/metrics.rs

@@ -173,11 +173,19 @@ pub struct GaugeGuard {
     delta: i64,
 }
 
+impl std::fmt::Debug for GaugeGuard {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        self.delta.fmt(f)
+    }
+}
+
 impl GaugeGuard {
-    pub fn from_gauge(gauge: &'static IntGauge, delta: i64) -> Self {
-        gauge.add(delta);
+    pub fn from_gauge(gauge: &'static IntGauge) -> Self {
+        Self { gauge, delta: 0i64 }
+    }
 
-        Self { gauge, delta }
+    pub fn get(&self) -> i64 {
+        self.delta
     }
 
     pub fn add(&mut self, delta: i64) {
@@ -245,46 +253,34 @@ impl Default for MemoryMetrics {
 #[derive(Clone)]
 pub struct InFlightDataGauges {
     pub doc_processor_mailbox: IntGauge,
+    pub index_writer: IntGauge,
     pub indexer_mailbox: IntGauge,
     pub ingest_router: IntGauge,
     pub rest_server: IntGauge,
     pub sources: InFlightDataSourceGauges,
 }
 
-const IN_FLIGHT_DATA_GAUGES_HELP: &str = "Amount of data in-flight in various buffers in bytes.";
-
 impl Default for InFlightDataGauges {
     fn default() -> Self {
+        let in_flight_gauge_vec = new_gauge_vec(
+            "in_flight_data_bytes",
+            "Amount of data in-flight in various buffers in bytes.",
+            "memory",
+            &[],
+            ["component"],
+        );
         Self {
-            doc_processor_mailbox: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "doc_processor_mailbox")],
-            ),
-            indexer_mailbox: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "indexer_mailbox")],
-            ),
-            ingest_router: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "ingest_router")],
-            ),
-            rest_server: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "rest_server")],
-            ),
-            sources: InFlightDataSourceGauges::default(),
+            doc_processor_mailbox: in_flight_gauge_vec.with_label_values(["doc_processor_mailbox"]),
+            index_writer: in_flight_gauge_vec.with_label_values(["index_writer"]),
+            indexer_mailbox: in_flight_gauge_vec.with_label_values(["indexer_mailbox"]),
+            ingest_router: in_flight_gauge_vec.with_label_values(["ingest_router"]),
+            rest_server: in_flight_gauge_vec.with_label_values(["rest_server"]),
+            sources: InFlightDataSourceGauges::new(&in_flight_gauge_vec),
         }
     }
 }
 
+/// TODO make those lazy.
 #[derive(Clone)]
 pub struct InFlightDataSourceGauges {
     pub file: IntGauge,
@@ -296,51 +292,16 @@ pub struct InFlightDataSourceGauges {
     pub other: IntGauge,
 }
 
-impl Default for InFlightDataSourceGauges {
-    fn default() -> Self {
+impl InFlightDataSourceGauges {
+    pub fn new(in_flight_gauge_vec: &IntGaugeVec<1>) -> Self {
         Self {
-            file: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "file_source")],
-            ),
-            ingest: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "ingest_source")],
-            ),
-            kafka: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "kafka_source")],
-            ),
-            kinesis: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "kinesis_source")],
-            ),
-            pubsub: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "pubsub_source")],
-            ),
-            pulsar: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "pulsar")],
-            ),
-            other: new_gauge(
-                "in_flight_data_bytes",
-                IN_FLIGHT_DATA_GAUGES_HELP,
-                "memory",
-                &[("component", "other")],
-            ),
+            file: in_flight_gauge_vec.with_label_values(["file_source"]),
+            ingest: in_flight_gauge_vec.with_label_values(["ingest_source"]),
+            kafka: in_flight_gauge_vec.with_label_values(["kafka_source"]),
+            kinesis: in_flight_gauge_vec.with_label_values(["kinesis_source"]),
+            pubsub: in_flight_gauge_vec.with_label_values(["pubsub_source"]),
+            pulsar: in_flight_gauge_vec.with_label_values(["pulsar_source"]),
+            other: in_flight_gauge_vec.with_label_values(["other"]),
         }
     }
 }

quickwit/quickwit-control-plane/src/control_plane.rs

@@ -55,7 +55,7 @@ use quickwit_proto::metastore::{
 };
 use quickwit_proto::types::{IndexUid, NodeId, ShardId, SourceUid};
 use serde::Serialize;
-use tracing::{error, info};
+use tracing::{debug, error, info};
 
 use crate::debouncer::Debouncer;
 use crate::indexing_scheduler::{IndexingScheduler, IndexingSchedulerState};
@@ -513,11 +513,12 @@ impl Handler<DeleteIndexRequest> for ControlPlane {
         ctx: &ActorContext<Self>,
     ) -> Result<Self::Reply, ActorExitStatus> {
         let index_uid: IndexUid = request.index_uid().clone();
-        info!(index=%index_uid, "delete index");
+        debug!(%index_uid, "deleting index");
 
         if let Err(metastore_error) = self.metastore.delete_index(request).await {
             return convert_metastore_error(metastore_error);
         };
+        info!(%index_uid, "deleted index");
 
         let ingester_needing_resync: BTreeSet<NodeId> = self
             .model
@@ -558,16 +559,18 @@ impl Handler<AddSourceRequest> for ControlPlane {
                     return Ok(Err(ControlPlaneError::from(error)));
                 }
             };
-        info!(index=%index_uid, source_config=?source_config, "add source");
+        let source_id = source_config.source_id.clone();
+        debug!(%index_uid, source_id, "adding source");
 
         if let Err(error) = self.metastore.add_source(request).await {
             return Ok(Err(ControlPlaneError::from(error)));
         };
-
         self.model
             .add_source(&index_uid, source_config)
             .context("failed to add source")?;
 
+        info!(%index_uid, source_id, "added source");
+
         // TODO: Refine the event. Notify index will have the effect to reload the entire state from
         // the metastore. We should update the state of the control plane.
         let _rebuild_plan_waiter = self.rebuild_plan_debounced(ctx);
@@ -591,12 +594,13 @@ impl Handler<ToggleSourceRequest> for ControlPlane {
         let index_uid: IndexUid = request.index_uid().clone();
         let source_id = request.source_id.clone();
         let enable = request.enable;
-
-        info!(index=%index_uid, source_id=%source_id, enable=enable, "toggle source");
+        debug!(%index_uid, source_id, enable, "toggling source");
 
         if let Err(error) = self.metastore.toggle_source(request).await {
             return Ok(Err(ControlPlaneError::from(error)));
         };
+        info!(%index_uid, source_id, enabled=enable, "toggled source");
+
         let mutation_occured = self.model.toggle_source(&index_uid, &source_id, enable)?;
 
         if mutation_occured {

quickwit/quickwit-indexing/src/actors/indexer.rs

@@ -33,6 +33,7 @@ use quickwit_actors::{
     Actor, ActorContext, ActorExitStatus, Command, Handler, Mailbox, QueueCapacity,
 };
 use quickwit_common::io::IoControls;
+use quickwit_common::metrics::GaugeGuard;
 use quickwit_common::runtimes::RuntimeType;
 use quickwit_common::temp_dir::TempDirectory;
 use quickwit_config::IndexingSettings;
@@ -151,10 +152,10 @@ impl IndexerState {
         other_split_opt: &'a mut Option<IndexedSplitBuilder>,
         counter: &'a mut IndexerCounters,
         ctx: &ActorContext<Indexer>,
-    ) -> anyhow::Result<&'a mut IndexedSplitBuilder> {
+    ) -> anyhow::Result<(&'a mut IndexedSplitBuilder, bool)> {
         let num_splits = splits.len();
         match splits.entry(partition_id) {
-            Entry::Occupied(indexed_split) => Ok(indexed_split.into_mut()),
+            Entry::Occupied(indexed_split) => Ok((indexed_split.into_mut(), false)),
             Entry::Vacant(vacant_entry) => {
                 if num_splits as u32 >= self.max_num_partitions.get() {
                     // In order to avoid exceeding max_num_partitions, we map the document to the
@@ -172,11 +173,11 @@ impl IndexerState {
                         )?;
                         *other_split_opt = Some(new_other_split);
                     }
-                    Ok(other_split_opt.as_mut().unwrap())
+                    Ok((other_split_opt.as_mut().unwrap(), true))
                 } else {
                     let indexed_split =
                         self.create_indexed_split_builder(partition_id, last_delete_opstamp, ctx)?;
-                    Ok(vacant_entry.insert(indexed_split))
+                    Ok((vacant_entry.insert(indexed_split), true))
                 }
             }
         }
@@ -236,7 +237,11 @@ impl IndexerState {
             publish_lock,
             publish_token_opt,
             last_delete_opstamp,
-            memory_usage: ByteSize(0),
+            memory_usage: GaugeGuard::from_gauge(
+                &quickwit_common::metrics::MEMORY_METRICS
+                    .in_flight_data
+                    .index_writer,
+            ),
         };
         Ok(workbench)
     }
@@ -294,7 +299,7 @@ impl IndexerState {
             .source_delta
             .extend(batch.checkpoint_delta)
             .context("batch delta does not follow indexer checkpoint")?;
-        let mut memory_usage_delta: u64 = 0;
+        let mut memory_usage_delta: i64 = 0;
         counters.num_doc_batches_in_workbench += 1;
         for doc in batch.docs {
             let ProcessedDoc {
@@ -304,7 +309,7 @@ impl IndexerState {
                 num_bytes,
             } = doc;
             counters.num_docs_in_workbench += 1;
-            let indexed_split: &mut IndexedSplitBuilder = self.get_or_create_indexed_split(
+            let (indexed_split, split_created) = self.get_or_create_indexed_split(
                 partition,
                 *last_delete_opstamp,
                 indexed_splits,
@@ -313,6 +318,11 @@ impl IndexerState {
                 ctx,
             )?;
             let mem_usage_before = indexed_split.index_writer.mem_usage() as u64;
+            if split_created {
+                // The split was just created. We need to account for the initial index writer's
+                // memory usage.
+                memory_usage_delta += mem_usage_before as i64;
+            }
             indexed_split.split_attrs.uncompressed_docs_size_in_bytes += num_bytes as u64;
             indexed_split.split_attrs.num_docs += 1;
             if let Some(timestamp) = timestamp_opt {
@@ -324,10 +334,10 @@ impl IndexerState {
                 .add_document(doc)
                 .context("failed to add document")?;
             let mem_usage_after = indexed_split.index_writer.mem_usage() as u64;
-            memory_usage_delta += mem_usage_after - mem_usage_before;
+            memory_usage_delta += mem_usage_after as i64 - mem_usage_before as i64;
             ctx.record_progress();
         }
-        *memory_usage = ByteSize(memory_usage.as_u64() + memory_usage_delta);
+        memory_usage.add(memory_usage_delta);
         Ok(())
     }
 }
@@ -353,7 +363,7 @@ struct IndexingWorkbench {
     // We use this value to set the `delete_opstamp` of the workbench splits.
     last_delete_opstamp: u64,
     // Number of bytes declared as used by tantivy.
-    memory_usage: ByteSize,
+    memory_usage: GaugeGuard,
 }
 
 pub struct Indexer {
@@ -570,9 +580,9 @@ impl Indexer {
 
     fn memory_usage(&self) -> ByteSize {
         if let Some(workbench) = &self.indexing_workbench_opt {
-            workbench.memory_usage
+            ByteSize(workbench.memory_usage.get() as u64)
         } else {
-            ByteSize(0)
+            ByteSize(0u64)
         }
     }
 
@@ -591,7 +601,8 @@ impl Indexer {
                 ctx,
             )
             .await?;
-        if self.memory_usage() >= self.indexer_state.indexing_settings.resources.heap_size {
+        let memory_usage = self.memory_usage();
+        if memory_usage >= self.indexer_state.indexing_settings.resources.heap_size {
             self.send_to_serializer(CommitTrigger::MemoryLimit, ctx)
                 .await?;
         }
@@ -623,6 +634,7 @@ impl Indexer {
             publish_token_opt,
             batch_parent_span,
             indexing_permit,
+            memory_usage,
             ..
         }) = self.indexing_workbench_opt.take()
         else {
@@ -674,6 +686,7 @@ impl Indexer {
                 publish_token_opt,
                 commit_trigger,
                 batch_parent_span,
+                memory_usage,
             },
         )
         .await?;
@@ -883,7 +896,7 @@ mod tests {
         let body_field = schema.get_field("body").unwrap();
         let indexing_directory = TempDirectory::for_test();
         let mut indexing_settings = IndexingSettings::for_test();
-        indexing_settings.resources.heap_size = ByteSize::mb(5);
+        indexing_settings.resources.heap_size = ByteSize::mb(16);
         let (index_serializer_mailbox, index_serializer_inbox) = universe.create_test_mailbox();
         let mut metastore = MetastoreServiceClient::mock();
         metastore.expect_publish_splits().never();
@@ -1214,7 +1227,8 @@ mod tests {
         let body_field = schema.get_field("body").unwrap();
 
         let indexing_directory = TempDirectory::for_test();
-        let indexing_settings = IndexingSettings::for_test();
+        let mut indexing_settings = IndexingSettings::for_test();
+        indexing_settings.resources.heap_size = ByteSize::mb(100);
         let (index_serializer_mailbox, index_serializer_inbox) = universe.create_test_mailbox();
         let mut metastore = MetastoreServiceClient::mock();
         metastore.expect_publish_splits().never();
@@ -1310,7 +1324,8 @@ mod tests {
             Arc::new(serde_json::from_str::<DefaultDocMapper>(DOCMAPPER_SIMPLE_JSON).unwrap());
         let body_field = doc_mapper.schema().get_field("body").unwrap();
         let indexing_directory = TempDirectory::for_test();
-        let indexing_settings = IndexingSettings::for_test();
+        let mut indexing_settings = IndexingSettings::for_test();
+        indexing_settings.resources.heap_size = ByteSize::gb(5);
         let mut metastore = MetastoreServiceClient::mock();
         metastore
             .expect_last_delete_opstamp()

quickwit/quickwit-indexing/src/models/indexed_split.rs

@@ -21,6 +21,7 @@ use std::fmt;
 use std::path::Path;
 
 use quickwit_common::io::IoControls;
+use quickwit_common::metrics::GaugeGuard;
 use quickwit_common::temp_dir::TempDirectory;
 use quickwit_metastore::checkpoint::IndexCheckpointDelta;
 use quickwit_proto::indexing::IndexingPipelineId;
@@ -183,6 +184,7 @@ pub struct IndexedSplitBatchBuilder {
     pub publish_token_opt: Option<PublishToken>,
     pub commit_trigger: CommitTrigger,
     pub batch_parent_span: Span,
+    pub memory_usage: GaugeGuard,
 }
 
 /// Sends notifications to the Publisher that the last batch of splits was emtpy.