perf(blobs): stream downloads, single-scan, and in-flight memory budget

crates/lance-context-core/src/rollout_store.rs

@@ -1132,6 +1132,42 @@ impl RolloutStore {
         }
     }
 
+    /// Fetch a row *together with* its `binary_payload` in a single base-first
+    /// scan, returning `(record, payload)`.
+    ///
+    /// Callers that need both the row metadata (e.g. `content_type` for a
+    /// download's `Content-Type`/filename) and the artifact bytes would
+    /// otherwise call [`Self::get_by_id`] then [`Self::get_blob`] — two
+    /// independent point scans over the same shard. This folds them into one
+    /// scan: the row is located once and `binary_payload` is materialized for
+    /// only that row (via a projected `take`), so it never reads back an entire
+    /// fragment's payloads. `payload` is `None` when the row carries no blob.
+    ///
+    /// Base-table-first with the same immutable-row reasoning as
+    /// [`Self::get_by_id`]: a hit in the base table returns immediately with no
+    /// MemWAL generation opened; only a base miss falls back to the flushed WAL.
+    pub async fn get_record_with_blob(
+        &self,
+        id: &str,
+    ) -> LanceResult<Option<(RolloutRecord, Option<Vec<u8>>)>> {
+        // Base table first — no manifest reads, no per-generation opens.
+        if let Some(record) = self.scan_one_by_id(id, ListSource::Fragments).await? {
+            let payload = Self::get_blob_from_dataset(&self.dataset, id)
+                .await?
+                .flatten();
+            return Ok(Some((record, payload)));
+        }
+
+        // Base miss: locate the row in the flushed generations. Reuse the
+        // NotFound-tolerant, bounded-parallel blob fallback and pair it with the
+        // WAL-sourced metadata scan.
+        let Some(record) = self.scan_one_by_id(id, ListSource::Wal).await? else {
+            return Ok(None);
+        };
+        let payload = self.get_blob(id).await?;
+        Ok(Some((record, payload)))
+    }
+
     /// Run an id-equality point scan against a single [`ListSource`] and return
     /// the first matching record. `Fragments` passes an empty snapshot set so it
     /// performs no MemWAL manifest discovery; `Wal`/`All` discover flushed
@@ -3832,4 +3868,54 @@ mod tests {
             );
         });
     }
+
+    #[test]
+    fn get_record_with_blob_returns_row_and_payload_in_one_scan() {
+        // get_record_with_blob folds the metadata point lookup and the payload
+        // fetch into a single base-first scan. It must return both the row and
+        // its bytes for an un-merged WAL row (base-miss -> fallback) and for a
+        // merged base row, and None for a missing id.
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let bytes = b"\x00\x01record-with-blob";
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+        runtime.block_on(async {
+            let mut store = RolloutStore::open_with_options(
+                &uri,
+                RolloutStoreOptions {
+                    storage_options: None,
+                    shard_id: Some("rollout-0".to_string()),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+            let artifact = artifact_record("row-rw", bytes);
+            store.add(std::slice::from_ref(&artifact)).await.unwrap();
+
+            // Un-merged: found via the WAL fallback, row + payload paired.
+            assert_eq!(flushed_generation_count(&store).await, 1);
+            let (record, payload) = store
+                .get_record_with_blob("row-rw")
+                .await
+                .unwrap()
+                .expect("row present in WAL");
+            assert_records_eq(&record, &artifact);
+            assert_eq!(payload.as_deref(), Some(&bytes[..]));
+
+            // After merge: found via the base table, still paired.
+            store.cleanup_own_shard().await.unwrap();
+            assert_eq!(flushed_generation_count(&store).await, 0);
+            let (record, payload) = store
+                .get_record_with_blob("row-rw")
+                .await
+                .unwrap()
+                .expect("row present in base");
+            assert_records_eq(&record, &artifact);
+            assert_eq!(payload.as_deref(), Some(&bytes[..]));
+
+            // Missing id: None, not an error.
+            assert!(store.get_record_with_blob("nope").await.unwrap().is_none());
+        });
+    }
 }

crates/lance-context-master/Cargo.toml

@@ -19,6 +19,7 @@ lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" }
 arrow-array = "58"
 arrow-schema = "58"
 axum = { version = "0.8", features = ["json"] }
+bytes = "1"
 chrono = { version = "0.4", default-features = false, features = ["clock"] }
 clap = { version = "4", features = ["derive", "env"] }
 etcd-client = { version = "0.19", features = ["tls"] }

crates/lance-context-master/src/routes.rs

@@ -209,16 +209,16 @@ pub async fn download_experiment_blob(
         .refresh_latest()
         .await
         .map_err(MasterError::from_lance)?;
-    let record = store
-        .get_by_id(&id)
+    // Single base-first scan returns the row metadata and its payload together,
+    // instead of a separate get_by_id + get_blob (two point scans over the same
+    // shard) as before.
+    let (record, payload) = store
+        .get_record_with_blob(&id)
         .await
         .map_err(MasterError::from_lance)?
         .ok_or_else(|| MasterError::NotFound(format!("record '{}' does not exist", id)))?;
-    let bytes = store
-        .get_blob(&id)
-        .await
-        .map_err(MasterError::from_lance)?
-        .ok_or_else(|| MasterError::NotFound(format!("record '{}' has no blob", id)))?;
+    let bytes =
+        payload.ok_or_else(|| MasterError::NotFound(format!("record '{}' has no blob", id)))?;
 
     let content_type = HeaderValue::from_str(&record.content_type)
         .unwrap_or_else(|_| HeaderValue::from_static("application/octet-stream"));
@@ -229,10 +229,29 @@ pub async fn download_experiment_blob(
     Response::builder()
         .header(header::CONTENT_TYPE, content_type)
         .header(header::CONTENT_DISPOSITION, disposition)
-        .body(Body::from(bytes))
+        .header(header::CONTENT_LENGTH, bytes.len())
+        .body(blob_stream_body(bytes))
         .map_err(|err| MasterError::Internal(err.to_string()))
 }
 
+/// Frame size for chunked blob download responses (see [`blob_stream_body`]).
+const BLOB_STREAM_CHUNK_BYTES: usize = 256 * 1024;
+
+/// Build a chunked [`Body`] from an owned blob so the HTTP send path never holds
+/// a second full-blob copy and a slow client applies backpressure at frame
+/// granularity. `Bytes` frames are refcounted slices of one allocation, so no
+/// per-frame copy of the payload is made.
+fn blob_stream_body(bytes: Vec<u8>) -> Body {
+    let buf = bytes::Bytes::from(bytes);
+    let chunks = (0..buf.len())
+        .step_by(BLOB_STREAM_CHUNK_BYTES.max(1))
+        .map(move |start| {
+            let end = (start + BLOB_STREAM_CHUNK_BYTES).min(buf.len());
+            Ok::<_, std::convert::Infallible>(buf.slice(start..end))
+        });
+    Body::from_stream(futures::stream::iter(chunks))
+}
+
 /// `POST /api/v1/rescan` — trigger one immediate full scan.
 pub async fn rescan(
     State(state): State<Arc<MasterState>>,

crates/lance-context-server/Cargo.toml

@@ -17,6 +17,8 @@ lance-context-core = { version = "0.6.3", path = "../lance-context-core" }
 lance-context-api = { version = "0.6.3", path = "../lance-context-api" }
 lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" }
 axum = { version = "0.8", features = ["json", "multipart"] }
+bytes = "1"
+futures = "0.3"
 metrics = "0.24"
 chrono = { version = "0.4", default-features = false, features = ["clock"] }
 clap = { version = "4", features = ["derive", "env"] }

crates/lance-context-server/src/config.rs

@@ -53,6 +53,22 @@ pub struct ServerConfig {
     /// back to the built-in default.
     #[arg(long, env = "ROLLOUT_CACHE_CAPACITY", default_value = "2000")]
     pub rollout_cache_capacity: usize,
+
+    /// Ceiling, in bytes, on the total artifact-blob payload held in memory
+    /// across all *concurrent* rollout uploads and downloads on this instance.
+    ///
+    /// Each blob request materializes its full payload as an in-memory buffer
+    /// (uploads buffer the request body; downloads materialize the row's
+    /// `binary_payload`). Without a global cap, N concurrent 1 GiB requests
+    /// would need N GiB and can OOM the process. This budget admits a request
+    /// only while enough of the budget is free — otherwise it is rejected with
+    /// `503 Service Unavailable` and a `Retry-After`, applying backpressure at
+    /// the edge instead of the allocator. A single request larger than the whole
+    /// budget is still admitted when the instance is otherwise idle (it reserves
+    /// the entire budget), so this bounds concurrency, not maximum blob size.
+    /// `0` (the default) disables the budget.
+    #[arg(long, env = "ROLLOUT_MAX_INFLIGHT_BLOB_BYTES", default_value = "0")]
+    pub rollout_max_inflight_blob_bytes: usize,
 }
 
 impl ServerConfig {

crates/lance-context-server/src/error.rs

@@ -11,6 +11,9 @@ pub enum AppError {
     InvalidRequest(String),
     Internal(String),
     CompactionInProgress,
+    /// The in-flight blob-byte budget is exhausted; the client should retry
+    /// later. Maps to `503 Service Unavailable`.
+    Overloaded(String),
 }
 
 impl AppError {
@@ -55,6 +58,7 @@ impl IntoResponse for AppError {
                 "COMPACTION_IN_PROGRESS",
                 "Compaction already in progress".to_string(),
             ),
+            AppError::Overloaded(msg) => (StatusCode::SERVICE_UNAVAILABLE, "OVERLOADED", msg),
         };
 
         let body = ErrorResponse {

crates/lance-context-server/src/routes/rollouts.rs

@@ -27,6 +27,71 @@ use crate::state::AppState;
 /// raises the ceiling well above it while still bounding memory.
 pub const MAX_ROLLOUT_UPLOAD_BYTES: usize = 1024 * 1024 * 1024;
 
+/// Frame size for chunked blob download responses. The in-memory `Vec<u8>` is
+/// sliced into frames of this size so the HTTP send path never holds an extra
+/// full-blob copy and a slow reader exerts backpressure at frame granularity
+/// rather than after the whole payload is queued. 256 KiB keeps per-frame
+/// overhead negligible while capping the amount buffered ahead of the socket.
+const BLOB_STREAM_CHUNK_BYTES: usize = 256 * 1024;
+
+/// Build a chunked [`Body`] from an owned blob. `Bytes` frames share one
+/// backing allocation (cheap refcounted slices — no per-frame copy), so this
+/// does not duplicate the payload; it only changes how it is fed to the socket.
+///
+/// `reservation` (the in-flight blob-budget guard, if any) is moved into the
+/// stream and released only when the last frame has been produced, so the
+/// budget accounts for a slow download for its full lifetime.
+fn blob_stream_body(
+    bytes: Vec<u8>,
+    mut reservation: Option<crate::state::BlobReservation>,
+) -> Body {
+    let buf = bytes::Bytes::from(bytes);
+    let len = buf.len();
+    let mut offset = 0usize;
+    let stream = futures::stream::poll_fn(move |_| {
+        if offset >= len {
+            // Drop the reservation exactly when the stream ends.
+            let _ = reservation.take();
+            return std::task::Poll::Ready(None);
+        }
+        let end = (offset + BLOB_STREAM_CHUNK_BYTES).min(len);
+        let frame = buf.slice(offset..end);
+        offset = end;
+        std::task::Poll::Ready(Some(Ok::<_, std::convert::Infallible>(frame)))
+    });
+    Body::from_stream(stream)
+}
+
+/// Parse the `Content-Length` header into a byte count, defaulting to `0` when
+/// absent or unparsable (a chunked upload without a declared length reserves
+/// nothing up-front; the body-size limit still caps it).
+fn content_length(headers: &header::HeaderMap) -> usize {
+    headers
+        .get(header::CONTENT_LENGTH)
+        .and_then(|value| value.to_str().ok())
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(0)
+}
+
+/// Reserve `bytes` from the instance's in-flight blob budget (if configured),
+/// returning the RAII guard to hold for the request's duration. Returns
+/// `503 Overloaded` when the budget cannot currently admit the request. When no
+/// budget is configured (`None`) this is a no-op that always admits.
+fn acquire_blob_budget(
+    state: &AppState,
+    bytes: usize,
+) -> Result<Option<crate::state::BlobReservation>, AppError> {
+    match &state.blob_budget {
+        None => Ok(None),
+        Some(budget) => budget.try_acquire(bytes).map(Some).ok_or_else(|| {
+            metrics::counter!("rollout_blob_budget_rejections_total").increment(1);
+            AppError::Overloaded(
+                "server is at its in-flight blob memory limit; retry shortly".to_string(),
+            )
+        }),
+    }
+}
+
 /// Response for the internal WAL-merge trigger: how many flushed generations
 /// this worker's shard folded into the base table.
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
@@ -170,6 +235,12 @@ pub async fn add_rollouts(
         .unwrap_or("")
         .to_string();
 
+    // Admit against the in-flight blob budget before buffering the body, using
+    // the declared Content-Length as the reservation size. Held for the whole
+    // handler so concurrent uploads cannot collectively exceed the budget and
+    // OOM the worker; dropped when the request completes.
+    let _budget = acquire_blob_budget(&state, content_length(req.headers()))?;
+
     let records = if content_type_is(&content_type, "multipart/form-data") {
         let multipart = Multipart::from_request(req, &state)
             .await
@@ -364,6 +435,11 @@ pub async fn get_rollout(
 /// Materialize a rollout row's `binary_payload` bytes. The artifact
 /// bytes are opaque, so they stream as `application/octet-stream`. `404` when
 /// the row is absent or carries no payload.
+///
+/// The payload is sent as a **chunked** body ([`blob_stream_body`]) rather than
+/// a single giant frame: the response yields fixed-size frames so the HTTP
+/// layer never re-buffers the whole blob and a slow client applies backpressure
+/// instead of forcing the server to hold an extra full copy in the send queue.
 pub async fn fetch_rollout_blob(
     State(state): State<Arc<AppState>>,
     Path((name, id)): Path<(String, String)>,
@@ -377,9 +453,18 @@ pub async fn fetch_rollout_blob(
         .map_err(AppError::from_lance)?
         .ok_or_else(|| AppError::NotFound(format!("Rollout '{}' has no payload", id)))?;
 
+    // Reserve now that the payload size is known, and hold the reservation for
+    // the whole streamed send (moved into the body): a slow client keeps the
+    // blob resident until the last frame flushes, so the budget must account for
+    // it until then. Reject with 503 if the budget is currently exhausted.
+    let reservation = acquire_blob_budget(&state, bytes.len())?;
+    drop(store);
+
+    let len = bytes.len();
     Response::builder()
         .header(header::CONTENT_TYPE, "application/octet-stream")
-        .body(Body::from(bytes))
+        .header(header::CONTENT_LENGTH, len)
+        .body(blob_stream_body(bytes, reservation))
         .map_err(|err| AppError::Internal(err.to_string()))
 }
 
@@ -620,6 +705,36 @@ mod tests {
     use super::*;
     use crate::state::AppState;
 
+    #[tokio::test]
+    async fn blob_stream_body_reassembles_across_chunk_boundaries() {
+        // A payload spanning several BLOB_STREAM_CHUNK_BYTES frames must
+        // reassemble byte-for-byte, and a reservation moved into the body is
+        // held until the stream is fully drained.
+        let payload: Vec<u8> = (0..(BLOB_STREAM_CHUNK_BYTES * 2 + 123))
+            .map(|i| (i % 251) as u8)
+            .collect();
+        let budget = crate::state::BlobBudget::new(payload.len());
+        let reservation = budget.try_acquire(payload.len());
+        assert!(reservation.is_some());
+        // Budget is now fully occupied.
+        assert!(budget.try_acquire(1).is_none());
+
+        let body = blob_stream_body(payload.clone(), reservation);
+        let collected = axum::body::to_bytes(body, usize::MAX).await.unwrap();
+        assert_eq!(&collected[..], &payload[..]);
+
+        // Once the body is consumed the reservation dropped, freeing the budget.
+        assert!(budget.try_acquire(payload.len()).is_some());
+    }
+
+    #[test]
+    fn content_length_parses_or_defaults_zero() {
+        let mut headers = header::HeaderMap::new();
+        assert_eq!(content_length(&headers), 0);
+        headers.insert(header::CONTENT_LENGTH, header::HeaderValue::from(4096));
+        assert_eq!(content_length(&headers), 4096);
+    }
+
     async fn rollout_state() -> (Arc<AppState>, TempDir) {
         let dir = TempDir::new().unwrap();
         let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await);