fix(query): drain the result socket while rendering

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<BatchMessage> {
+        /// 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.

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<String> = 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<QueryResponse,
     Ok(response)
 }
 
+/// Batches in flight between the socket reader and the renderer.
+type BatchQueue = tokio::sync::mpsc::Receiver<crate::client::sdk::BatchMessage>;
+
+/// 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<Option<arrow::array::RecordBatch>, 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,16 +945,15 @@ 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<u64, ApiError> {
     writeln!(out, "{}", columns.join(",")).map_err(write_failure)?;
     let mut written = 0u64;
     let mut buf: Vec<u8> = 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<i64>,
     out: &mut impl Write,
 ) -> Result<u64, ApiError> {
-    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<u8> = 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<String> = 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<u8> = 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<u8> = 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<u8> = 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<u8> = 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]