Reuse HTTP request buffers for native ingest

quickwit/Cargo.toml

@@ -126,6 +126,7 @@ libz-sys = "1.1"
 lru = "0.18"
 matches = "0.1"
 md5 = "0.8"
+memchr = "2"
 metrics = "0.24"
 metrics-opentelemetry = "=0.24032.8"
 metrics-exporter-prometheus = { version = "0.18", default-features = false }

quickwit/quickwit-ingest/Cargo.toml

@@ -19,6 +19,7 @@ fail = { workspace = true, optional = true }
 futures = { workspace = true }
 http = { workspace = true }
 itertools = { workspace = true }
+memchr = { workspace = true }
 mockall = { workspace = true, optional = true }
 mrecordlog = { workspace = true }
 prost = { workspace = true }

quickwit/quickwit-ingest/benches/mrecord_bench.rs

@@ -28,7 +28,8 @@ use std::time::{Duration, Instant};
 use bytes::{BufMut, Bytes};
 use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
 use mrecordlog::MultiRecordLog;
-use quickwit_ingest::MRecord;
+use quickwit_ingest::{DocBatchV2Builder, MRecord};
+use quickwit_proto::types::DocUidGenerator;
 
 /// Representative document sizes (bytes): a small log line up to a large structured event.
 const DOC_SIZES: [usize; 4] = [128, 1_024, 8_192, 65_536];
@@ -180,5 +181,44 @@ fn bench_append_batch(criterion: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_encode, bench_decode, bench_append_batch);
+fn bench_doc_batch_builder(criterion: &mut Criterion) {
+    let mut group = criterion.benchmark_group("doc_batch_v2_builder");
+    let docs: Vec<Bytes> = (0..10_000)
+        .map(|index| make_doc(BATCH_DOC_SIZES[index % BATCH_DOC_SIZES.len()]))
+        .collect();
+    let total_bytes: usize = docs.iter().map(Bytes::len).sum();
+    group.throughput(Throughput::Bytes(total_bytes as u64));
+
+    group.bench_function("default", |bencher| {
+        bencher.iter(|| {
+            let mut builder = DocBatchV2Builder::default();
+            let mut doc_uid_generator = DocUidGenerator::default();
+            for doc in &docs {
+                builder.add_doc(doc_uid_generator.next_doc_uid(), doc);
+            }
+            black_box(builder.build())
+        });
+    });
+
+    group.bench_function("preallocated", |bencher| {
+        bencher.iter(|| {
+            let mut builder = DocBatchV2Builder::with_capacity(total_bytes);
+            let mut doc_uid_generator = DocUidGenerator::default();
+            for doc in &docs {
+                builder.add_doc(doc_uid_generator.next_doc_uid(), doc);
+            }
+            black_box(builder.build())
+        });
+    });
+
+    group.finish();
+}
+
+criterion_group!(
+    benches,
+    bench_encode,
+    bench_decode,
+    bench_append_batch,
+    bench_doc_batch_builder
+);
 criterion_main!(benches);

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

@@ -153,6 +153,14 @@ pub struct DocBatchV2Builder {
 }
 
 impl DocBatchV2Builder {
+    pub fn with_capacity(doc_buffer_capacity: usize) -> Self {
+        Self {
+            doc_uids: Vec::new(),
+            doc_buffer: BytesMut::with_capacity(doc_buffer_capacity),
+            doc_lengths: Vec::new(),
+        }
+    }
+
     /// Adds a document to the batch.
     pub fn add_doc(&mut self, doc_uid: DocUid, doc: &[u8]) {
         self.doc_uids.push(doc_uid);

quickwit/quickwit-ingest/src/lib.rs

@@ -24,6 +24,7 @@ mod ingest_v2;
 mod memory_capacity;
 mod metrics;
 mod mrecordlog_async;
+mod ndjson;
 mod notifications;
 mod position;
 mod queue;
@@ -39,6 +40,7 @@ pub use ingest_api_service::{GetMemoryCapacity, GetPartitionId, IngestApiService
 pub use ingest_service::*;
 pub use ingest_v2::*;
 pub use memory_capacity::MemoryCapacity;
+pub use ndjson::{doc_batch_v2_from_ndjson, split_ndjson_lines};
 pub use position::Position;
 pub use queue::Queues;
 use quickwit_actors::{Mailbox, Universe};

quickwit/quickwit-ingest/src/ndjson.rs

@@ -0,0 +1,151 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use bytes::Bytes;
+use quickwit_proto::ingest::DocBatchV2;
+use quickwit_proto::types::DocUidGenerator;
+
+/// Iterates over the non-blank lines of an NDJSON body.
+pub fn split_ndjson_lines(ndjson_body: &[u8]) -> impl Iterator<Item = &[u8]> {
+    let mut line_start = 0;
+    memchr::memchr_iter(b'\n', ndjson_body)
+        .chain(std::iter::once(ndjson_body.len()))
+        .map(move |line_end| {
+            let line = &ndjson_body[line_start..line_end];
+            line_start = line_end + 1;
+            line
+        })
+        .filter(|line| !is_empty_or_blank_line(line))
+}
+
+/// Builds a [`DocBatchV2`] whose doc buffer is the NDJSON body itself, without copying. The docs
+/// must cover the whole buffer, so newlines and blank lines stay attached to neighboring docs.
+pub fn doc_batch_v2_from_ndjson(ndjson_body: Bytes) -> Option<DocBatchV2> {
+    let mut doc_uids = Vec::new();
+    let mut doc_lengths = Vec::new();
+    let mut doc_uid_generator = DocUidGenerator::default();
+    let mut segment_start = 0usize;
+    let mut line_start = 0usize;
+
+    for position in memchr::memchr_iter(b'\n', &ndjson_body) {
+        let line = &ndjson_body[line_start..position];
+        if !is_empty_or_blank_line(line) {
+            doc_uids.push(doc_uid_generator.next_doc_uid());
+            doc_lengths.push((position + 1 - segment_start) as u32);
+            segment_start = position + 1;
+        }
+        line_start = position + 1;
+    }
+
+    let line = &ndjson_body[line_start..];
+    if !is_empty_or_blank_line(line) {
+        doc_uids.push(doc_uid_generator.next_doc_uid());
+        doc_lengths.push((ndjson_body.len() - segment_start) as u32);
+        segment_start = ndjson_body.len();
+    }
+
+    if doc_uids.is_empty() {
+        return None;
+    }
+    if segment_start < ndjson_body.len() {
+        let trailing_whitespace_len = ndjson_body.len() - segment_start;
+        let last_doc_len = doc_lengths
+            .last_mut()
+            .expect("doc lengths should not be empty");
+        *last_doc_len += trailing_whitespace_len as u32;
+    }
+
+    Some(DocBatchV2 {
+        doc_uids,
+        doc_buffer: ndjson_body,
+        doc_lengths,
+    })
+}
+
+#[inline]
+fn is_empty_or_blank_line(line: &[u8]) -> bool {
+    line.is_empty() || line.iter().all(|ch| ch.is_ascii_whitespace())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str;
+
+    use super::*;
+
+    #[test]
+    fn test_split_ndjson_lines() {
+        let test_cases = [
+            // an empty line is inserted before the metadata action and the doc
+            (&b"\n{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n{\"id\": 1, \"message\": \"push\"}"[..], 2),
+            // a blank line is inserted before the metadata action and the doc
+            (&b"       \n{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n{\"id\": 1, \"message\": \"push\"}"[..], 2),
+            // an empty line is inserted after the metadata action and before the doc
+            (&b"{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n\n{\"id\": 1, \"message\": \"push\"}"[..], 2),
+            // a blank line is inserted after the metadata action and before the doc
+            (&b"{ \"create\" : { \"_index\" : \"my-index-1\", \"_id\" : \"1\"} }\n     \n{\"id\": 1, \"message\": \"push\"}"[..], 2),
+        ];
+        for &(input, expected_count) in &test_cases {
+            assert_eq!(split_ndjson_lines(input).count(), expected_count);
+        }
+        let lines: Vec<&[u8]> = split_ndjson_lines(b"a\n\nbc\n  \nd").collect();
+        assert_eq!(lines, [&b"a"[..], b"bc", b"d"]);
+
+        let lines: Vec<&[u8]> = split_ndjson_lines(b"a\n").collect();
+        assert_eq!(lines, [&b"a"[..]]);
+
+        assert_eq!(split_ndjson_lines(b"").count(), 0);
+        assert_eq!(split_ndjson_lines(b"\n \n\t").count(), 0);
+    }
+
+    #[test]
+    fn test_doc_batch_v2_from_ndjson_zero_copy() {
+        let body = Bytes::from_static(b"\n  {\"id\":1}\n\n{\"id\":2}\n   \n");
+        let doc_batch = doc_batch_v2_from_ndjson(body.clone()).unwrap();
+        assert_eq!(doc_batch.num_docs(), 2);
+        assert_eq!(doc_batch.doc_buffer, body);
+
+        let docs: Vec<Bytes> = doc_batch.docs().map(|(_doc_uid, doc)| doc).collect();
+        assert_eq!(str::from_utf8(&docs[0]).unwrap(), "\n  {\"id\":1}\n");
+        assert_eq!(str::from_utf8(&docs[1]).unwrap(), "\n{\"id\":2}\n   \n");
+    }
+
+    #[test]
+    fn test_doc_batch_v2_from_ndjson_no_trailing_newline() {
+        let body = Bytes::from_static(b"{\"id\":1}\n{\"id\":2}");
+        let doc_batch = doc_batch_v2_from_ndjson(body).unwrap();
+        let docs: Vec<Bytes> = doc_batch.docs().map(|(_doc_uid, doc)| doc).collect();
+        assert_eq!(docs, [&b"{\"id\":1}\n"[..], b"{\"id\":2}"]);
+    }
+
+    #[test]
+    fn test_doc_batch_v2_from_blank_ndjson() {
+        assert!(doc_batch_v2_from_ndjson(Bytes::new()).is_none());
+        assert!(doc_batch_v2_from_ndjson(Bytes::from_static(b"\n \n\t")).is_none());
+    }
+
+    #[test]
+    fn test_doc_batch_v2_from_ndjson_matches_split_ndjson_lines() {
+        let body = Bytes::from_static(b"  \n{\"a\":1}\n\n \t\n{\"b\":2}\n{\"c\":3}\n\n");
+        let doc_batch = doc_batch_v2_from_ndjson(body.clone()).unwrap();
+        let docs: Vec<Bytes> = doc_batch.docs().map(|(_doc_uid, doc)| doc).collect();
+        let lines: Vec<&[u8]> = split_ndjson_lines(&body).collect();
+        assert_eq!(docs.len(), lines.len());
+        for (doc, line) in docs.iter().zip(&lines) {
+            assert_eq!(doc.trim_ascii(), *line);
+        }
+        let doc_uids: Vec<_> = doc_batch.docs().map(|(doc_uid, _doc)| doc_uid).collect();
+        assert!(doc_uids.windows(2).all(|pair| pair[0] < pair[1]));
+    }
+}

quickwit/quickwit-serve/src/elasticsearch_api/bulk.rs

@@ -18,6 +18,7 @@ use std::time::Instant;
 use bytesize::ByteSize;
 use quickwit_ingest::{
     CommitType, DocBatchBuilder, IngestRequest, IngestService, IngestServiceClient,
+    split_ndjson_lines,
 };
 use quickwit_proto::ingest::router::IngestRouterServiceClient;
 use quickwit_proto::types::IndexId;
@@ -29,7 +30,6 @@ use crate::elasticsearch_api::filter::{elastic_bulk_filter, elastic_index_bulk_f
 use crate::elasticsearch_api::make_elastic_api_response;
 use crate::elasticsearch_api::model::{BulkAction, ElasticBulkOptions, ElasticsearchError};
 use crate::format::extract_format_from_qs;
-use crate::ingest_api::lines;
 use crate::rest::recover_fn;
 use crate::{Body, with_arg};
 
@@ -111,7 +111,7 @@ async fn elastic_ingest_bulk(
     }
     let now = Instant::now();
     let mut doc_batch_builders = HashMap::new();
-    let mut lines = lines(&body.content).enumerate();
+    let mut lines = split_ndjson_lines(&body.content).enumerate();
 
     while let Some((line_number, line)) = lines.next() {
         let action = serde_json::from_slice::<BulkAction>(line).map_err(|error| {

quickwit/quickwit-serve/src/elasticsearch_api/bulk_v2.rs

@@ -17,7 +17,7 @@ use std::time::Instant;
 
 use quickwit_common::rate_limited_error;
 use quickwit_config::{INGEST_V2_SOURCE_ID, validate_identifier};
-use quickwit_ingest::IngestRequestV2Builder;
+use quickwit_ingest::{IngestRequestV2Builder, split_ndjson_lines};
 use quickwit_proto::ingest::CommitTypeV2;
 use quickwit_proto::ingest::router::{
     IngestFailureReason, IngestResponseV2, IngestRouterService, IngestRouterServiceClient,
@@ -29,7 +29,6 @@ use warp::hyper::StatusCode;
 use super::model::ElasticException;
 use crate::Body;
 use crate::elasticsearch_api::model::{BulkAction, ElasticBulkOptions, ElasticsearchError};
-use crate::ingest_api::lines;
 
 #[derive(Debug, Default, Serialize, Deserialize)]
 pub(crate) struct ElasticBulkResponse {
@@ -88,7 +87,7 @@ pub(crate) async fn elastic_bulk_ingest_v2(
 ) -> Result<ElasticBulkResponse, ElasticsearchError> {
     let now = Instant::now();
     let mut ingest_request_builder = IngestRequestV2Builder::default();
-    let mut lines = lines(&body.content).enumerate();
+    let mut lines = split_ndjson_lines(&body.content).enumerate();
     let mut per_subrequest_doc_handles: HashMap<u32, Vec<DocHandle>> = HashMap::new();
     let mut action_count = 0;
     let mut invalid_index_id_items = Vec::new();

quickwit/quickwit-serve/src/ingest_api/mod.rs

@@ -16,7 +16,7 @@ mod response;
 mod rest_handler;
 
 pub use response::{RestIngestResponse, RestParseFailure};
+pub(crate) use rest_handler::ingest_api_handlers;
 #[cfg(test)]
 pub(crate) use rest_handler::tests::setup_ingest_v1_service;
 pub use rest_handler::{IngestApi, IngestApiSchemas};
-pub(crate) use rest_handler::{ingest_api_handlers, lines};