From 75eac36dc97aa3d377dff8ac34ca692cf292d1a9 Mon Sep 17 00:00:00 2001 From: Beinan Date: Tue, 14 Jul 2026 07:26:14 +0000 Subject: [PATCH] feat(server): scale rollout stores to hundreds of thousands via registry + LRU + global sweeper Per-experiment physical partitioning (one dataset per experiment) requires the server runtime to decouple resident memory and background tasks from the number of stores. Three changes make that safe at 100k+ experiments: - Durable directory (`RolloutRegistry`, a single Lance table) is the source of truth for which stores exist. `create`/`list`/`delete`/existence checks go through it instead of the in-memory map, so eviction can no longer cause a false 404 and `list` enumerates every store, not just resident ones. - Bounded LRU replaces the unbounded `HashMap` of store handles; evicted stores are transparently reopened on the next request. - A single process-wide WAL-cleanup sweeper replaces the former one-tokio-timer-per-store model (which would spawn 100k+ timers). It sweeps the resident set with the same per-pass timeout as the periodic cleanup. `RolloutStoreInfo.version` becomes `Option` (omitted in list responses to avoid opening every dataset; populated for single-store get/create). Adds `--rollout-cache-capacity` / `ROLLOUT_CACHE_CAPACITY`. Co-Authored-By: Claude Opus 4 --- Cargo.lock | 1 + crates/lance-context-api/src/lib.rs | 7 +- crates/lance-context-client/src/lib.rs | 4 +- crates/lance-context-core/src/lib.rs | 2 + crates/lance-context-core/src/registry.rs | 310 +++++++++++++++++ crates/lance-context-server/Cargo.toml | 1 + crates/lance-context-server/src/config.rs | 9 + crates/lance-context-server/src/main.rs | 13 +- .../src/routes/records.rs | 28 +- .../src/routes/rollouts.rs | 172 +++++++--- .../lance-context-server/src/routes/search.rs | 18 +- crates/lance-context-server/src/state.rs | 315 +++++++++++++++--- 12 files changed, 748 insertions(+), 132 deletions(-) create mode 100644 crates/lance-context-core/src/registry.rs diff --git a/Cargo.lock b/Cargo.lock index 31da8498..008b065c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5447,6 +5447,7 @@ dependencies = [ "clap", "lance-context-api", "lance-context-core", + "lru", "serde", "serde_json", "tempfile", diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index b57d6324..b9bfbbf3 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -580,7 +580,12 @@ pub struct CreateRolloutStoreRequest { pub struct RolloutStoreInfo { pub name: String, pub uri: String, - pub version: u64, + /// Dataset version. `None` in list responses, which are served from the + /// durable registry without opening each dataset (a store may not be + /// resident in memory). Single-store lookups (`get`/`create`) always + /// populate it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub version: Option, } #[derive(Debug, Serialize, Deserialize)] diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index e13ea0ee..1ece4c36 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -264,7 +264,7 @@ impl RemoteRolloutStore { Ok(Self { client, store_name: store_name.to_string(), - cached_version: info.version, + cached_version: info.version.unwrap_or(0), }) } @@ -281,7 +281,7 @@ impl RemoteRolloutStore { Ok(Self { client, store_name: req.name.clone(), - cached_version: info.version, + cached_version: info.version.unwrap_or(0), }) } } diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 54cfe0d1..41114323 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -8,6 +8,7 @@ mod export; mod id; mod namespace; mod record; +mod registry; mod rollout; mod rollout_store; pub mod serde; @@ -31,6 +32,7 @@ pub use record::{ RetrieveResult, SearchResult, StateMetadata, UpdateResult, UpsertResult, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED, }; +pub use registry::{RegistryEntry, RolloutRegistry}; pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL}; pub use rollout_store::{rollout_schema, RolloutStore, RolloutStoreOptions}; pub use store::{ diff --git a/crates/lance-context-core/src/registry.rs b/crates/lance-context-core/src/registry.rs new file mode 100644 index 00000000..ce20852a --- /dev/null +++ b/crates/lance-context-core/src/registry.rs @@ -0,0 +1,310 @@ +//! Persistent directory of rollout stores. +//! +//! # Why this exists +//! +//! With `experiment_name` used as the physical partition key, a deployment may +//! hold **tens or hundreds of thousands** of independent rollout datasets — one +//! per experiment. The server can no longer keep every opened store resident in +//! memory (it uses a bounded LRU), so "is this store in the in-memory map?" is +//! no longer a valid existence check, and `list` can no longer enumerate stores +//! by walking the cache. +//! +//! [`RolloutRegistry`] is a single, small Lance dataset — one row per rollout +//! store (`name`, `uri`, `created_at`) — that serves as the durable source of +//! truth for *which* stores exist. It is consulted on a cache miss to decide +//! whether to lazily open a store (vs. return 404), and it backs the `list` +//! endpoint without touching object storage for every experiment. +//! +//! It deliberately holds only cheap directory metadata, never rollout rows. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator, StringArray}; +use arrow_schema::{ArrowError, DataType, Field, Schema}; +use chrono::Utc; +use futures::TryStreamExt; +use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; +use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; +use lance::{Error as LanceError, Result as LanceResult}; + +/// One entry in the rollout-store directory. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RegistryEntry { + /// Logical store name (the experiment name); unique key. + pub name: String, + /// Physical dataset URI, as produced by the server's `rollout_uri`. + pub uri: String, + /// Creation time, Unix milliseconds. + pub created_at: i64, +} + +/// Durable directory of rollout stores, backed by a single Lance dataset. +/// +/// All mutations (`upsert`, `remove`) take `&mut self` and are expected to be +/// serialized by the caller (the server wraps this in a lock), so the registry +/// itself does no internal locking. Reads (`list`, `contains`) take `&self`. +pub struct RolloutRegistry { + dataset: Dataset, + uri: String, + storage_options: Option>, +} + +fn registry_schema() -> Schema { + Schema::new(vec![ + Field::new("name", DataType::Utf8, false), + Field::new("uri", DataType::Utf8, false), + Field::new("created_at", DataType::Int64, false), + ]) +} + +impl RolloutRegistry { + /// Open the registry dataset at `uri`, creating an empty one if it does not + /// yet exist. Idempotent across process restarts. + pub async fn open_or_create( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + let dataset = match Self::load(uri, storage_options.clone()).await { + Ok(dataset) => dataset, + Err(LanceError::DatasetNotFound { .. }) => { + Self::create(uri, storage_options.clone()).await? + } + Err(err) => return Err(err), + }; + Ok(Self { + dataset, + uri: uri.to_string(), + storage_options, + }) + } + + async fn load( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + if let Some(options) = storage_options { + DatasetBuilder::from_uri(uri) + .with_storage_options(options) + .load() + .await + } else { + Dataset::open(uri).await + } + } + + async fn create( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + let schema = Arc::new(registry_schema()); + let empty = RecordBatch::new_empty(schema.clone()); + let batches = RecordBatchIterator::new( + vec![Ok::(empty)].into_iter(), + schema.clone(), + ); + let params = Self::write_params(WriteMode::Create, storage_options); + Dataset::write(batches, uri, Some(params)).await + } + + fn write_params( + mode: WriteMode, + storage_options: Option>, + ) -> WriteParams { + let mut params = WriteParams { + mode, + ..Default::default() + }; + if let Some(options) = storage_options { + params.store_params = Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new( + StorageOptionsAccessor::with_static_options(options), + )), + ..Default::default() + }); + } + params + } + + /// Insert or replace the directory entry for `name`. + /// + /// Implemented as delete-same-name-then-append so it is **idempotent**: a + /// `create` retried after a crash (dataset on disk, registry row possibly + /// present) converges to exactly one row. Callers must serialize mutations. + pub async fn upsert(&mut self, name: &str, uri: &str) -> LanceResult<()> { + self.delete_row(name).await?; + let schema = Arc::new(registry_schema()); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![name])), + Arc::new(StringArray::from(vec![uri])), + Arc::new(Int64Array::from(vec![Utc::now().timestamp_millis()])), + ], + )?; + let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema); + let params = Self::write_params(WriteMode::Append, self.storage_options.clone()); + self.dataset.append(reader, Some(params)).await?; + Ok(()) + } + + /// Remove the directory entry for `name`, if present. No-op when absent. + pub async fn remove(&mut self, name: &str) -> LanceResult<()> { + self.delete_row(name).await + } + + async fn delete_row(&mut self, name: &str) -> LanceResult<()> { + let escaped = name.replace('\'', "''"); + self.dataset + .delete(&format!("name = '{}'", escaped)) + .await?; + Ok(()) + } + + /// Whether a store named `name` exists in the directory. + pub async fn contains(&self, name: &str) -> LanceResult { + let escaped = name.replace('\'', "''"); + let mut scanner = self.dataset.scan(); + scanner.project(&["name"])?; + scanner.filter(&format!("name = '{}'", escaped))?; + scanner.limit(Some(1), None)?; + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + if batch.num_rows() > 0 { + return Ok(true); + } + } + Ok(false) + } + + /// All directory entries, ordered as stored (unspecified). The registry is a + /// narrow three-column table, so even hundreds of thousands of rows scan + /// quickly; pagination can be layered on later if needed. + pub async fn list(&self) -> LanceResult> { + let mut scanner = self.dataset.scan(); + scanner.project(&["name", "uri", "created_at"])?; + let mut stream = scanner.try_into_stream().await?; + let mut out = Vec::new(); + while let Some(batch) = stream.try_next().await? { + let names = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + LanceError::from(ArrowError::InvalidArgumentError( + "registry 'name' column is not Utf8".into(), + )) + })?; + let uris = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + LanceError::from(ArrowError::InvalidArgumentError( + "registry 'uri' column is not Utf8".into(), + )) + })?; + let created = batch + .column(2) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + LanceError::from(ArrowError::InvalidArgumentError( + "registry 'created_at' column is not Int64".into(), + )) + })?; + for i in 0..batch.num_rows() { + out.push(RegistryEntry { + name: names.value(i).to_string(), + uri: uris.value(i).to_string(), + created_at: created.value(i), + }); + } + } + Ok(out) + } + + /// The registry dataset URI. + #[must_use] + pub fn uri(&self) -> &str { + &self.uri + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + async fn new_registry(dir: &TempDir) -> RolloutRegistry { + let uri = dir.path().join("_registry.rollout.lance"); + RolloutRegistry::open_or_create(uri.to_str().unwrap(), None) + .await + .unwrap() + } + + #[tokio::test] + async fn upsert_is_idempotent() { + let dir = TempDir::new().unwrap(); + let mut reg = new_registry(&dir).await; + reg.upsert("exp-a", "/data/exp-a.rollout.lance") + .await + .unwrap(); + // Re-upserting the same name must not create a duplicate row. + reg.upsert("exp-a", "/data/exp-a.rollout.lance") + .await + .unwrap(); + let entries = reg.list().await.unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].name, "exp-a"); + assert!(reg.contains("exp-a").await.unwrap()); + } + + #[tokio::test] + async fn contains_and_remove() { + let dir = TempDir::new().unwrap(); + let mut reg = new_registry(&dir).await; + assert!(!reg.contains("missing").await.unwrap()); + reg.upsert("exp-b", "/data/exp-b.rollout.lance") + .await + .unwrap(); + assert!(reg.contains("exp-b").await.unwrap()); + reg.remove("exp-b").await.unwrap(); + assert!(!reg.contains("exp-b").await.unwrap()); + assert!(reg.list().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn list_returns_all_entries() { + let dir = TempDir::new().unwrap(); + let mut reg = new_registry(&dir).await; + for i in 0..5 { + reg.upsert(&format!("exp-{i}"), &format!("/data/exp-{i}.rollout.lance")) + .await + .unwrap(); + } + let mut names: Vec = reg + .list() + .await + .unwrap() + .into_iter() + .map(|e| e.name) + .collect(); + names.sort(); + assert_eq!(names, vec!["exp-0", "exp-1", "exp-2", "exp-3", "exp-4"]); + } + + #[tokio::test] + async fn survives_reopen() { + let dir = TempDir::new().unwrap(); + { + let mut reg = new_registry(&dir).await; + reg.upsert("persist", "/data/persist.rollout.lance") + .await + .unwrap(); + } + // Reopen the same path: the entry must still be present. + let reg = new_registry(&dir).await; + assert!(reg.contains("persist").await.unwrap()); + } +} diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index 7e945fb9..d7ce72c1 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -18,6 +18,7 @@ lance-context-api = { version = "0.6.3", path = "../lance-context-api" } axum = { version = "0.8", features = ["json", "multipart"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } clap = { version = "4", features = ["derive", "env"] } +lru = "0.12" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index db618633..8337fb8c 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -50,6 +50,15 @@ pub struct ServerConfig { /// set. Defaults to `1` (clean up whenever any generation is present). #[arg(long, env = "ROLLOUT_CLEANUP_MIN_GENERATIONS", default_value = "1")] pub rollout_cleanup_min_generations: usize, + + /// Upper bound on the number of resident rollout-store handles kept in + /// memory (an LRU). With one physical dataset per experiment a deployment + /// may hold hundreds of thousands of stores; this bounds how many stay open + /// at once. Size it for peak *concurrent* experiments, not the total count. + /// Evicted stores are transparently reopened on the next request. `0` falls + /// back to the built-in default. + #[arg(long, env = "ROLLOUT_CACHE_CAPACITY", default_value = "2000")] + pub rollout_cache_capacity: usize, } impl ServerConfig { diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index 286a26e1..38b14a3f 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -32,7 +32,18 @@ async fn main() { std::process::exit(1); } - let state = Arc::new(AppState::new(config)); + let state = match AppState::new(config).await { + Ok(state) => Arc::new(state), + Err(e) => { + tracing::error!("Failed to initialize server state: {:?}", e); + std::process::exit(1); + } + }; + + // Single process-wide WAL-cleanup sweeper (replaces one timer per store). + // The handle is detached: it runs for the server's lifetime and stops when + // the last `Arc` is dropped. + let _sweeper = state.spawn_global_sweeper(); let app = routes::router() .with_state(state) diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 98a453c5..0a7395e6 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -538,7 +538,6 @@ fn record_from_add_request( #[cfg(test)] mod tests { - use std::collections::HashMap; use std::sync::Arc; use axum::extract::{Path, Query, State}; @@ -563,17 +562,12 @@ mod tests { .to_string_lossy() .to_string(); let store = ContextStore::open(&uri).await.unwrap(); - let mut stores = HashMap::new(); - stores.insert(context_name.to_string(), Arc::new(RwLock::new(store))); - let state = Arc::new(AppState { - stores: RwLock::new(stores), - rollout_stores: RwLock::new(HashMap::new()), - base_path: dir.path().to_path_buf(), - instance_id: None, - rollout_merge_after_generations: 0, - rollout_cleanup_interval_secs: 0, - rollout_cleanup_min_generations: 1, - }); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + state + .stores + .write() + .await + .insert(context_name.to_string(), Arc::new(RwLock::new(store))); (state, dir) } @@ -1220,15 +1214,7 @@ mod tests { .expect("write on instance A"); // Instance B: fresh AppState over the same data dir, empty cache. - let state_b = Arc::new(AppState { - stores: RwLock::new(HashMap::new()), - rollout_stores: RwLock::new(HashMap::new()), - base_path: dir.path().to_path_buf(), - instance_id: None, - rollout_merge_after_generations: 0, - rollout_cleanup_interval_secs: 0, - rollout_cleanup_min_generations: 1, - }); + let state_b = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); let Json(list) = list_records( State(state_b.clone()), diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 54d41a4f..ad6769aa 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -30,14 +30,20 @@ pub async fn create_rollout_store( State(state): State>, Json(req): Json, ) -> Result<(StatusCode, Json), AppError> { - let stores = state.rollout_stores.read().await; - if stores.contains_key(&req.name) { + // Existence is tracked durably in the registry, not by cache membership. + if state + .rollout_registry + .read() + .await + .contains(&req.name) + .await + .map_err(AppError::from_lance)? + { return Err(AppError::AlreadyExists(format!( "Rollout store '{}' already exists", req.name ))); } - drop(stores); let uri = state.rollout_uri(&req.name); let options = RolloutStoreOptions { @@ -56,40 +62,44 @@ pub async fn create_rollout_store( let version = store.version(); let store = Arc::new(RwLock::new(store)); - // Start the periodic per-shard WAL-cleanup timer (no-op when the interval is - // disabled). The handle is detached: it is aborted when the store is dropped - // on delete, and otherwise runs for the server's lifetime. - let _cleanup = RolloutStore::spawn_periodic_cleanup(store.clone()); - - let mut stores = state.rollout_stores.write().await; - stores.insert(req.name.clone(), store); + // Record in the durable registry and the LRU. WAL cleanup is handled by the + // single process-wide sweeper (see `AppState::spawn_global_sweeper`), so no + // per-store timer is started here. + state.register_rollout(&req.name, &uri, store).await?; Ok(( StatusCode::CREATED, Json(RolloutStoreInfo { name: req.name, uri, - version, + version: Some(version), }), )) } pub async fn list_rollout_stores( State(state): State>, -) -> Json { - let stores = state.rollout_stores.read().await; - let mut out = Vec::with_capacity(stores.len()); - - for (name, store_lock) in stores.iter() { - let store = store_lock.read().await; - out.push(RolloutStoreInfo { - name: name.clone(), - uri: state.rollout_uri(name), - version: store.version(), - }); - } +) -> Result, AppError> { + // Served from the durable registry so it enumerates *all* stores, not just + // those currently resident in the LRU. `version` is omitted here to avoid + // opening every dataset. + let entries = state + .rollout_registry + .read() + .await + .list() + .await + .map_err(AppError::from_lance)?; + let out = entries + .into_iter() + .map(|entry| RolloutStoreInfo { + name: entry.name, + uri: entry.uri, + version: None, + }) + .collect(); - Json(ListRolloutStoresResponse { stores: out }) + Ok(Json(ListRolloutStoresResponse { stores: out })) } pub async fn get_rollout_store( @@ -102,7 +112,7 @@ pub async fn get_rollout_store( Ok(Json(RolloutStoreInfo { name: name.clone(), uri: state.rollout_uri(&name), - version: store.version(), + version: Some(store.version()), })) } @@ -110,8 +120,9 @@ pub async fn delete_rollout_store( State(state): State>, Path(name): Path, ) -> Result { - let mut stores = state.rollout_stores.write().await; - if stores.remove(&name).is_none() { + // Removes from the durable registry and evicts any resident handle; 404 when + // the store is not registered. + if !state.unregister_rollout(&name).await? { return Err(AppError::NotFound(format!( "Rollout store '{}' does not exist", name @@ -532,27 +543,17 @@ fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { #[cfg(test)] mod tests { - use std::collections::HashMap; use std::sync::Arc; use lance_context_api::CreateRolloutStoreRequest; use tempfile::TempDir; - use tokio::sync::RwLock; use super::*; use crate::state::AppState; async fn rollout_state() -> (Arc, TempDir) { let dir = TempDir::new().unwrap(); - let state = Arc::new(AppState { - stores: RwLock::new(HashMap::new()), - rollout_stores: RwLock::new(HashMap::new()), - base_path: dir.path().to_path_buf(), - instance_id: None, - rollout_merge_after_generations: 0, - rollout_cleanup_interval_secs: 0, - rollout_cleanup_min_generations: 1, - }); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); let (_status, _info) = create_rollout_store( State(state.clone()), Json(CreateRolloutStoreRequest { @@ -803,15 +804,13 @@ mod tests { // Pod B: a fresh AppState over the SAME data dir, with an empty cache and // a different instance id (its own shard). It never saw the `create`. - let state_b = Arc::new(AppState { - stores: RwLock::new(HashMap::new()), - rollout_stores: RwLock::new(HashMap::new()), - base_path: dir.path().to_path_buf(), - instance_id: Some("rollout-1".to_string()), - rollout_merge_after_generations: 0, - rollout_cleanup_interval_secs: 0, - rollout_cleanup_min_generations: 1, - }); + let state_b = Arc::new( + AppState::new_for_test_with_instance( + dir.path().to_path_buf(), + Some("rollout-1".to_string()), + ) + .await, + ); // Read routed to B must lazily open the store, not 404. let Json(list) = list_rollouts( @@ -853,4 +852,83 @@ mod tests { .unwrap_err(); assert!(matches!(err, AppError::NotFound(_))); } + + /// With a capacity-1 LRU, creating a second store must evict the first — but + /// the evicted store is still reachable, transparently reopened from storage + /// via the durable registry, not falsely 404'd. + #[tokio::test] + async fn evicted_store_is_transparently_reopened() { + let dir = TempDir::new().unwrap(); + // Force a capacity-1 cache so the second create evicts the first. + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + state.rollout_stores = + tokio::sync::Mutex::new(lru::LruCache::new(std::num::NonZeroUsize::new(1).unwrap())); + let state = Arc::new(state); + + for name in ["exp-a", "exp-b"] { + let _ = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: name.to_string(), + storage_options: None, + }), + ) + .await + .expect("create"); + } + + // exp-a was evicted (cache holds only exp-b now), yet a lookup succeeds. + assert_eq!(state.rollout_stores.lock().await.len(), 1); + let info = get_rollout_store(State(state.clone()), Path("exp-a".to_string())) + .await + .expect("evicted store reopens instead of 404"); + assert_eq!(info.name, "exp-a"); + } + + /// `list` is served from the durable registry, so it enumerates every store + /// that exists — including ones evicted from (or never resident in) the LRU. + #[tokio::test] + async fn list_enumerates_all_registered_stores() { + let dir = TempDir::new().unwrap(); + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + state.rollout_stores = + tokio::sync::Mutex::new(lru::LruCache::new(std::num::NonZeroUsize::new(1).unwrap())); + let state = Arc::new(state); + + for name in ["e0", "e1", "e2"] { + let _ = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: name.to_string(), + storage_options: None, + }), + ) + .await + .expect("create"); + } + // Only one store is resident, but list must return all three. + assert_eq!(state.rollout_stores.lock().await.len(), 1); + let Json(listed) = list_rollout_stores(State(state.clone())) + .await + .expect("list"); + let mut names: Vec = listed.stores.into_iter().map(|s| s.name).collect(); + names.sort(); + assert_eq!(names, vec!["e0", "e1", "e2"]); + } + + /// Delete removes the store from the durable registry, so a subsequent read + /// 404s and it no longer appears in `list`. + #[tokio::test] + async fn delete_unregisters_store() { + let (state, _dir) = rollout_state().await; + delete_rollout_store(State(state.clone()), Path("rl".to_string())) + .await + .expect("delete"); + let err = get_rollout_store(State(state.clone()), Path("rl".to_string())) + .await + .unwrap_err(); + assert!(matches!(err, AppError::NotFound(_))); + let Json(listed) = list_rollout_stores(State(state.clone())).await.unwrap(); + assert!(listed.stores.is_empty()); + } } diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index e8bbe25a..1efc79cd 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -106,7 +106,6 @@ pub async fn retrieve( #[cfg(test)] mod tests { - use std::collections::HashMap; use std::sync::Arc; use axum::extract::{Path, State}; @@ -141,17 +140,12 @@ mod tests { ) .await .unwrap(); - let mut stores = HashMap::new(); - stores.insert(CTX.to_string(), Arc::new(RwLock::new(store))); - let state = Arc::new(AppState { - stores: RwLock::new(stores), - rollout_stores: RwLock::new(HashMap::new()), - base_path: dir.path().to_path_buf(), - instance_id: None, - rollout_merge_after_generations: 0, - rollout_cleanup_interval_secs: 0, - rollout_cleanup_min_generations: 1, - }); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + state + .stores + .write() + .await + .insert(CTX.to_string(), Arc::new(RwLock::new(store))); (state, dir) } diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 5049e6a8..6f8fc520 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -1,16 +1,38 @@ -use std::collections::HashMap; +use std::num::NonZeroUsize; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; -use lance_context_core::{ContextStore, ContextStoreOptions, RolloutStore, RolloutStoreOptions}; -use tokio::sync::RwLock; +use lance_context_core::{ + ContextStore, ContextStoreOptions, RolloutRegistry, RolloutStore, RolloutStoreOptions, +}; +use lru::LruCache; +use tokio::sync::{Mutex, RwLock}; +use tokio::task::JoinHandle; use crate::config::ServerConfig; use crate::error::AppError; +/// Default upper bound on resident rollout-store handles when the config does +/// not specify one. Sized for peak *concurrent* experiments, not the total +/// number of experiments (which may be hundreds of thousands). +pub const DEFAULT_ROLLOUT_CACHE_CAPACITY: usize = 2000; + pub struct AppState { - pub stores: RwLock>>>, - pub rollout_stores: RwLock>>>, + pub stores: RwLock>>>, + /// Bounded LRU of resident rollout-store handles. + /// + /// With one physical dataset per experiment, the deployment may hold + /// hundreds of thousands of stores; keeping them all resident would exhaust + /// memory and (formerly) spawn one background timer each. This cache bounds + /// residency: on overflow the least-recently-used handle is evicted and its + /// `Arc>` dropped. Existence is tracked durably by + /// [`Self::rollout_registry`], not by membership in this cache. + pub rollout_stores: Mutex>>>, + /// Durable directory of which rollout stores exist. Consulted on a cache + /// miss (existence check) and to back the list endpoint. Guarded by a lock + /// because mutations take `&mut`. + pub rollout_registry: RwLock, pub base_path: PathBuf, /// Stable identity of this server instance, used as the MemWAL shard key for /// rollout writes so each instance owns exactly one shard. `None` falls back @@ -20,7 +42,7 @@ pub struct AppState { /// disables it. See `RolloutStoreOptions::merge_after_generations`. pub rollout_merge_after_generations: usize, /// Periodic per-shard WAL-cleanup interval in seconds; `0` disables the - /// background timer. See `RolloutStoreOptions::cleanup_interval_secs`. + /// global sweeper. See [`Self::spawn_global_sweeper`]. pub rollout_cleanup_interval_secs: u64, /// Minimum flushed generations before a periodic cleanup tick merges. See /// `RolloutStoreOptions::cleanup_min_generations`. @@ -28,17 +50,30 @@ pub struct AppState { } impl AppState { - pub fn new(config: ServerConfig) -> Self { + /// Build the shared server state, opening (or creating) the rollout registry + /// under `data_dir`. Async because opening the registry touches storage. + pub async fn new(config: ServerConfig) -> Result { let instance_id = config.resolved_instance_id(); - Self { - stores: RwLock::new(HashMap::new()), - rollout_stores: RwLock::new(HashMap::new()), - base_path: PathBuf::from(&config.data_dir), + let base_path = PathBuf::from(&config.data_dir); + let registry_uri = base_path + .join("_registry.rollout.lance") + .to_string_lossy() + .to_string(); + let registry = RolloutRegistry::open_or_create(®istry_uri, None) + .await + .map_err(AppError::from_lance)?; + let capacity = NonZeroUsize::new(config.rollout_cache_capacity) + .unwrap_or_else(|| NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap()); + Ok(Self { + stores: RwLock::new(std::collections::HashMap::new()), + rollout_stores: Mutex::new(LruCache::new(capacity)), + rollout_registry: RwLock::new(registry), + base_path, instance_id, rollout_merge_after_generations: config.rollout_merge_after_generations, rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, rollout_cleanup_min_generations: config.rollout_cleanup_min_generations, - } + }) } pub fn context_uri(&self, name: &str) -> String { @@ -48,6 +83,41 @@ impl AppState { .to_string() } + /// Build a default-configured `AppState` rooted at `base_path`, for tests. + /// Opens a fresh registry under the directory; cleanup interval disabled. + #[cfg(test)] + pub async fn new_for_test(base_path: PathBuf) -> Self { + Self::new_for_test_with_instance(base_path, None).await + } + + /// Like [`Self::new_for_test`] but with an explicit MemWAL instance id, so a + /// second in-process "instance" can present its own shard. + #[cfg(test)] + pub async fn new_for_test_with_instance( + base_path: PathBuf, + instance_id: Option, + ) -> Self { + let registry_uri = base_path + .join("_registry.rollout.lance") + .to_string_lossy() + .to_string(); + let registry = RolloutRegistry::open_or_create(®istry_uri, None) + .await + .expect("open test registry"); + Self { + stores: RwLock::new(std::collections::HashMap::new()), + rollout_stores: Mutex::new(LruCache::new( + NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), + )), + rollout_registry: RwLock::new(registry), + base_path, + instance_id, + rollout_merge_after_generations: 0, + rollout_cleanup_interval_secs: 0, + rollout_cleanup_min_generations: 1, + } + } + /// Rollout datasets live under a distinct `.rollout.lance` suffix so a /// rollout store and a context store may share the same name without /// colliding on disk. @@ -58,56 +128,115 @@ impl AppState { .to_string() } + fn rollout_store_options(&self) -> RolloutStoreOptions { + RolloutStoreOptions { + // No request body on the read path: object-store credentials come + // from the pod's workload-identity environment. + storage_options: None, + shard_id: self.instance_id.clone(), + merge_after_generations: (self.rollout_merge_after_generations > 0) + .then_some(self.rollout_merge_after_generations), + cleanup_interval_secs: (self.rollout_cleanup_interval_secs > 0) + .then_some(self.rollout_cleanup_interval_secs), + cleanup_min_generations: Some(self.rollout_cleanup_min_generations), + } + } + + /// Record that a rollout store exists, in both the durable registry and the + /// in-memory LRU. Called by the create route after the dataset is written. + pub async fn register_rollout( + &self, + name: &str, + uri: &str, + store: Arc>, + ) -> Result<(), AppError> { + self.rollout_registry + .write() + .await + .upsert(name, uri) + .await + .map_err(AppError::from_lance)?; + // Insertion may evict the LRU's least-recently-used entry; dropping the + // returned handle releases it (no per-store timer to abort — cleanup is + // now global). + self.rollout_stores + .lock() + .await + .put(name.to_string(), store); + Ok(()) + } + + /// Remove a rollout store from the durable registry and evict any resident + /// handle. Returns whether the store existed. + pub async fn unregister_rollout(&self, name: &str) -> Result { + let existed = self + .rollout_registry + .read() + .await + .contains(name) + .await + .map_err(AppError::from_lance)?; + if !existed { + return Ok(false); + } + self.rollout_registry + .write() + .await + .remove(name) + .await + .map_err(AppError::from_lance)?; + self.rollout_stores.lock().await.pop(name); + Ok(true) + } + /// Look up a rollout store by name, lazily loading it from object storage on /// a local cache miss. /// - /// The server caches each opened store in this process's memory. In a - /// multi-replica deployment a store `create`d on pod A only enters A's map, - /// so a read/write that the load balancer routes to pod B would otherwise - /// 404 even though the dataset exists on shared object storage. This helper - /// closes that gap: on a miss it *loads* (never creates — see - /// [`RolloutStore::open_existing_with_options`], which returns a - /// `DatasetNotFound`/404 for a genuinely-absent store rather than silently - /// materializing an empty table) and caches the handle for subsequent hits. + /// Unlike the previous implementation, a cache miss is **not** a 404: the + /// bounded LRU may have evicted a store that still exists. Existence is + /// resolved against the durable [`RolloutRegistry`]; only a store absent + /// from the registry yields [`AppError::NotFound`]. This closes the + /// multi-replica gap (create on pod A, read on pod B) *and* the + /// eviction-induced false-404 that per-experiment partitioning introduces. pub async fn get_or_open_rollout_store( &self, name: &str, ) -> Result>, AppError> { - // Fast path: already cached in this process. - if let Some(store) = self.rollout_stores.read().await.get(name) { + // Fast path: resident in this process's LRU (updates recency). + if let Some(store) = self.rollout_stores.lock().await.get(name) { return Ok(store.clone()); } - // Slow path: load from object storage WITHOUT holding the map lock, so a - // slow open does not block other stores' requests. + // Existence is the registry's job, not the cache's. + let exists = self + .rollout_registry + .read() + .await + .contains(name) + .await + .map_err(AppError::from_lance)?; + if !exists { + return Err(AppError::NotFound(format!( + "Rollout store '{}' does not exist", + name + ))); + } + + // Load from storage WITHOUT holding the LRU lock, so a slow open does + // not block other stores' requests. let uri = self.rollout_uri(name); - let options = RolloutStoreOptions { - // No request body on the read path: object-store credentials come - // from the pod's workload-identity environment, exactly as they do - // for the `create` route in production. - storage_options: None, - shard_id: self.instance_id.clone(), - merge_after_generations: (self.rollout_merge_after_generations > 0) - .then_some(self.rollout_merge_after_generations), - cleanup_interval_secs: (self.rollout_cleanup_interval_secs > 0) - .then_some(self.rollout_cleanup_interval_secs), - cleanup_min_generations: Some(self.rollout_cleanup_min_generations), - }; - let opened = RolloutStore::open_existing_with_options(&uri, options) + let opened = RolloutStore::open_existing_with_options(&uri, self.rollout_store_options()) .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); - // Insert under the write lock, re-checking for a store another request - // may have opened concurrently while we were loading. - let mut stores = self.rollout_stores.write().await; - if let Some(existing) = stores.get(name) { + // Insert under the lock, re-checking for a store another request may + // have opened concurrently while we were loading. + let mut cache = self.rollout_stores.lock().await; + if let Some(existing) = cache.get(name) { return Ok(existing.clone()); } - // A lazily-opened store must also get the background WAL-cleanup timer - // that the `create` route starts, or these stores would never merge. - let _cleanup = RolloutStore::spawn_periodic_cleanup(opened.clone()); - stores.insert(name.to_string(), opened.clone()); + cache.put(name.to_string(), opened.clone()); Ok(opened) } @@ -124,9 +253,6 @@ impl AppState { } let uri = self.context_uri(name); - // Load-only: existing schema (embedding dim, blob columns, distance - // metric) is read from the persisted dataset, so no create-time options - // are needed. Credentials come from the pod environment. let opened = ContextStore::open_existing_with_options(&uri, ContextStoreOptions::default()) .await .map_err(AppError::from_lance)?; @@ -139,4 +265,97 @@ impl AppState { stores.insert(name.to_string(), opened.clone()); Ok(opened) } + + /// Spawn the single, process-wide WAL-cleanup sweeper. + /// + /// This replaces the former one-timer-per-store model, which does not scale + /// to hundreds of thousands of per-experiment datasets. On each tick the + /// sweeper snapshots the *resident* rollout stores (those in the LRU) and + /// folds each one's flushed MemWAL generations into its base table, guarding + /// every pass with a timeout so a wedged store cannot stall the sweeper (see + /// [`RolloutStore::cleanup_own_shard`] and the timeout added in the periodic + /// cleanup hardening). + /// + /// Cold stores (evicted from the LRU) are not swept: having no new writes, + /// they accumulate no new generations, and are swept again the moment a + /// request re-opens them. Returns `None` when the interval is `0`. + pub fn spawn_global_sweeper(self: &Arc) -> Option> { + let interval_secs = self.rollout_cleanup_interval_secs; + if interval_secs == 0 { + return None; + } + let interval = Duration::from_secs(interval_secs); + // Abandon any single pass that outruns five intervals (min 30s) so one + // stuck store cannot wedge cleanup for every other store. + let pass_timeout = interval.saturating_mul(5).max(Duration::from_secs(30)); + let weak = Arc::downgrade(self); + Some(tokio::spawn(async move { + let mut ticker = tokio::time::interval(interval); + ticker.tick().await; // skip the immediate first tick + loop { + ticker.tick().await; + let Some(state) = weak.upgrade() else { + return; + }; + // Snapshot resident stores without holding the LRU lock across + // the awaits below. + let resident: Vec<(String, Arc>)> = { + let cache = state.rollout_stores.lock().await; + cache + .iter() + .map(|(name, store)| (name.clone(), store.clone())) + .collect() + }; + for (name, store) in resident { + let mut guard = store.write().await; + match tokio::time::timeout(pass_timeout, guard.cleanup_own_shard()).await { + Ok(Ok(0)) => {} + Ok(Ok(n)) => tracing::info!( + store = %name, + reclaimed = n, + "global sweeper merged flushed generations" + ), + Ok(Err(e)) => tracing::warn!( + store = %name, + error = %e, + "global sweeper WAL cleanup failed" + ), + Err(_elapsed) => tracing::warn!( + store = %name, + "global sweeper WAL cleanup timed out; abandoning this store this tick" + ), + } + } + } + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + async fn state_with_interval(dir: &TempDir, secs: u64) -> Arc { + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + state.rollout_cleanup_interval_secs = secs; + Arc::new(state) + } + + /// The sweeper is a single detached task, gated on the cleanup interval: + /// spawned when non-zero, `None` when disabled. This is the whole-process + /// replacement for the former one-timer-per-store model. + #[tokio::test] + async fn global_sweeper_is_gated_on_interval() { + let dir = TempDir::new().unwrap(); + let disabled = state_with_interval(&dir, 0).await; + assert!(disabled.spawn_global_sweeper().is_none()); + + let dir2 = TempDir::new().unwrap(); + let enabled = state_with_interval(&dir2, 3600).await; + let handle = enabled + .spawn_global_sweeper() + .expect("sweeper spawns when interval > 0"); + handle.abort(); + } }