From ae3dd45129bab16138c2d6d8bdf7596c73600d31 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 15 Sep 2026 17:03:21 +0530 Subject: [PATCH 1/3] fix(query): drain the result socket while rendering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 0.34.0 renders each batch on the thread that reads the socket, so the connection sits idle for the whole of every render. Locally that is invisible — rendering keeps up at 0.3us/row — but over a WAN link it is a slow consumer: throughput collapses against the receive window, and a connection held open for minutes is one an intermediary drops. Measured against a deployed server on a 20M-row result: 301s and `error decoding response body` after 1,966,081 of 20,000,001 rows, where the buffering client 0.34.0 replaced completes the same query in 41s. A reader on the runtime now drains the socket into a four-batch bounded queue and the renderer pulls from it, so reading never waits on rendering. Peak memory stays bounded by the queue rather than the result, which is the property the streaming change exists for. Same query after the change: 17.2s, 26MB, all 20,000,001 rows — 2.4x faster than the buffering client at 1/65th its memory. --- src/client/sdk.rs | 48 ++++++++++++++++++++++++++++++++++ src/commands/query.rs | 61 ++++++++++++++++++++++++++++++------------- 2 files changed, 91 insertions(+), 18 deletions(-) diff --git a/src/client/sdk.rs b/src/client/sdk.rs index 976ae08..94f09e2 100644 --- a/src/client/sdk.rs +++ b/src/client/sdk.rs @@ -957,6 +957,54 @@ impl Api { .map_err(ApiError::from_arrow) } + /// Drain a result stream on the runtime, handing batches back through a + /// bounded queue. + /// + /// Rendering a batch takes far longer than reading one, so a caller that + /// decodes and reads on the same thread leaves the socket idle for most of + /// the transfer. Over a WAN link that is a slow consumer: throughput + /// collapses against the receive window, and a connection held open for + /// minutes is one an intermediary is entitled to drop — measured at 20M + /// rows as a 7x slowdown ending in `error decoding response body`. + /// + /// Reading on the runtime while the caller renders keeps the socket drained + /// at full speed. `QUEUE_DEPTH` is what keeps that from being the buffering + /// this all exists to avoid: the reader blocks once the queue is full, so + /// resident batches are bounded by the queue rather than by the result. + pub fn stream_result_batches( + &self, + mut stream: hotdata::ArrowResultStream, + ) -> tokio::sync::mpsc::Receiver> { + /// Batches held between the reader and the renderer. Enough to cover a + /// render while the next reads, small enough that peak memory stays a + /// handful of batches. + const QUEUE_DEPTH: usize = 4; + + // A tokio channel, not `std::sync::mpsc`: the reader is async and its + // send must yield when the queue is full, not block a runtime worker. + let (tx, rx) = tokio::sync::mpsc::channel(QUEUE_DEPTH); + rt().spawn(async move { + loop { + match stream.next_batch().await { + Ok(Some(batch)) => { + // A send error means the renderer is gone — it failed, + // or stopped early. Stop reading rather than draining + // the rest of the result into a dropped queue. + if tx.send(Ok(batch)).await.is_err() { + return; + } + } + Ok(None) => return, + Err(e) => { + let _ = tx.send(Err(ApiError::from_arrow(e))).await; + return; + } + } + } + }); + rx + } + // --- Sample migrated call (workspace.rs uses this) ----------------------- /// List workspaces visible to the authenticated principal. diff --git a/src/commands/query.rs b/src/commands/query.rs index 6c74a79..0cdfc0c 100644 --- a/src/commands/query.rs +++ b/src/commands/query.rs @@ -759,7 +759,7 @@ pub(crate) fn print_streamed_result( eprintln!("warning: {warning}"); } - let mut stream = api + let stream = api .open_result_arrow_when_ready(&meta.result_id, None, RESULT_READY_TIMEOUT) .map_err(StreamFailure::early)?; let columns: Vec = stream @@ -777,9 +777,14 @@ pub(crate) fn print_streamed_result( // Past this point bytes are on their way to stdout, so a failure can no // longer be swapped for a fallback rendering — it is reported as one. + // Hand the stream to a reader on the runtime and render from its queue, so + // the socket keeps draining while a batch is being written. + let total_row_count = stream.total_row_count(); + let mut batches = api.stream_result_batches(stream); + let written = match format { - "csv" => stream_csv(api, &mut stream, &columns, &mut out), - "json" => stream_json(api, &mut stream, &columns, &meta, &mut out), + "csv" => stream_csv(&mut batches, &columns, &mut out), + "json" => stream_json(&mut batches, &columns, &meta, total_row_count, &mut out), _ => unreachable!("streamed formats are csv and json"), } .map_err(StreamFailure::mid_output)?; @@ -792,7 +797,7 @@ pub(crate) fn print_streamed_result( // The server reports the full result size independently of what was read, so // a body that ended early is caught even though the rows themselves looked // well-formed. - if let Some(total) = stream.total_row_count() { + if let Some(total) = total_row_count { let total = u64::try_from(total).unwrap_or(written); if written != total { eprintln!( @@ -864,6 +869,22 @@ fn fetch_capped(api: &Api, meta: &StreamMeta, cap: i64) -> Result>; + +/// Take the next batch, waiting for the reader. +/// +/// `blocking_recv` is correct here and only here: the renderer is a plain +/// synchronous loop, and the reader it waits on runs on the shared runtime. +fn next_queued(batches: &mut BatchQueue) -> Result, ApiError> { + match batches.blocking_recv() { + Some(Ok(batch)) => Ok(Some(batch)), + Some(Err(e)) => Err(e), + // The reader finished and dropped its sender. + None => Ok(None), + } +} + /// A stdout write that failed. Reported as a transport error so it exits the /// same way every other unrecoverable render failure does. fn write_failure(err: std::io::Error) -> ApiError { @@ -913,8 +934,7 @@ where /// Write a result as CSV, one batch at a time. Returns the rows written. fn stream_csv( - api: &Api, - stream: &mut hotdata::ArrowResultStream, + batches: &mut BatchQueue, columns: &[String], out: &mut impl Write, ) -> Result { @@ -922,7 +942,7 @@ fn stream_csv( let mut written = 0u64; let mut buf: Vec = Vec::new(); let mut line = String::new(); - while let Some(batch) = api.next_result_batch(stream)? { + while let Some(batch) = next_queued(batches)? { written += for_each_row(&batch, &mut buf, |cells| { line.clear(); for (i, cell) in cells.iter().enumerate() { @@ -957,13 +977,13 @@ fn csv_field(cell: &JsonCell) -> String { /// `serde_json::to_string_pretty` of the equivalent [`QueryResponse`] in /// `streamed_json_is_byte_identical_to_the_buffered_envelope`. fn stream_json( - api: &Api, - stream: &mut hotdata::ArrowResultStream, + batches: &mut BatchQueue, columns: &[String], meta: &StreamMeta, + reported_total: Option, out: &mut impl Write, ) -> Result { - let total_row_count = stream.total_row_count().and_then(|t| u64::try_from(t).ok()); + let total_row_count = reported_total.and_then(|t| u64::try_from(t).ok()); write!(out, "{{\n \"query_run_id\": ").map_err(write_failure)?; write_compact(out, &meta.query_run_id)?; @@ -975,7 +995,7 @@ fn stream_json( let mut written = 0u64; let mut buf: Vec = Vec::new(); - while let Some(batch) = api.next_result_batch(stream)? { + while let Some(batch) = next_queued(batches)? { for_each_row(&batch, &mut buf, |cells| { if written > 0 { write!(out, ",").map_err(write_failure)?; @@ -2062,7 +2082,7 @@ mod tests { let meta = meta_for("res_1"); // Streamed: batch by batch, off the socket. - let mut stream = api + let stream = api .open_result_arrow_when_ready("res_1", None, RESULT_READY_TIMEOUT) .unwrap(); let columns: Vec = stream @@ -2071,13 +2091,15 @@ mod tests { .iter() .map(|f| f.name().clone()) .collect(); + let reported_total = stream.total_row_count(); + let mut batches = api.stream_result_batches(stream); let mut streamed: Vec = Vec::new(); match format { "csv" => { - stream_csv(&api, &mut stream, &columns, &mut streamed).unwrap(); + stream_csv(&mut batches, &columns, &mut streamed).unwrap(); } "json" => { - stream_json(&api, &mut stream, &columns, &meta, &mut streamed).unwrap(); + stream_json(&mut batches, &columns, &meta, reported_total, &mut streamed).unwrap(); } _ => unreachable!(), } @@ -2154,19 +2176,22 @@ mod tests { let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); let meta = meta_for("res_1"); - let mut stream = api + let stream = api .open_result_arrow_when_ready("res_1", None, RESULT_READY_TIMEOUT) .unwrap(); let columns = vec!["id".to_string()]; + let mut batches = api.stream_result_batches(stream); let mut csv: Vec = Vec::new(); - stream_csv(&api, &mut stream, &columns, &mut csv).unwrap(); + stream_csv(&mut batches, &columns, &mut csv).unwrap(); assert_eq!(String::from_utf8(csv).unwrap(), "id\n"); - let mut stream = api + let stream = api .open_result_arrow_when_ready("res_1", None, RESULT_READY_TIMEOUT) .unwrap(); + let reported_total = stream.total_row_count(); + let mut batches = api.stream_result_batches(stream); let mut json: Vec = Vec::new(); - stream_json(&api, &mut stream, &columns, &meta, &mut json).unwrap(); + stream_json(&mut batches, &columns, &meta, reported_total, &mut json).unwrap(); let expected = QueryResponse { query_run_id: meta.query_run_id.clone(), result_id: Some("res_1".to_string()), From 4748e7e309290895146a288b85d85e9250b4658b Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 15 Sep 2026 17:10:11 +0530 Subject: [PATCH 2/3] fix(query): tell a dead reader apart from a finished one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reader task signalled the end of the stream by dropping its sender, which is also what a panic looks like: `spawn` stores the panic in a JoinHandle this code drops, so a reader that dies inside arrow decoding is invisible except for its silence. The renderer read that silence as the end of the result, and with no X-Total-Row-Count the row-count guard never fired — partial rows on stdout and exit 0. The stream now ends with an explicit `BatchMessage::Done`. A queue that closes without one is reported as an incomplete read. Covers the mid-stream failure path with a test that serves a body cut inside a batch and no row-count header, so nothing but the stream error can reveal the truncation. --- src/client/sdk.rs | 34 +++++++++++++++++++++---- src/commands/query.rs | 58 +++++++++++++++++++++++++++++++++++++++---- 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/src/client/sdk.rs b/src/client/sdk.rs index 94f09e2..b7b82ef 100644 --- a/src/client/sdk.rs +++ b/src/client/sdk.rs @@ -483,6 +483,19 @@ async fn apply_seam_headers( req } +/// What the result reader hands to the renderer. +/// +/// `Done` is explicit because a dropped sender is ambiguous: it is what a clean +/// finish looks like *and* what a panicked reader looks like, and reading the +/// second as the first turns a truncated download into a short result reported +/// as complete. +pub enum BatchMessage { + Batch(arrow::array::RecordBatch), + /// The stream ended cleanly. Nothing follows. + Done, + Failed(ApiError), +} + /// How often to re-check a result the server is still writing. const RESULT_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(500); /// Upper bound on a single wait, so a generous `Retry-After` cannot stall the @@ -969,12 +982,17 @@ impl Api { /// /// Reading on the runtime while the caller renders keeps the socket drained /// at full speed. `QUEUE_DEPTH` is what keeps that from being the buffering - /// this all exists to avoid: the reader blocks once the queue is full, so + /// this all exists to avoid: the reader waits once the queue is full, so /// resident batches are bounded by the queue rather than by the result. + /// + /// The stream ends with an explicit [`BatchMessage::Done`]. A closed queue + /// without one means the reader stopped without finishing — a panic inside + /// decoding, which `spawn` swallows into a dropped `JoinHandle` — and the + /// consumer reports that rather than reading it as the end of the result. pub fn stream_result_batches( &self, mut stream: hotdata::ArrowResultStream, - ) -> tokio::sync::mpsc::Receiver> { + ) -> tokio::sync::mpsc::Receiver { /// Batches held between the reader and the renderer. Enough to cover a /// render while the next reads, small enough that peak memory stays a /// handful of batches. @@ -990,13 +1008,19 @@ impl Api { // A send error means the renderer is gone — it failed, // or stopped early. Stop reading rather than draining // the rest of the result into a dropped queue. - if tx.send(Ok(batch)).await.is_err() { + if tx.send(BatchMessage::Batch(batch)).await.is_err() { return; } } - Ok(None) => return, + // Say so, rather than letting the dropped sender say it: a + // closed queue alone cannot be told apart from this task + // dying mid-decode. + Ok(None) => { + let _ = tx.send(BatchMessage::Done).await; + return; + } Err(e) => { - let _ = tx.send(Err(ApiError::from_arrow(e))).await; + let _ = tx.send(BatchMessage::Failed(ApiError::from_arrow(e))).await; return; } } diff --git a/src/commands/query.rs b/src/commands/query.rs index 0cdfc0c..597aed4 100644 --- a/src/commands/query.rs +++ b/src/commands/query.rs @@ -870,18 +870,29 @@ fn fetch_capped(api: &Api, meta: &StreamMeta, cap: i64) -> Result>; +type BatchQueue = tokio::sync::mpsc::Receiver; /// Take the next batch, waiting for the reader. /// /// `blocking_recv` is correct here and only here: the renderer is a plain /// synchronous loop, and the reader it waits on runs on the shared runtime. +/// +/// A queue that closes without [`BatchMessage::Done`] is an error, not an end. +/// `spawn` stores a task panic in a `JoinHandle` this code drops, so a reader +/// that dies mid-decode is invisible except for its silence — and reading that +/// silence as "the result ended" is how a truncated download becomes a short +/// result reported as complete. fn next_queued(batches: &mut BatchQueue) -> Result, ApiError> { + use crate::client::sdk::BatchMessage; match batches.blocking_recv() { - Some(Ok(batch)) => Ok(Some(batch)), - Some(Err(e)) => Err(e), - // The reader finished and dropped its sender. - None => Ok(None), + Some(BatchMessage::Batch(batch)) => Ok(Some(batch)), + Some(BatchMessage::Done) => Ok(None), + Some(BatchMessage::Failed(e)) => Err(e), + None => Err(ApiError::Transport( + "the result reader stopped before the end of the stream; the rows \ + read so far are incomplete" + .to_string(), + )), } } @@ -2336,6 +2347,43 @@ mod tests { assert_eq!(result_exit_code(&result), 0); } + /// A body that dies mid-stream must surface as an error, not as the end of + /// the result. + /// + /// This is the path that failed in production: the reader gets + /// `error decoding response body` partway through and the rows already + /// written are a fraction of the result. The error travels through the + /// queue as `BatchMessage::Failed`, so the renderer stops and the process + /// exits non-zero rather than reporting a short result as complete. + #[test] + fn a_body_that_dies_mid_stream_is_an_error_not_an_ending() { + let ipc = awkward_ipc(); + // Cut inside the second batch: the schema and first batch decode, then + // the stream ends where a message should continue. + let truncated = ipc[..ipc.len() - 40].to_vec(); + + let mut server = mockito::Server::new(); + // No X-Total-Row-Count, so nothing but the stream error can reveal the + // truncation — the row-count guard cannot fire. + let _m = serve_ipc(&mut server, truncated, None); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); + + let stream = api + .open_result_arrow_when_ready("res_1", None, RESULT_READY_TIMEOUT) + .expect("the schema is intact, so opening succeeds"); + let columns = vec!["id".to_string()]; + let mut batches = api.stream_result_batches(stream); + let mut out: Vec = Vec::new(); + + let err = stream_csv(&mut batches, &columns, &mut out) + .expect_err("a body that ends mid-message must not read as a clean end of stream"); + let msg = err.message(); + assert!( + !msg.is_empty(), + "the failure has to name something a caller can act on" + ); + } + /// The row count the server advertises is checked against what was read, so /// a body that ends early is reported rather than passed off as complete. #[test] From 3b532367e6aba4516ddeefe6c163c22b2a261ca6 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 15 Sep 2026 17:21:01 +0530 Subject: [PATCH 3/3] test(query): assert the decode failure, not just any failure --- src/commands/query.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/commands/query.rs b/src/commands/query.rs index 597aed4..e809645 100644 --- a/src/commands/query.rs +++ b/src/commands/query.rs @@ -2378,9 +2378,19 @@ mod tests { let err = stream_csv(&mut batches, &columns, &mut out) .expect_err("a body that ends mid-message must not read as a clean end of stream"); let msg = err.message(); + // Both error paths carry a message, so an emptiness check would pass on + // either. This has to be the decode failure travelling as + // `BatchMessage::Failed` — if the reader panicked instead, the queue + // would close with no marker and `next_queued` would produce its own + // reader-stopped error, which is a different bug wearing the same exit + // code. assert!( - !msg.is_empty(), - "the failure has to name something a caller can act on" + !msg.contains("the result reader stopped"), + "the reader died instead of reporting the decode failure: {msg}" + ); + assert!( + msg.to_lowercase().contains("arrow") || msg.to_lowercase().contains("ipc"), + "the failure should name the decode that failed, got {msg}" ); }