perf(eap): Stream RowBinary INSERTs instead of buffering

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"] }

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<KafkaPayload> 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<RowData>) -> 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<dyn ProcessingStrategy<BytesInsertBatch<RowData>>> =
             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<RowData>, small_batch: Message<BytesInsertBatch<RowData>>| {
-                Ok(batch.merge(small_batch.into_payload()))
-            },
-        );
+                );
 
-        let compute_batch_size: fn(&BytesInsertBatch<RowData>) -> 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::<RowData>::new(
-                    RowData::default(),
-                    None,
-                    None,
-                    None,
-                    Default::default(),
-                    CogsData::default(),
+                #[allow(clippy::result_large_err)]
+                let accumulator = Arc::new(
+                    |batch: BytesInsertBatch<RowData>,
+                     small_batch: Message<BytesInsertBatch<RowData>>| {
+                        Ok(batch.merge(small_batch.into_payload()))
+                    },
+                );
+
+                Box::new(
+                    Reduce::new(
+                        writer,
+                        accumulator,
+                        Arc::new(move || {
+                            BytesInsertBatch::<RowData>::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

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;