Use detailed mrecordlog responses; Add queue size tracking to ingester shard

quickwit/Cargo.toml

@@ -133,7 +133,7 @@ metrics-util = "0.20"
 mime_guess = "2.0"
 mini-moka = "0.10"
 mockall = "0.14"
-mrecordlog = { git = "https://github.com/quickwit-oss/mrecordlog", rev = "3b3562ef" }
+mrecordlog = { git = "https://github.com/quickwit-oss/mrecordlog", rev = "582ddb131181a053635e6904d3190491a3ee556c" }
 new_string_template = "1.5"
 nom = "8.0"
 numfmt = "1.2"

quickwit/quickwit-cli/src/source.rs

@@ -638,10 +638,12 @@ mod tests {
             flatten_json(json!(JsonValue::Null)),
             vec![("".to_string(), JsonValue::Null)]
         );
+        let mut flattened = flatten_json(
+            json!({"foo": {"bar": JsonValue::Bool(true)}, "baz": JsonValue::Bool(false)}),
+        );
+        flattened.sort_by(|left, right| left.0.cmp(&right.0));
         assert_eq!(
-            flatten_json(
-                json!({"foo": {"bar": JsonValue::Bool(true)}, "baz": JsonValue::Bool(false)})
-            ),
+            flattened,
             vec![
                 ("baz".to_string(), JsonValue::Bool(false)),
                 ("foo.bar".to_string(), JsonValue::Bool(true)),

quickwit/quickwit-indexing/src/source/ingest/mod.rs

@@ -1560,7 +1560,7 @@ mod tests {
             from_position_exclusive: Some(Position::offset(11u64)),
             to_position_inclusive: Some(Position::offset(14u64)),
         };
-        let batch_size = fetch_payload.estimate_size();
+        let batch_size = fetch_payload.buffer_size();
         let fetch_message = FetchMessage::new_payload(fetch_payload);
         let in_flight_value =
             InFlightValue::new(fetch_message, batch_size, &IN_FLIGHT_FETCH_STREAM);
@@ -1574,7 +1574,7 @@ mod tests {
             from_position_exclusive: Some(Position::offset(22u64)),
             to_position_inclusive: Some(Position::offset(23u64)),
         };
-        let batch_size = fetch_payload.estimate_size();
+        let batch_size = fetch_payload.buffer_size();
         let fetch_message = FetchMessage::new_payload(fetch_payload);
         let in_flight_value =
             InFlightValue::new(fetch_message, batch_size, &IN_FLIGHT_FETCH_STREAM);
@@ -1653,7 +1653,7 @@ mod tests {
             from_position_exclusive: Some(Position::offset(14u64)),
             to_position_inclusive: Some(Position::offset(15u64)),
         };
-        let batch_size = fetch_payload.estimate_size();
+        let batch_size = fetch_payload.buffer_size();
         let fetch_message = FetchMessage::new_payload(fetch_payload);
         let in_flight_value =
             InFlightValue::new(fetch_message, batch_size, &IN_FLIGHT_FETCH_STREAM);
@@ -1795,7 +1795,7 @@ mod tests {
             from_position_exclusive: Some(Position::offset(11u64)),
             to_position_inclusive: Some(Position::offset(13u64)),
         };
-        let batch_size = fetch_payload.estimate_size();
+        let batch_size = fetch_payload.buffer_size();
         let fetch_message = FetchMessage::new_payload(fetch_payload);
         let in_flight_value =
             InFlightValue::new(fetch_message, batch_size, &IN_FLIGHT_FETCH_STREAM);

quickwit/quickwit-ingest/src/ingest_v2/doc_mapper.rs

@@ -85,7 +85,7 @@ fn is_document_validation_enabled() -> bool {
     !quickwit_common::get_bool_from_env_cached!("QW_DISABLE_DOCUMENT_VALIDATION", false)
 }
 
-#[instrument(name = "ingester.validate_doc_batch", skip_all, fields(num_docs = doc_batch.num_docs(), num_bytes = doc_batch.num_bytes()))]
+#[instrument(name = "ingester.validate_doc_batch", skip_all, fields(num_docs = doc_batch.num_docs(), num_bytes = doc_batch.doc_buffer.len()))]
 async fn validate_doc_batch_cpu_intensive(
     doc_batch: DocBatchV2,
     doc_mapper: Arc<DocMapper>,

quickwit/quickwit-ingest/src/ingest_v2/fetch.rs

@@ -170,7 +170,7 @@ impl FetchStreamTask {
                     mrecord_buffer: mrecord_buffer.freeze(),
                     mrecord_lengths,
                 };
-                let batch_size = mrecord_batch.estimate_size();
+                let batch_size = mrecord_batch.buffer_size();
                 let fetch_payload = FetchPayload {
                     index_uid: Some(self.index_uid.clone()),
                     source_id: self.source_id.clone(),
@@ -490,7 +490,7 @@ async fn fetch_stream_once(
         match fetch_message_result {
             Ok(fetch_message) => match &fetch_message.message {
                 Some(fetch_message::Message::Payload(fetch_payload)) => {
-                    let batch_size = fetch_payload.estimate_size();
+                    let batch_size = fetch_payload.buffer_size();
                     let to_position_inclusive = fetch_payload.to_position_inclusive();
                     let in_flight_value = InFlightValue::new(
                         fetch_message,

quickwit/quickwit-ingest/src/ingest_v2/ingester.rs

@@ -51,11 +51,11 @@ use super::fetch::FetchStreamTask;
 use super::idle::CloseIdleShardsTask;
 use super::models::IngesterShard;
 use super::mrecordlog_utils::{
-    AppendDocBatchError, append_non_empty_doc_batch, check_enough_capacity, wal_stats,
+    AppendDocBatchError, append_non_empty_doc_batch, check_enough_capacity, doc_batch_size,
+    wal_stats,
 };
 use super::rate_meter::RateMeter;
 use super::state::{IngesterState, InnerIngesterState, WeakIngesterState};
-use crate::estimate_size;
 use crate::ingest_v2::doc_mapper::get_or_try_build_doc_mapper;
 use crate::ingest_v2::metrics::{
     DECOMMISSION_FAILED, DECOMMISSION_SUCCEEDED, RESET_SHARDS_OPERATIONS_TOTAL, STATUS,
@@ -472,10 +472,10 @@ impl Ingester {
                         DocBatchV2::default()
                     }
                 };
-                let requested_capacity = estimate_size(&doc_batch);
+                let requested_capacity = doc_batch_size(&doc_batch, force_commit);
 
                 if let Err(error) = check_enough_capacity(
-                    &state_guard.mrecordlog,
+                    &state_guard.mrecordlog.resource_usage(),
                     self.disk_capacity,
                     self.memory_capacity,
                     requested_capacity + total_requested_capacity,
@@ -513,7 +513,7 @@ impl Ingester {
                 }
 
                 // Total number of bytes (valid and invalid documents)
-                let original_batch_num_bytes = doc_batch.num_bytes() as u64;
+                let original_batch_num_bytes = doc_batch.doc_buffer.len() as u64;
 
                 let (valid_doc_batch, parse_failures) = if validate_docs {
                     validate_doc_batch(doc_batch, doc_mapper).await?
@@ -545,6 +545,7 @@ impl Ingester {
                     continue;
                 };
 
+                let valid_batch_num_bytes = valid_doc_batch.doc_buffer.len() as u64;
                 counter!(
                     parent: DOCS_TOTAL,
                     labels: [label_values!(VALIDITY => "valid")],
@@ -554,7 +555,7 @@ impl Ingester {
                     parent: DOCS_BYTES_TOTAL,
                     labels: [label_values!(VALIDITY => "valid")],
                 )
-                .inc_by(valid_doc_batch.num_bytes() as u64);
+                .inc_by(valid_batch_num_bytes);
                 if !parse_failures.is_empty() {
                     counter!(
                         parent: DOCS_TOTAL,
@@ -565,9 +566,8 @@ impl Ingester {
                         parent: DOCS_BYTES_TOTAL,
                         labels: [label_values!(VALIDITY => "invalid")],
                     )
-                    .inc_by(original_batch_num_bytes - valid_doc_batch.num_bytes() as u64);
+                    .inc_by(original_batch_num_bytes - valid_batch_num_bytes);
                 }
-                let valid_batch_num_bytes = valid_doc_batch.num_bytes() as u64;
                 shard.rate_meter.update(valid_batch_num_bytes);
                 total_requested_capacity += requested_capacity;
 
@@ -602,8 +602,8 @@ impl Ingester {
                 )
                 .await;
 
-                let current_position_inclusive = match append_result {
-                    Ok(current_position_inclusive) => current_position_inclusive,
+                let (current_position_inclusive, queue_size) = match append_result {
+                    Ok(append_outcome) => append_outcome,
                     Err(append_error) => {
                         let reason = match &append_error {
                             AppendDocBatchError::Io(io_error) => {
@@ -633,11 +633,12 @@ impl Ingester {
                     }
                 };
 
-                state_guard
+                let shard = state_guard
                     .shards
                     .get_mut(&queue_id)
-                    .expect("shard should exist")
-                    .set_replication_position_inclusive(current_position_inclusive.clone(), now);
+                    .expect("shard should exist");
+                shard.queue_size = queue_size;
+                shard.set_replication_position_inclusive(current_position_inclusive.clone(), now);
 
                 let persist_success = PersistSuccess {
                     subrequest_id: subrequest.subrequest_id,
@@ -915,7 +916,9 @@ impl IngesterService for Ingester {
             .subrequests
             .iter()
             .flat_map(|subrequest| match &subrequest.doc_batch {
-                Some(doc_batch) if doc_batch.doc_buffer.is_unique() => Some(doc_batch.num_bytes()),
+                Some(doc_batch) if doc_batch.doc_buffer.is_unique() => {
+                    Some(doc_batch.doc_buffer.len())
+                }
                 _ => None,
             })
             .sum::<usize>();
@@ -1566,6 +1569,51 @@ mod tests {
         assert!(state_guard.mrecordlog.queue_exists(&queue_id));
     }
 
+    #[tokio::test]
+    async fn test_persist_and_truncate_update_queue_size() {
+        let (ingester_ctx, ingester) = IngesterForTest::default().build().await;
+        let index_uid = IndexUid::for_test("test-index", 0);
+        let source_id = SourceId::from("test-source");
+        let doc_mapper = try_build_doc_mapper("{}").unwrap();
+        let shard = IngesterShard::builder(index_uid.clone(), source_id.clone(), ShardId::from(1))
+            .with_doc_mapper(doc_mapper)
+            .build();
+        let queue_id = shard.queue_id();
+        let mut state_guard = ingester.state.lock_fully("test").await.unwrap();
+        state_guard.shards.insert(queue_id.clone(), shard);
+        state_guard
+            .mrecordlog
+            .create_queue(&queue_id)
+            .await
+            .unwrap();
+        drop(state_guard);
+
+        let persist_request = PersistRequest {
+            ingester_id: ingester_ctx.node_id.to_string(),
+            commit_type: CommitTypeV2::Auto as i32,
+            subrequests: vec![PersistSubrequest {
+                subrequest_id: 0,
+                index_uid: Some(index_uid),
+                source_id,
+                doc_batch: Some(DocBatchV2::for_test(["test-doc-foo", "test-doc-barbaz"])),
+            }],
+        };
+        let persist_response = ingester.persist(persist_request).await.unwrap();
+        assert_eq!(persist_response.successes.len(), 1);
+
+        // Each doc record is prefixed with a 2-byte mrecord header.
+        let mut state_guard = ingester.state.lock_fully("test").await.unwrap();
+        assert_eq!(
+            state_guard.shards[&queue_id].queue_size,
+            ByteSize::b(14 + 17)
+        );
+
+        state_guard
+            .truncate_shard(&queue_id, Position::offset(0u64), "test")
+            .await;
+        assert_eq!(state_guard.shards[&queue_id].queue_size, ByteSize::b(17));
+    }
+
     #[tokio::test]
     async fn test_ingester_persist() {
         let (ingester_ctx, ingester) = IngesterForTest::default().build().await;

quickwit/quickwit-ingest/src/ingest_v2/mod.rs

@@ -57,7 +57,6 @@ use workbench::pending_subrequests;
 
 pub use self::fetch::{FetchStreamError, MultiFetchStream};
 pub use self::ingester::Ingester;
-use self::mrecord::MRECORD_HEADER_LEN;
 pub use self::mrecord::{MRecord, decoded_mrecords};
 pub use self::router::IngestRouter;
 
@@ -278,11 +277,6 @@ impl IngestRequestV2Builder {
     }
 }
 
-pub(super) fn estimate_size(doc_batch: &DocBatchV2) -> ByteSize {
-    let estimate = doc_batch.num_bytes() + doc_batch.num_docs() * MRECORD_HEADER_LEN;
-    ByteSize(estimate as u64)
-}
-
 #[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Ord, PartialOrd)]
 pub struct RateMibPerSec(pub u16);
 
@@ -333,7 +327,6 @@ mod tests {
         let doc_batch = doc_batch_builder.build().unwrap();
 
         assert_eq!(doc_batch.num_docs(), 2);
-        assert_eq!(doc_batch.num_bytes(), 21);
         assert_eq!(doc_batch.doc_lengths, [7, 6]);
         assert_eq!(doc_batch.doc_buffer, Bytes::from(&b"Hello, World!"[..]));
     }
@@ -384,14 +377,6 @@ mod tests {
                 .num_docs(),
             2
         );
-        assert_eq!(
-            ingest_request.subrequests[0]
-                .doc_batch
-                .as_ref()
-                .unwrap()
-                .num_bytes(),
-            21
-        );
         assert_eq!(
             ingest_request.subrequests[0]
                 .doc_batch
@@ -427,14 +412,6 @@ mod tests {
                 .num_docs(),
             2
         );
-        assert_eq!(
-            ingest_request.subrequests[1]
-                .doc_batch
-                .as_ref()
-                .unwrap()
-                .num_bytes(),
-            20
-        );
         assert_eq!(
             ingest_request.subrequests[1]
                 .doc_batch
@@ -460,21 +437,4 @@ mod tests {
             [hola_doc_uid, mundo_doc_uid]
         );
     }
-
-    #[test]
-    fn test_estimate_size() {
-        let doc_batch = DocBatchV2 {
-            doc_buffer: Vec::new().into(),
-            doc_lengths: Vec::new(),
-            doc_uids: Vec::new(),
-        };
-        assert_eq!(estimate_size(&doc_batch), ByteSize(0));
-
-        let doc_batch = DocBatchV2 {
-            doc_buffer: vec![0u8; 100].into(),
-            doc_lengths: vec![10, 20, 30],
-            doc_uids: Vec::new(),
-        };
-        assert_eq!(estimate_size(&doc_batch), ByteSize(118));
-    }
 }