From 1a628e48eea68d6ab7be7450eb74b72abb32178d Mon Sep 17 00:00:00 2001 From: Pierre Massat Date: Thu, 10 Sep 2026 14:12:01 -0700 Subject: [PATCH 1/2] perf(eap): Stream RowBinary INSERTs instead of buffering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace Reduce → RowBinaryWriterStep with a streaming writer so EAP consumers compress and POST RowBinary blocks as messages arrive instead of holding the full uncompressed batch. Unlike #8001, keep clickhouse_concurrency in-flight slots and only retry after the body is complete so a mid-stream failure cannot ACK a truncated insert. Refs LINEAR-EAP-517 --- rust_snuba/Cargo.lock | 77 +- rust_snuba/Cargo.toml | 3 +- rust_snuba/src/factory_v2.rs | 101 ++- rust_snuba/src/strategies/clickhouse/mod.rs | 2 + .../strategies/clickhouse/streaming_lz4.rs | 202 +++++ .../strategies/clickhouse/streaming_writer.rs | 849 ++++++++++++++++++ .../src/strategies/clickhouse/writer_v2.rs | 260 +++++- rust_snuba/src/types.rs | 18 + 8 files changed, 1434 insertions(+), 78 deletions(-) create mode 100644 rust_snuba/src/strategies/clickhouse/streaming_lz4.rs create mode 100644 rust_snuba/src/strategies/clickhouse/streaming_writer.rs diff --git a/rust_snuba/Cargo.lock b/rust_snuba/Cargo.lock index 8039e9b5384..ee13580e3e8 100644 --- a/rust_snuba/Cargo.lock +++ b/rust_snuba/Cargo.lock @@ -1213,9 +1213,9 @@ dependencies = [ [[package]] name = "futures" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" dependencies = [ "futures-channel", "futures-core", @@ -1228,9 +1228,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", "futures-sink", @@ -1238,15 +1238,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" [[package]] name = "futures-executor" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" dependencies = [ "futures-core", "futures-task", @@ -1255,9 +1255,9 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" [[package]] name = "futures-lite" @@ -1289,26 +1289,26 @@ dependencies = [ [[package]] name = "futures-macro" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.5", ] [[package]] name = "futures-sink" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" [[package]] name = "futures-task" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" [[package]] name = "futures-timer" @@ -1318,9 +1318,9 @@ checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" [[package]] name = "futures-util" -version = "0.3.30" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-channel", "futures-core", @@ -1330,7 +1330,6 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "slab", ] @@ -1620,7 +1619,7 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "socket2 0.4.10", + "socket2 0.5.6", "tokio", "tower-service", "tracing", @@ -3402,12 +3401,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-native-tls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", ] @@ -4276,6 +4277,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -5056,12 +5068,14 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.42" +version = "0.4.64" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76bc14366121efc8dbb487ab05bcc9d346b3b5ec0eaa76e46594cabbe51762c0" +checksum = "e9c5522b3a28661442748e09d40924dfb9ca614b21c00d3fd135720e48b67db8" dependencies = [ "cfg-if", + "futures-util", "js-sys", + "once_cell", "wasm-bindgen", "web-sys", ] @@ -5120,6 +5134,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" @@ -5134,9 +5161,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.69" +version = "0.3.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77afa9a11836342370f4817622a2f0f418b134426d91a82dfb48f532d2ec13ef" +checksum = "854ba17bb104abfb26ba36da9729addc7ce7f06f5c0f90f3c391f8461cca21f9" dependencies = [ "js-sys", "wasm-bindgen", diff --git a/rust_snuba/Cargo.toml b/rust_snuba/Cargo.toml index da627a2a9b9..6c913d7a185 100644 --- a/rust_snuba/Cargo.toml +++ b/rust_snuba/Cargo.toml @@ -28,7 +28,7 @@ chrono = { version = "0.4.26", features = ["serde"] } cityhash-rs = "1.0.1" ctrlc = { version = "3.2.5", features = ["termination"] } data-encoding = "2.5.0" -futures = "0.3.21" +futures = "0.3.31" hyper = "1.2.0" json-schema-diff = "0.1.7" lz4_flex = "0.11" @@ -51,6 +51,7 @@ reqwest = { version = "0.13", default-features = false, features = [ "http2", "native-tls", "query", + "stream", "system-proxy", ] } schemars = { version = "0.8.16", features = ["uuid1"] } diff --git a/rust_snuba/src/factory_v2.rs b/rust_snuba/src/factory_v2.rs index 1d3f006ec66..27956aa8189 100644 --- a/rust_snuba/src/factory_v2.rs +++ b/rust_snuba/src/factory_v2.rs @@ -24,7 +24,8 @@ use crate::metrics::global_tags::set_global_tag; use crate::processors::{self, get_cogs_label}; use crate::strategies::accountant::RecordCogs; -use crate::strategies::clickhouse::writer_v2::{JsonWriterStep, RowBinaryWriterStep}; +use crate::strategies::clickhouse::streaming_writer::StreamingClickhouseWriter; +use crate::strategies::clickhouse::writer_v2::{ClickhouseClient, InsertFormat, JsonWriterStep}; use crate::strategies::commit_log::ProduceCommitLog; use crate::strategies::dlq_by_age::DlqByAge; use crate::strategies::healthcheck; @@ -163,63 +164,77 @@ impl ProcessingStrategyFactory for ConsumerStrategyFactoryV2 { Some(Duration::from_millis(self.join_timeout_ms.unwrap_or(0))), ); - // Pick the writer by wire format; RowBinary also needs the column list. + let compute_batch_size: fn(&BytesInsertBatch) -> usize = + match self.max_batch_size_calculation { + config::BatchSizeCalculation::Bytes => |batch| batch.num_bytes(), + config::BatchSizeCalculation::Rows => |batch| batch.len(), + }; + + // RowBinary streams compressed blocks onto an in-flight POST as + // rows arrive (no uncompressed Reduce buffer). JSON stays on the + // buffered writer — same pipeline we've shipped. let next_step: Box>> = if self.use_row_binary { + tracing::info!("Using streaming ClickHouse writer (RowBinary)"); let columns = insert_columns.expect("use_row_binary resolves a column list above"); - Box::new(RowBinaryWriterStep::new( + let client = Arc::new(ClickhouseClient::new( + &self.storage_config.clickhouse_cluster, + &self.storage_config.clickhouse_table_name, + self.storage_config.name.clone(), + InsertFormat::RowBinary, + Some(columns), + )); + Box::new(StreamingClickhouseWriter::new( next_step, - self.storage_config.clickhouse_cluster.clone(), - self.storage_config.clickhouse_table_name.clone(), + client, self.skip_write, - &self.clickhouse_concurrency, - self.storage_config.name.clone(), - columns, + self.clickhouse_concurrency.handle(), + self.clickhouse_concurrency.concurrency, + self.max_batch_size, + self.max_batch_time, + compute_batch_size, )) } else { - Box::new(JsonWriterStep::new( + let writer = JsonWriterStep::new( next_step, self.storage_config.clickhouse_cluster.clone(), self.storage_config.clickhouse_table_name.clone(), self.skip_write, &self.clickhouse_concurrency, self.storage_config.name.clone(), - )) - }; - - #[allow(clippy::result_large_err)] - let accumulator = Arc::new( - |batch: BytesInsertBatch, small_batch: Message>| { - Ok(batch.merge(small_batch.into_payload())) - }, - ); + ); - let compute_batch_size: fn(&BytesInsertBatch) -> usize = - match self.max_batch_size_calculation { - config::BatchSizeCalculation::Bytes => |batch| batch.num_bytes(), - config::BatchSizeCalculation::Rows => |batch| batch.len(), - }; - - let next_step = Reduce::new( - next_step, - accumulator, - Arc::new(move || { - BytesInsertBatch::::new( - RowData::default(), - None, - None, - None, - Default::default(), - CogsData::default(), + #[allow(clippy::result_large_err)] + let accumulator = Arc::new( + |batch: BytesInsertBatch, + small_batch: Message>| { + Ok(batch.merge(small_batch.into_payload())) + }, + ); + + Box::new( + Reduce::new( + writer, + accumulator, + Arc::new(move || { + BytesInsertBatch::::new( + RowData::default(), + None, + None, + None, + Default::default(), + CogsData::default(), + ) + }), + self.max_batch_size, + self.max_batch_time, + compute_batch_size, + // we need to enable this to deal with storages where we skip 100% of values. + // we still need to commit there + ) + .flush_empty_batches(true), ) - }), - self.max_batch_size, - self.max_batch_time, - compute_batch_size, - // we need to enable this to deal with storages where we skip 100% of values. - // we still need to commit there - ) - .flush_empty_batches(true); + }; // RowBinary can only be emitted by the Rust processor (the Python path // always returns JSONEachRow bytes). If the storage opted into diff --git a/rust_snuba/src/strategies/clickhouse/mod.rs b/rust_snuba/src/strategies/clickhouse/mod.rs index 45a3d852766..2a3a428333d 100644 --- a/rust_snuba/src/strategies/clickhouse/mod.rs +++ b/rust_snuba/src/strategies/clickhouse/mod.rs @@ -1,2 +1,4 @@ pub mod rowbinary; +pub mod streaming_lz4; +pub mod streaming_writer; pub mod writer_v2; diff --git a/rust_snuba/src/strategies/clickhouse/streaming_lz4.rs b/rust_snuba/src/strategies/clickhouse/streaming_lz4.rs new file mode 100644 index 00000000000..0900d6f1437 --- /dev/null +++ b/rust_snuba/src/strategies/clickhouse/streaming_lz4.rs @@ -0,0 +1,202 @@ +//! Incremental ClickHouse-native LZ4 compressor. +//! +//! Sister to `writer_v2::lz4_compress`. Both emit the same wire format +//! (concatenated CityHash128 + 9-byte header + LZ4 blocks chunked at +//! `LZ4_BLOCK_SIZE`), but this one accepts bytes incrementally and emits +//! a `Bytes` for each complete block as soon as it fills. Used by the +//! streaming writer so we never buffer a full uncompressed batch. + +use bytes::Bytes; + +use super::writer_v2::{ch_compression_checksum, LZ4_BLOCK_SIZE, LZ4_METHOD_BYTE}; + +pub struct StreamingLz4Compressor { + pending: Vec, +} + +impl StreamingLz4Compressor { + pub fn new() -> Self { + Self { + pending: Vec::with_capacity(LZ4_BLOCK_SIZE), + } + } + + /// Append `data` to the internal buffer. Returns one `Bytes` per + /// complete block that filled. Each returned chunk is a self-contained + /// ClickHouse native block ready to push onto an HTTP body stream. + pub fn push(&mut self, data: &[u8]) -> Vec { + self.pending.extend_from_slice(data); + let mut out = Vec::new(); + while self.pending.len() >= LZ4_BLOCK_SIZE { + let block: Vec = self.pending.drain(..LZ4_BLOCK_SIZE).collect(); + out.push(Bytes::from(encode_block(&block))); + } + out + } + + /// Emit any buffered remainder as a final (partial) block. Returns + /// `None` if the compressor saw no data — important because ClickHouse + /// rejects an empty native body, so callers must not append `None` + /// to the wire. + pub fn finish(mut self) -> Option { + if self.pending.is_empty() { + None + } else { + let last = std::mem::take(&mut self.pending); + Some(Bytes::from(encode_block(&last))) + } + } +} + +impl Default for StreamingLz4Compressor { + fn default() -> Self { + Self::new() + } +} + +/// Encode `chunk` as a single ClickHouse native compressed block: +/// `[16 B CityHash128 || 0x82 || u32 LE compressed-with-header || u32 LE +/// uncompressed || LZ4 raw bytes]`. The checksum spans the 9-byte header +/// and the compressed payload. +fn encode_block(chunk: &[u8]) -> Vec { + let compressed = lz4_flex::block::compress(chunk); + let compressed_with_header = 9u32 + compressed.len() as u32; + let uncompressed_size = chunk.len() as u32; + + let mut out = Vec::with_capacity(25 + compressed.len()); + out.extend_from_slice(&[0u8; 16]); + out.push(LZ4_METHOD_BYTE); + out.extend_from_slice(&compressed_with_header.to_le_bytes()); + out.extend_from_slice(&uncompressed_size.to_le_bytes()); + out.extend_from_slice(&compressed); + + let checksum = ch_compression_checksum(&out[16..]); + out[..16].copy_from_slice(&checksum); + out +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Walks a concatenation of ClickHouse-native compressed blocks and + /// returns the decoded payload. Mirrors the decoder in `writer_v2`'s + /// tests so the streaming output is verified against the same wire + /// expectations the buffered path is. + fn decode_native_blocks(buf: &[u8]) -> Vec { + let mut decoded = Vec::new(); + let mut pos = 0; + while pos < buf.len() { + assert!(buf.len() - pos >= 25, "truncated block header"); + let stored_checksum: [u8; 16] = buf[pos..pos + 16].try_into().unwrap(); + assert_eq!(buf[pos + 16], LZ4_METHOD_BYTE, "wrong compression method"); + let compressed_with_header = + u32::from_le_bytes(buf[pos + 17..pos + 21].try_into().unwrap()) as usize; + let uncompressed_size = + u32::from_le_bytes(buf[pos + 21..pos + 25].try_into().unwrap()) as usize; + let block_end = pos + 16 + compressed_with_header; + assert!(block_end <= buf.len(), "block size overruns buffer"); + + let computed = ch_compression_checksum(&buf[pos + 16..block_end]); + assert_eq!(computed, stored_checksum, "checksum mismatch"); + + let chunk = lz4_flex::block::decompress(&buf[pos + 25..block_end], uncompressed_size) + .expect("decompress"); + assert_eq!(chunk.len(), uncompressed_size); + decoded.extend_from_slice(&chunk); + pos = block_end; + } + decoded + } + + fn drain_to_bytes(chunks: Vec) -> Vec { + let mut out = Vec::new(); + for c in chunks { + out.extend_from_slice(&c); + } + out + } + + #[test] + fn push_under_block_size_buffers_until_finish() { + let mut c = StreamingLz4Compressor::new(); + assert!(c.push(b"hello, world").is_empty()); + let last = c.finish().expect("partial final block"); + assert_eq!(decode_native_blocks(&last), b"hello, world"); + } + + #[test] + fn empty_compressor_finishes_to_none() { + assert!(StreamingLz4Compressor::new().finish().is_none()); + } + + #[test] + fn exactly_one_full_block_emits_immediately_no_remainder() { + let input = vec![0xABu8; LZ4_BLOCK_SIZE]; + let mut c = StreamingLz4Compressor::new(); + let emitted = c.push(&input); + assert_eq!(emitted.len(), 1); + assert!(c.finish().is_none(), "should not have a partial tail"); + assert_eq!(decode_native_blocks(&drain_to_bytes(emitted)), input); + } + + #[test] + fn cross_boundary_push_splits_into_full_block_plus_remainder() { + let input: Vec = (0..(LZ4_BLOCK_SIZE + LZ4_BLOCK_SIZE / 2)) + .map(|i| (i % 251) as u8) + .collect(); + let mut c = StreamingLz4Compressor::new(); + let mut combined = drain_to_bytes(c.push(&input)); + combined.extend_from_slice(&c.finish().expect("half-block remainder")); + assert_eq!(decode_native_blocks(&combined), input); + } + + #[test] + fn many_small_pushes_recombine_at_block_boundary() { + // Stream 2.25 blocks worth of bytes one byte at a time. Two full + // blocks should emit during pushes (after byte LZ4_BLOCK_SIZE and + // byte 2*LZ4_BLOCK_SIZE), the rest comes out on finish. + let total = LZ4_BLOCK_SIZE * 2 + LZ4_BLOCK_SIZE / 4; + let mut c = StreamingLz4Compressor::new(); + let mut expected = Vec::with_capacity(total); + let mut wire = Vec::new(); + let mut full_blocks_seen = 0; + for i in 0..total { + let byte = (i % 251) as u8; + expected.push(byte); + let chunks = c.push(&[byte]); + full_blocks_seen += chunks.len(); + wire.extend_from_slice(&drain_to_bytes(chunks)); + } + assert_eq!( + full_blocks_seen, 2, + "should emit at each full-block boundary" + ); + wire.extend_from_slice(&c.finish().expect("trailing partial block")); + assert_eq!(decode_native_blocks(&wire), expected); + } + + #[test] + fn single_push_of_multiple_blocks_emits_all_at_once() { + // 3 full blocks in one push call — all three should emit; no remainder. + let input: Vec = (0..(LZ4_BLOCK_SIZE * 3)).map(|i| (i % 251) as u8).collect(); + let mut c = StreamingLz4Compressor::new(); + let emitted = c.push(&input); + assert_eq!(emitted.len(), 3); + assert!(c.finish().is_none()); + assert_eq!(decode_native_blocks(&drain_to_bytes(emitted)), input); + } + + #[test] + fn streaming_matches_buffered_lz4_compress() { + let input: Vec = (0..(LZ4_BLOCK_SIZE + 1234)) + .map(|i| (i % 251) as u8) + .collect(); + let mut c = StreamingLz4Compressor::new(); + let mut streamed = drain_to_bytes(c.push(&input)); + if let Some(last) = c.finish() { + streamed.extend_from_slice(&last); + } + assert_eq!(streamed, super::super::writer_v2::lz4_compress(&input)); + } +} diff --git a/rust_snuba/src/strategies/clickhouse/streaming_writer.rs b/rust_snuba/src/strategies/clickhouse/streaming_writer.rs new file mode 100644 index 00000000000..4c30c42964c --- /dev/null +++ b/rust_snuba/src/strategies/clickhouse/streaming_writer.rs @@ -0,0 +1,849 @@ +//! Streaming write path for ClickHouse INSERTs. +//! +//! Replaces the `Reduce → RowBinaryWriterStep` pair for the RowBinary +//! path. Encoded rows are LZ4-compressed as they arrive and pushed onto +//! an in-flight HTTP POST; the uncompressed bytes are dropped per +//! message, so we never hold a full uncompressed batch. +//! +//! The POST starts on the first non-empty payload of a batch (live +//! streaming) rather than at flush. Closing the body channel is what +//! ClickHouse treats as end-of-INSERT, so we only drop the sender in +//! `flush_pending`. Abandoning a batch aborts the HTTP task first so a +//! connection reset cannot land as a successful partial write. +//! +//! If the live attempt fails mid-batch, retries wait for +//! `body_complete` so they replay the full compressed body, not a +//! prefix. That was the truncation bug that forced #8001 to POST at +//! flush time. +//! +//! Inter-batch concurrency is `max_in_flight` (the same +//! `clickhouse_concurrency` the buffered writer uses). #8001 capped this +//! at one and the resulting submit backpressure showed up as consumer +//! pause time; we keep N slots and only reject once they are all busy. +//! In-flight batches complete in order so a later offset cannot commit +//! before an earlier write. + +use std::collections::{BTreeMap, VecDeque}; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use parking_lot::Mutex; +use reqwest::Response; +use sentry_arroyo::processing::strategies::{ + merge_commit_request, CommitRequest, MessageRejected, ProcessingStrategy, StrategyError, + SubmitError, +}; +use sentry_arroyo::types::{Message, Partition}; +use sentry_arroyo::utils::timing::Deadline; +use sentry_arroyo::{counter, gauge, timer}; +use tokio::runtime::Handle; +use tokio::sync::mpsc::UnboundedSender; +use tokio::task::JoinHandle; + +use super::streaming_lz4::StreamingLz4Compressor; +use super::writer_v2::ClickhouseClient; +use crate::types::{BytesInsertBatch, RowData}; + +struct PendingBatch { + batch_start: Deadline, + batch_size: usize, + num_rows: usize, + num_bytes: usize, + compressed_bytes: usize, + offsets: BTreeMap, + meta: BytesInsertBatch<()>, + write_start: Instant, + compressor: StreamingLz4Compressor, + chunks: std::sync::Arc>>, + tx: Option>, + complete_tx: Option>, + handle: Option>>, +} + +impl PendingBatch { + fn push_chunks(&mut self, chunks: Vec) { + if chunks.is_empty() { + return; + } + let mut buf = self.chunks.lock(); + for chunk in chunks { + self.compressed_bytes += chunk.len(); + if let Some(tx) = &self.tx { + let _ = tx.send(chunk.clone()); + } + buf.push(chunk); + } + } +} + +enum DrainMode { + NonBlocking, + BlockForever, + BlockUpTo(Duration), +} + +struct InFlightBatch { + handle: Option>>, + num_rows: usize, + num_bytes: usize, + compressed_bytes: usize, + offsets: BTreeMap, + meta: BytesInsertBatch<()>, + write_start: Instant, +} + +pub struct StreamingClickhouseWriter { + next_step: N, + client: std::sync::Arc, + skip_write: bool, + runtime: Handle, + max_batch_size: usize, + max_batch_time: Duration, + max_in_flight: usize, + compute_batch_size: fn(&BytesInsertBatch) -> usize, + + pending: Option, + in_flight: VecDeque, + + message_carried_over: Option>>, + commit_request_carried_over: Option, +} + +impl StreamingClickhouseWriter +where + N: ProcessingStrategy> + 'static, +{ + #[allow(clippy::too_many_arguments)] + pub fn new( + next_step: N, + client: std::sync::Arc, + skip_write: bool, + runtime: Handle, + max_in_flight: usize, + max_batch_size: usize, + max_batch_time: Duration, + compute_batch_size: fn(&BytesInsertBatch) -> usize, + ) -> Self { + StreamingClickhouseWriter { + next_step, + client, + skip_write, + runtime, + max_batch_size, + max_batch_time, + max_in_flight: max_in_flight.max(1), + compute_batch_size, + pending: None, + in_flight: VecDeque::new(), + message_carried_over: None, + commit_request_carried_over: None, + } + } + + fn ensure_pending(&mut self) -> &mut PendingBatch { + self.pending.get_or_insert_with(|| PendingBatch { + batch_start: Deadline::new(self.max_batch_time), + batch_size: 0, + num_rows: 0, + num_bytes: 0, + compressed_bytes: 0, + offsets: BTreeMap::new(), + meta: BytesInsertBatch::<()>::default(), + write_start: Instant::now(), + compressor: StreamingLz4Compressor::new(), + chunks: std::sync::Arc::new(Mutex::new(Vec::new())), + tx: None, + complete_tx: None, + handle: None, + }) + } + + /// Open the live HTTP POST for `pending`. No-op if we already have a + /// sender, we're skipping writes, or there is nothing to send. + fn start_stream(&mut self) { + if self.skip_write { + return; + } + let retry_buf = match &self.pending { + Some(pending) if pending.tx.is_none() && pending.handle.is_none() => { + pending.chunks.clone() + } + _ => return, + }; + + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let (complete_tx, complete_rx) = tokio::sync::oneshot::channel(); + let client = self.client.clone(); + let extra = self.max_batch_time; + let handle = self.runtime.spawn(async move { + let stream = futures::stream::unfold(rx, |mut rx| async move { + rx.recv() + .await + .map(|b| (Ok::(b), rx)) + }); + client + .send_streamed(stream, retry_buf, complete_rx, extra) + .await + }); + + let pending = self.pending.as_mut().expect("checked above"); + pending.tx = Some(tx); + pending.complete_tx = Some(complete_tx); + pending.handle = Some(handle); + } + + /// Finish the compressor, close the live body (clean EOF), signal + /// `body_complete` so retries may fire, and move the batch onto the + /// in-flight queue. + fn flush_pending(&mut self) { + let Some(pending) = self.pending.take() else { + return; + }; + let PendingBatch { + compressor, + chunks, + tx, + complete_tx, + handle, + mut compressed_bytes, + num_rows, + num_bytes, + offsets, + meta, + write_start, + .. + } = pending; + + if let Some(last) = compressor.finish() { + compressed_bytes += last.len(); + if let Some(ref tx) = tx { + let _ = tx.send(last.clone()); + } + chunks.lock().push(last); + } + + // Close the live body only here. Dropping `tx` is end-of-INSERT + // for ClickHouse; doing it earlier would ACK a truncated batch. + drop(tx); + if let Some(complete_tx) = complete_tx { + let _ = complete_tx.send(()); + } + + self.in_flight.push_back(InFlightBatch { + handle, + num_rows, + num_bytes, + compressed_bytes, + offsets, + meta, + write_start, + }); + } + + /// Abort the HTTP task *before* dropping the body sender, so the + /// connection resets instead of completing a partial INSERT. + fn abandon_pending(&mut self) { + let Some(mut pending) = self.pending.take() else { + return; + }; + if let Some(handle) = pending.handle.take() { + handle.abort(); + } + tracing::warn!( + "Abandoning pending streaming batch ({} rows, {} uncompressed bytes)", + pending.num_rows, + pending.num_bytes + ); + // `tx` / `complete_tx` drop after the abort: the request is + // already cancelled, so EOF cannot land as a successful write. + } + + fn try_drain_front(&mut self, mode: DrainMode) -> Result<(), StrategyError> { + let Some(mut in_flight) = self.in_flight.pop_front() else { + return Ok(()); + }; + match in_flight.handle.take() { + None => {} + Some(mut handle) => match mode { + DrainMode::NonBlocking => { + if !handle.is_finished() { + in_flight.handle = Some(handle); + self.in_flight.push_front(in_flight); + return Ok(()); + } + match self.runtime.block_on(&mut handle) { + Ok(Ok(_response)) => {} + Ok(Err(e)) => return Err(StrategyError::Other(e.into())), + Err(e) => return Err(StrategyError::Other(Box::new(e))), + } + } + DrainMode::BlockForever => match self.runtime.block_on(&mut handle) { + Ok(Ok(_response)) => {} + Ok(Err(e)) => return Err(StrategyError::Other(e.into())), + Err(e) => return Err(StrategyError::Other(Box::new(e))), + }, + DrainMode::BlockUpTo(max_wait) => { + let timeout_res = self + .runtime + .block_on(async { tokio::time::timeout(max_wait, &mut handle).await }); + match timeout_res { + Ok(Ok(Ok(_response))) => {} + Ok(Ok(Err(e))) => return Err(StrategyError::Other(e.into())), + Ok(Err(e)) => return Err(StrategyError::Other(Box::new(e))), + Err(_elapsed) => { + handle.abort(); + tracing::warn!( + "Streaming HTTP write exceeded {:?}; aborted in-flight task. \ + Batch metadata not committed downstream — the next consumer \ + instance will retry from the last committed offset.", + max_wait + ); + return Ok(()); + } + } + } + }, + } + + timer!("insertions.batch_write_ms", in_flight.write_start.elapsed()); + counter!("insertions.batch_write_bytes", in_flight.num_bytes as i64); + counter!("insertions.batch_write_msgs", in_flight.num_rows as i64); + gauge!( + "insertions.streaming_writer.last_batch_uncompressed_bytes", + in_flight.num_bytes as u64 + ); + gauge!( + "insertions.streaming_writer.last_batch_compressed_bytes", + in_flight.compressed_bytes as u64 + ); + in_flight.meta.record_message_latency(); + in_flight.meta.emit_item_type_metrics(); + tracing::info!("Inserted {} rows (streamed)", in_flight.num_rows); + + let message = Message::new_any_message(in_flight.meta, in_flight.offsets); + match self.next_step.submit(message) { + Ok(()) => Ok(()), + Err(SubmitError::MessageRejected(MessageRejected { message })) => { + self.message_carried_over = Some(message); + Ok(()) + } + Err(SubmitError::InvalidMessage(e)) => Err(e.into()), + } + } + + fn pending_ready_to_flush(&self) -> bool { + let Some(pending) = &self.pending else { + return false; + }; + pending.batch_size >= self.max_batch_size || pending.batch_start.has_elapsed() + } + + fn try_resubmit_carried_over(&mut self) -> Result<(), StrategyError> { + let Some(message) = self.message_carried_over.take() else { + return Ok(()); + }; + match self.next_step.submit(message) { + Ok(()) => Ok(()), + Err(SubmitError::MessageRejected(MessageRejected { message })) => { + self.message_carried_over = Some(message); + Ok(()) + } + Err(SubmitError::InvalidMessage(e)) => Err(e.into()), + } + } + + fn resident_compressed_bytes(&self) -> usize { + let pending = self.pending.as_ref().map_or(0, |p| p.compressed_bytes); + let in_flight = self + .in_flight + .iter() + .map(|b| b.compressed_bytes) + .sum::(); + pending + in_flight + } +} + +impl ProcessingStrategy> for StreamingClickhouseWriter +where + N: ProcessingStrategy> + 'static, +{ + fn poll(&mut self) -> Result, StrategyError> { + let commit_request = self.next_step.poll()?; + self.commit_request_carried_over = + merge_commit_request(self.commit_request_carried_over.take(), commit_request); + + self.try_resubmit_carried_over()?; + + while self.message_carried_over.is_none() && !self.in_flight.is_empty() { + let before = self.in_flight.len(); + self.try_drain_front(DrainMode::NonBlocking)?; + if self.in_flight.len() == before { + break; + } + } + + if self.message_carried_over.is_none() && self.pending_ready_to_flush() { + self.flush_pending(); + self.try_drain_front(DrainMode::NonBlocking)?; + } + + gauge!( + "insertions.streaming_writer.resident_compressed_bytes", + self.resident_compressed_bytes() as u64 + ); + gauge!( + "insertions.streaming_writer.in_flight", + self.in_flight.len() as u64 + ); + + Ok(self.commit_request_carried_over.take()) + } + + fn submit( + &mut self, + message: Message>, + ) -> Result<(), SubmitError>> { + if self.message_carried_over.is_some() + || (self.pending.is_none() && self.in_flight.len() >= self.max_in_flight) + { + return Err(SubmitError::MessageRejected(MessageRejected { message })); + } + + let commitables: Vec<(Partition, u64)> = message.committable().collect(); + let batch_size_inc = (self.compute_batch_size)(message.payload()); + let payload = message.into_payload(); + let (row_data, msg_meta) = payload.take(); + let RowData { + encoded_rows, + num_rows, + } = row_data; + let row_bytes_len = encoded_rows.len(); + let skip_write = self.skip_write; + + self.ensure_pending(); + if !encoded_rows.is_empty() && !skip_write { + self.start_stream(); + } + + let pending = self.pending.as_mut().expect("ensure_pending just ran"); + pending.batch_size += batch_size_inc; + pending.num_rows += num_rows; + pending.num_bytes += row_bytes_len; + for (partition, offset) in commitables { + pending.offsets.insert(partition, offset); + } + let prev_meta = std::mem::take(&mut pending.meta); + pending.meta = prev_meta.merge_meta(msg_meta); + + if !encoded_rows.is_empty() && !skip_write { + let new_chunks = pending.compressor.push(&encoded_rows); + pending.push_chunks(new_chunks); + } + + Ok(()) + } + + fn terminate(&mut self) { + self.abandon_pending(); + while let Some(in_flight) = self.in_flight.pop_front() { + if let Some(handle) = in_flight.handle { + handle.abort(); + } + } + self.next_step.terminate(); + } + + fn join(&mut self, timeout: Option) -> Result, StrategyError> { + let deadline = timeout.map(Deadline::new); + + if deadline.is_some_and(|d| d.has_elapsed()) { + self.abandon_pending(); + } else { + self.flush_pending(); + } + + while !self.in_flight.is_empty() || self.message_carried_over.is_some() { + if deadline.is_some_and(|d| d.has_elapsed()) { + while let Some(in_flight) = self.in_flight.pop_front() { + if let Some(handle) = in_flight.handle { + handle.abort(); + } + tracing::warn!( + "Timeout {:?} reached during streaming-writer join; aborted in-flight batch ({} rows)", + timeout, + in_flight.num_rows + ); + } + self.message_carried_over = None; + break; + } + + let commit = self.next_step.poll()?; + self.commit_request_carried_over = + merge_commit_request(self.commit_request_carried_over.take(), commit); + + self.try_resubmit_carried_over()?; + if self.message_carried_over.is_none() && !self.in_flight.is_empty() { + let mode = match deadline { + Some(d) => DrainMode::BlockUpTo(d.remaining()), + None => DrainMode::BlockForever, + }; + self.try_drain_front(mode)?; + } else if self.message_carried_over.is_some() { + std::thread::sleep(Duration::from_millis(10)); + } + } + + let next_commit = self.next_step.join(deadline.map(|d| d.remaining()))?; + Ok(merge_commit_request( + self.commit_request_carried_over.take(), + next_commit, + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use parking_lot::Mutex; + use sentry_arroyo::processing::strategies::{ + CommitRequest, MessageRejected, ProcessingStrategy, StrategyError, SubmitError, + }; + use sentry_arroyo::types::{BrokerMessage, InnerMessage, Message, Partition, Topic}; + use tokio::runtime::Handle; + + use crate::config::ClickhouseConfig; + use crate::strategies::clickhouse::writer_v2::{ClickhouseClient, InsertFormat}; + use crate::types::{BytesInsertBatch, RowData}; + + use super::*; + + struct RecordingStep { + batches: Arc>>>, + } + + impl ProcessingStrategy> for RecordingStep { + fn poll(&mut self) -> Result, StrategyError> { + Ok(None) + } + fn submit( + &mut self, + message: Message>, + ) -> Result<(), SubmitError>> { + self.batches.lock().push(message.into_payload()); + Ok(()) + } + fn terminate(&mut self) {} + fn join(&mut self, _: Option) -> Result, StrategyError> { + Ok(None) + } + } + + fn unreachable_client() -> std::sync::Arc { + std::sync::Arc::new(ClickhouseClient::new( + &ClickhouseConfig { + host: "127.0.0.1".to_string(), + port: 1, + secure: false, + user: "default".to_string(), + password: "".to_string(), + database: "default".to_string(), + }, + "test_table", + "test_storage".to_string(), + InsertFormat::RowBinary, + Some(&["col"]), + )) + } + + fn make_message( + payload: BytesInsertBatch, + partition: Partition, + offset: u64, + ) -> Message> { + Message { + inner_message: InnerMessage::BrokerMessage(BrokerMessage::new( + payload, + partition, + offset, + chrono::Utc::now(), + )), + } + } + + fn batch_with(rows: usize, bytes: usize) -> BytesInsertBatch { + BytesInsertBatch::::from_rows(RowData { + encoded_rows: vec![0u8; bytes], + num_rows: rows, + }) + .with_num_bytes(bytes) + } + + fn hung_in_flight(handle: Handle) -> InFlightBatch { + let stuck = handle.spawn(async { + tokio::time::sleep(Duration::from_secs(3600)).await; + anyhow::bail!("would have resolved if we ever got here") + }); + InFlightBatch { + handle: Some(stuck), + num_rows: 0, + num_bytes: 0, + compressed_bytes: 0, + offsets: BTreeMap::new(), + meta: BytesInsertBatch::<()>::default(), + write_start: Instant::now(), + } + } + + /// With `skip_write=true` and a row-based size limit of 2, three + /// messages produce one full flush + one trailing flush on join. + #[tokio::test] + async fn row_based_batching_skip_write() { + crate::testutils::initialize_python(); + let runtime = Handle::current(); + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + true, + runtime, + 2, + 2, + Duration::from_secs(3600), + |b| b.len(), + ); + + let partition = Partition::new(Topic::new("t"), 0); + + for i in 0..3 { + strategy + .submit(make_message(batch_with(1, 100), partition, i)) + .expect("submit should be accepted"); + let _ = strategy.poll(); + } + strategy + .join(Some(Duration::from_secs(5))) + .expect("join should not error"); + + let batches = recorded.lock(); + assert_eq!(batches.len(), 2, "size-flush + join-flush"); + assert_eq!(batches[0].num_bytes(), 200, "two messages of 100B"); + assert_eq!(batches[1].num_bytes(), 100, "trailing single message"); + } + + #[tokio::test] + async fn submit_rejected_while_downstream_backpressures() { + crate::testutils::initialize_python(); + let runtime = Handle::current(); + + struct AlwaysReject; + impl ProcessingStrategy> for AlwaysReject { + fn poll(&mut self) -> Result, StrategyError> { + Ok(None) + } + fn submit( + &mut self, + message: Message>, + ) -> Result<(), SubmitError>> { + Err(SubmitError::MessageRejected(MessageRejected { message })) + } + fn terminate(&mut self) {} + fn join( + &mut self, + _: Option, + ) -> Result, StrategyError> { + Ok(None) + } + } + + let mut strategy = StreamingClickhouseWriter::new( + AlwaysReject, + unreachable_client(), + true, + runtime, + 2, + 1, + Duration::from_secs(3600), + |b| b.len(), + ); + + let partition = Partition::new(Topic::new("t"), 0); + strategy + .submit(make_message(batch_with(1, 50), partition, 0)) + .expect("first submit should be accepted"); + let _ = strategy.poll(); + + let res = strategy.submit(make_message(batch_with(1, 50), partition, 1)); + assert!(matches!(res, Err(SubmitError::MessageRejected(_)))); + } + + #[test] + fn submit_accepted_when_a_slot_is_free() { + crate::testutils::initialize_python(); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .expect("build tokio runtime"); + let handle = runtime.handle().clone(); + + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + true, + handle.clone(), + 2, + 10, + Duration::from_secs(3600), + |b| b.len(), + ); + strategy.in_flight.push_back(hung_in_flight(handle)); + + let partition = Partition::new(Topic::new("t"), 0); + strategy + .submit(make_message(batch_with(1, 50), partition, 0)) + .expect("one busy slot of two must not reject submits"); + } + + #[test] + fn submit_rejected_when_all_slots_full() { + crate::testutils::initialize_python(); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .expect("build tokio runtime"); + let handle = runtime.handle().clone(); + + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + true, + handle.clone(), + 2, + 10, + Duration::from_secs(3600), + |b| b.len(), + ); + strategy.in_flight.push_back(hung_in_flight(handle.clone())); + strategy.in_flight.push_back(hung_in_flight(handle)); + + let partition = Partition::new(Topic::new("t"), 0); + let res = strategy.submit(make_message(batch_with(1, 50), partition, 0)); + assert!(matches!(res, Err(SubmitError::MessageRejected(_)))); + } + + #[test] + fn join_respects_timeout_with_hung_http() { + crate::testutils::initialize_python(); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .expect("build tokio runtime"); + let handle = runtime.handle().clone(); + + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + true, + handle.clone(), + 2, + 10, + Duration::from_secs(3600), + |b| b.len(), + ); + strategy.in_flight.push_back(hung_in_flight(handle)); + + let join_timeout = Duration::from_millis(200); + let started = std::time::Instant::now(); + strategy + .join(Some(join_timeout)) + .expect("join should not error on timeout"); + let elapsed = started.elapsed(); + + assert!( + elapsed < Duration::from_secs(2), + "join took {elapsed:?}; should have honored the {join_timeout:?} budget" + ); + assert!(recorded.lock().is_empty()); + } + + #[tokio::test] + async fn join_with_elapsed_deadline_abandons_pending() { + crate::testutils::initialize_python(); + let runtime = Handle::current(); + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + false, + runtime, + 2, + 10, + Duration::from_secs(3600), + |b| b.len(), + ); + + let partition = Partition::new(Topic::new("t"), 0); + strategy + .submit(make_message(batch_with(1, 50), partition, 0)) + .expect("submit should be accepted"); + + strategy + .join(Some(Duration::ZERO)) + .expect("join should not error"); + + assert!( + recorded.lock().is_empty(), + "elapsed-deadline join must not propagate the pending batch downstream" + ); + assert!(strategy.pending.is_none()); + assert!(strategy.in_flight.is_empty()); + } + + #[tokio::test] + async fn join_with_no_messages_returns_immediately() { + crate::testutils::initialize_python(); + let runtime = Handle::current(); + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + true, + runtime, + 2, + 10, + Duration::from_secs(3600), + |b| b.len(), + ); + + strategy + .join(Some(Duration::from_secs(1))) + .expect("join should not error"); + assert!(recorded.lock().is_empty()); + } +} diff --git a/rust_snuba/src/strategies/clickhouse/writer_v2.rs b/rust_snuba/src/strategies/clickhouse/writer_v2.rs index 118c7196a30..3e22eefbbdb 100644 --- a/rust_snuba/src/strategies/clickhouse/writer_v2.rs +++ b/rust_snuba/src/strategies/clickhouse/writer_v2.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; +use bytes::Bytes; use reqwest::header::{HeaderMap, HeaderValue, ACCEPT_ENCODING, CONNECTION}; use reqwest::{Client, Response}; use sentry_arroyo::processing::strategies::run_task_in_threads::{ @@ -195,13 +196,17 @@ where impl_writer_delegate!(JsonWriterStep); -/// Writer for the `RowBinary` wire format. `columns` is required: RowBinary is -/// positional, so the explicit column list maps wire order to the table's -/// columns (see `EAPItemRow::COLUMN_NAMES`). +/// Buffered writer for the `RowBinary` wire format. Production RowBinary +/// inserts go through [`super::streaming_writer::StreamingClickhouseWriter`]; +/// this remains as a fallback that POSTs a fully-buffered body. +/// `columns` is required: RowBinary is positional, so the explicit column +/// list maps wire order to the table's columns (see `EAPItemRow::COLUMN_NAMES`). +#[allow(dead_code)] pub struct RowBinaryWriterStep { inner: WriterInner, } +#[allow(dead_code)] impl RowBinaryWriterStep where N: ProcessingStrategy> + 'static, @@ -326,18 +331,34 @@ impl ClickhouseClient { async fn send_once( &self, - body: bytes::Bytes, + body: Bytes, attempt: usize, max_retries: usize, + ) -> Result { + self.send_once_body( + reqwest::Body::from(body), + attempt, + max_retries, + get_clickhouse_write_client_timeouts(&self.storage_name).request, + ) + .await + } + + async fn send_once_body( + &self, + body: reqwest::Body, + attempt: usize, + max_retries: usize, + timeout: Duration, ) -> Result { let started = Instant::now(); let res = self .client .post(self.build_url()) .headers(self.headers.clone()) - .timeout(get_clickhouse_write_client_timeouts(&self.storage_name).request) + .timeout(timeout) .query(&[("query", &self.query)]) - .body(reqwest::Body::from(body)) + .body(body) .send() .await; @@ -411,15 +432,109 @@ impl ClickhouseClient { ) }) } + + /// Stream `body_stream` as the POST body for the first attempt; on + /// failure, wait until `body_complete` fires (the caller has finished + /// the batch, so `retry_buf` is immutable) and retry from `retry_buf`. + /// + /// The caller MUST tee every chunk it pushes into `body_stream` into + /// `retry_buf` in the same order, and MUST only close the stream when + /// the entire batch has been pushed. Closing early lands a truncated + /// body that ClickHouse will accept as a complete INSERT. + /// + /// If the first attempt fails while the caller is still producing + /// chunks, retrying immediately would replay a partial `retry_buf`. + /// Waiting on `body_complete` is what makes live streaming safe. + /// Dropping the oneshot sender (batch abandoned) cancels retries. + /// + /// `first_attempt_extra` is added to the configured request timeout + /// for attempt 0 so the open connection can cover batch accumulation + /// plus the ClickHouse write. Retries use the plain request timeout: + /// the body is already complete. + pub async fn send_streamed( + &self, + body_stream: S, + retry_buf: Arc>>, + body_complete: tokio::sync::oneshot::Receiver<()>, + first_attempt_extra: Duration, + ) -> anyhow::Result + where + S: futures::stream::Stream> + Send + 'static, + { + let retry_policy = get_clickhouse_write_retry_policy(&self.storage_name); + let request_timeout = get_clickhouse_write_client_timeouts(&self.storage_name).request; + let first_timeout = request_timeout.saturating_add(first_attempt_extra); + + let first = self + .send_once_body( + reqwest::Body::wrap_stream(body_stream), + 0, + retry_policy.max_retries, + first_timeout, + ) + .await; + + match first { + Ok(response) => return Ok(response), + Err(failure) if retry_policy.max_retries == 0 => { + return Err(anyhow::anyhow!( + "error writing to clickhouse after 1 attempts: {}", + failure.detail + )); + } + Err(_) => { + if body_complete.await.is_err() { + anyhow::bail!( + "streamed insert cancelled before the body was complete; \ + not retrying a partial batch" + ); + } + } + } + + for attempt in 1..retry_policy.max_retries { + tokio::time::sleep(retry_policy.backoff(attempt - 1)).await; + let body = retry_body(&retry_buf); + if let Ok(response) = self + .send_once_body(body, attempt, retry_policy.max_retries, request_timeout) + .await + { + return Ok(response); + } + } + + tokio::time::sleep(retry_policy.backoff(retry_policy.max_retries.saturating_sub(1))).await; + self.send_once_body( + retry_body(&retry_buf), + retry_policy.max_retries, + retry_policy.max_retries, + request_timeout, + ) + .await + .map_err(|failure| { + anyhow::anyhow!( + "error writing to clickhouse after {} attempts: {}", + retry_policy.max_retries + 1, + failure.detail + ) + }) + } +} + +fn retry_body(retry_buf: &parking_lot::Mutex>) -> reqwest::Body { + let chunks: Vec = retry_buf.lock().clone(); + reqwest::Body::wrap_stream(futures::stream::iter( + chunks.into_iter().map(Ok::<_, std::io::Error>), + )) } /// ClickHouse native compressed-block size cap. Matches the server's /// `max_compress_block_size` default; sending larger blocks risks tripping /// server-side decompress limits. -const LZ4_BLOCK_SIZE: usize = 1024 * 1024; +pub(super) const LZ4_BLOCK_SIZE: usize = 1024 * 1024; /// ClickHouse compression method identifier for LZ4 in the native block header. -const LZ4_METHOD_BYTE: u8 = 0x82; +pub(super) const LZ4_METHOD_BYTE: u8 = 0x82; /// CityHash128 over `data` in the wire layout ClickHouse's /// `CompressedReadBuffer` reads: 8 little-endian bytes of the low 64-bit half @@ -435,7 +550,7 @@ const LZ4_METHOD_BYTE: u8 = 0x82; /// We use CityHash 1.0.2 — that's the variant ClickHouse bundles for /// compression checksums; the 110 variant is reserved for newer hash columns /// and is NOT interchangeable here. -fn ch_compression_checksum(data: &[u8]) -> [u8; 16] { +pub(super) fn ch_compression_checksum(data: &[u8]) -> [u8; 16] { cityhash_rs::cityhash_102_128(data) .rotate_left(64) .to_le_bytes() @@ -860,4 +975,131 @@ mod tests { ); } } + + fn streamed_test_client(port: u16, storage_name: &str) -> ClickhouseClient { + ClickhouseClient::new( + &ClickhouseConfig { + host: "127.0.0.1".to_string(), + port, + secure: false, + user: "default".to_string(), + password: "".to_string(), + database: "default".to_string(), + }, + "test_table", + storage_name.to_string(), + InsertFormat::RowBinary, + Some(&["col"]), + ) + } + + #[tokio::test] + async fn test_send_streamed_retries_from_buffered_chunks() { + crate::testutils::initialize_python(); + init_options(); + let _guard = override_options(&[( + "snuba", + "clickhouse_write_retry_policy", + json!({ + "streamed_retry_test": { + "initial_backoff_ms": 50.0, + "max_retries": 2, + "jitter_factor": 0.0 + } + }), + )]) + .unwrap(); + + let client = streamed_test_client(1, "streamed_retry_test"); + let retry_buf = Arc::new(parking_lot::Mutex::new(vec![ + Bytes::from_static(b"chunk-a"), + Bytes::from_static(b"chunk-b"), + ])); + let stream_chunks = retry_buf.lock().clone(); + let body_stream = + futures::stream::iter(stream_chunks.into_iter().map(Ok::<_, std::io::Error>)); + let (complete_tx, complete_rx) = tokio::sync::oneshot::channel(); + complete_tx.send(()).unwrap(); + + let start = Instant::now(); + let result = client + .send_streamed(body_stream, retry_buf, complete_rx, Duration::ZERO) + .await; + let elapsed = start.elapsed(); + + assert!(result.is_err()); + // Two backoffs of 50ms (attempts 1 and 2 after the live attempt). + assert!(elapsed >= Duration::from_millis(90)); + let err = result.unwrap_err().to_string(); + assert!(err.contains("after 3 attempts"), "got: {err}"); + } + + #[tokio::test] + async fn test_send_streamed_waits_for_body_complete_before_retry() { + crate::testutils::initialize_python(); + init_options(); + let _guard = override_options(&[( + "snuba", + "clickhouse_write_retry_policy", + json!({ + "streamed_wait_complete_test": { + "initial_backoff_ms": 10.0, + "max_retries": 1, + "jitter_factor": 0.0 + } + }), + )]) + .unwrap(); + + let client = streamed_test_client(1, "streamed_wait_complete_test"); + let retry_buf = Arc::new(parking_lot::Mutex::new(vec![Bytes::from_static(b"chunk")])); + // Connection is refused immediately, so the live attempt fails + // without consuming the (never-ending) stream. Retries must then + // wait for `body_complete` rather than replaying a partial buf. + let body_stream = futures::stream::pending::>(); + let (complete_tx, complete_rx) = tokio::sync::oneshot::channel(); + + let handle = tokio::spawn(async move { + client + .send_streamed(body_stream, retry_buf, complete_rx, Duration::ZERO) + .await + }); + + tokio::time::sleep(Duration::from_millis(150)).await; + assert!( + !handle.is_finished(), + "must not retry until the caller signals body complete" + ); + + complete_tx.send(()).unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), handle) + .await + .expect("retries should finish after body_complete") + .expect("task should not panic"); + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("after 2 attempts"), "got: {err}"); + } + + #[tokio::test] + async fn test_send_streamed_cancelled_if_body_complete_dropped() { + crate::testutils::initialize_python(); + init_options(); + let client = streamed_test_client(1, "streamed_cancel_test"); + let retry_buf = Arc::new(parking_lot::Mutex::new(vec![Bytes::from_static(b"chunk")])); + let body_stream = futures::stream::iter(std::iter::once(Ok::<_, std::io::Error>( + Bytes::from_static(b"chunk"), + ))); + let (complete_tx, complete_rx) = tokio::sync::oneshot::channel(); + drop(complete_tx); + + let result = client + .send_streamed(body_stream, retry_buf, complete_rx, Duration::ZERO) + .await; + let err = result.unwrap_err().to_string(); + assert!( + err.contains("cancelled before the body was complete"), + "got: {err}" + ); + } } diff --git a/rust_snuba/src/types.rs b/rust_snuba/src/types.rs index 1be51301905..714b96ff075 100644 --- a/rust_snuba/src/types.rs +++ b/rust_snuba/src/types.rs @@ -451,6 +451,24 @@ impl BytesInsertBatch { } } +impl BytesInsertBatch<()> { + /// Merge offsets, latency timers, cogs, item-type metrics and byte + /// counts without touching row payloads. Used by the streaming writer, + /// which drops encoded rows per message and only carries metadata + /// forward to commit. + pub fn merge_meta(mut self, other: BytesInsertBatch<()>) -> Self { + self.num_bytes += other.num_bytes; + self.commit_log_offsets.merge(other.commit_log_offsets); + self.message_timestamp.merge(other.message_timestamp); + self.origin_timestamp.merge(other.origin_timestamp); + self.sentry_received_timestamp + .merge(other.sentry_received_timestamp); + self.cogs_data.merge(other.cogs_data); + self.item_type_metrics.merge(other.item_type_metrics); + self + } +} + impl BytesInsertBatch { pub fn len(&self) -> usize { self.rows.num_rows From 417bd32454b825e513dac9ff38a57e210d417848 Mon Sep 17 00:00:00 2001 From: Pierre Massat Date: Thu, 10 Sep 2026 15:49:43 -0700 Subject: [PATCH 2/2] fix(eap): Flush streaming inserts when submit fills the batch RunTaskInThreads can deliver many completed messages in one cycle with no intervening poll. Close the INSERT body in submit once max_batch_size is reached so a single streamed write cannot grow unbounded. --- .../strategies/clickhouse/streaming_writer.rs | 74 ++++++++++++++++--- 1 file changed, 63 insertions(+), 11 deletions(-) diff --git a/rust_snuba/src/strategies/clickhouse/streaming_writer.rs b/rust_snuba/src/strategies/clickhouse/streaming_writer.rs index 4c30c42964c..932e9ff922d 100644 --- a/rust_snuba/src/strategies/clickhouse/streaming_writer.rs +++ b/rust_snuba/src/strategies/clickhouse/streaming_writer.rs @@ -425,19 +425,29 @@ where self.start_stream(); } - let pending = self.pending.as_mut().expect("ensure_pending just ran"); - pending.batch_size += batch_size_inc; - pending.num_rows += num_rows; - pending.num_bytes += row_bytes_len; - for (partition, offset) in commitables { - pending.offsets.insert(partition, offset); + { + let pending = self.pending.as_mut().expect("ensure_pending just ran"); + pending.batch_size += batch_size_inc; + pending.num_rows += num_rows; + pending.num_bytes += row_bytes_len; + for (partition, offset) in commitables { + pending.offsets.insert(partition, offset); + } + let prev_meta = std::mem::take(&mut pending.meta); + pending.meta = prev_meta.merge_meta(msg_meta); + + if !encoded_rows.is_empty() && !skip_write { + let new_chunks = pending.compressor.push(&encoded_rows); + pending.push_chunks(new_chunks); + } } - let prev_meta = std::mem::take(&mut pending.meta); - pending.meta = prev_meta.merge_meta(msg_meta); - if !encoded_rows.is_empty() && !skip_write { - let new_chunks = pending.compressor.push(&encoded_rows); - pending.push_chunks(new_chunks); + // Close the body as soon as the batch is full. RunTaskInThreads can + // deliver many completed messages in one poll cycle with no + // intervening `poll` here, so waiting until poll would let a single + // INSERT grow well past max_batch_size. + if self.pending_ready_to_flush() { + self.flush_pending(); } Ok(()) @@ -636,6 +646,48 @@ mod tests { assert_eq!(batches[1].num_bytes(), 100, "trailing single message"); } + /// RunTaskInThreads can `submit` many completed messages in one cycle + /// with no intervening `poll`. Size-based flush must happen in `submit` + /// so a single INSERT cannot grow past `max_batch_size`. + #[tokio::test] + async fn submit_flushes_full_batch_without_poll() { + crate::testutils::initialize_python(); + let runtime = Handle::current(); + let recorded = Arc::new(Mutex::new(Vec::new())); + let next_step = RecordingStep { + batches: recorded.clone(), + }; + let mut strategy = StreamingClickhouseWriter::new( + next_step, + unreachable_client(), + true, + runtime, + 2, + 2, + Duration::from_secs(3600), + |b| b.len(), + ); + + let partition = Partition::new(Topic::new("t"), 0); + for i in 0..4 { + strategy + .submit(make_message(batch_with(1, 100), partition, i)) + .expect("submit should be accepted"); + } + + assert!( + strategy.pending.is_none(), + "the 2nd and 4th submits should have flushed on size" + ); + assert_eq!(strategy.in_flight.len(), 2); + + let _ = strategy.poll(); + let batches = recorded.lock(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].num_bytes(), 200); + assert_eq!(batches[1].num_bytes(), 200); + } + #[tokio::test] async fn submit_rejected_while_downstream_backpressures() { crate::testutils::initialize_python();