diff --git a/src/client/sdk.rs b/src/client/sdk.rs index 976ae08..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 @@ -957,6 +970,65 @@ 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 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 { + /// 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(BatchMessage::Batch(batch)).await.is_err() { + 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(BatchMessage::Failed(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..e809645 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,33 @@ 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. +/// +/// 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(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(), + )), + } +} + /// 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 +945,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 +953,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 +988,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 +1006,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 +2093,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 +2102,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 +2187,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()), @@ -2311,6 +2347,53 @@ 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(); + // 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.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}" + ); + } + /// 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]