diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 509f0ff..9571e2e 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -14,12 +14,13 @@ pub struct ServerConfig { pub data_dir: String, /// Stable identity of this server instance, used as the MemWAL shard key for - /// rollout writes. Each instance must present a distinct, stable value so it - /// owns exactly one shard and never contends with peers (see + /// context, rollout, datagen, and generic writes. Each instance must present + /// a distinct, stable value so it owns exactly one shard per store and never + /// contends with peers (see /// `docs/src/specs/rollout-deployment.md`). In Kubernetes, set this to the /// StatefulSet pod ordinal hostname (e.g. `rollout-0`). Defaults to the /// `INSTANCE_ID` env var, then the pod/host `HOSTNAME`; if neither is set, - /// rollout writes fall back to a single shared `default` shard, which is + /// writes fall back to a single shared `default` shard per store, which is /// only safe for a single-instance deployment. #[arg(long, env = "INSTANCE_ID")] pub instance_id: Option, @@ -122,9 +123,9 @@ pub struct ServerConfig { } impl ServerConfig { - /// Resolve the instance id used for rollout MemWAL sharding: the explicit - /// `--instance-id`/`INSTANCE_ID` if provided, otherwise the `HOSTNAME` - /// environment variable (stable per-pod under a StatefulSet). + /// Resolve the instance id used for server-managed MemWAL sharding: the + /// explicit `--instance-id`/`INSTANCE_ID` if provided, otherwise the + /// `HOSTNAME` environment variable (stable per-pod under a StatefulSet). #[must_use] pub fn resolved_instance_id(&self) -> Option { self.instance_id diff --git a/crates/lance-context-server/src/routes/contexts.rs b/crates/lance-context-server/src/routes/contexts.rs index bbfb831..d1dd1ef 100644 --- a/crates/lance-context-server/src/routes/contexts.rs +++ b/crates/lance-context-server/src/routes/contexts.rs @@ -52,7 +52,7 @@ pub async fn create_context( blob_columns, id_index_type, distance_metric, - ..Default::default() + ..state.context_store_options() }; let store = ContextStore::open_with_options(&uri, options) @@ -128,6 +128,9 @@ pub async fn delete_context( #[cfg(test)] mod tests { use super::*; + use crate::routes::records::add_records; + use lance_context_api::{AddRecordRequest, AddRecordsRequest}; + use std::time::Duration; use tempfile::TempDir; fn create_request(name: &str) -> CreateContextRequest { @@ -166,4 +169,69 @@ mod tests { }; assert!(matches!(err, AppError::InvalidRequest(_))); } + + #[tokio::test] + async fn server_instances_write_contexts_to_distinct_wal_shards() { + let dir = TempDir::new().unwrap(); + let context_name = "shared"; + let state_a = Arc::new( + AppState::new_for_test_with_instance( + dir.path().to_path_buf(), + Some("context-0".to_string()), + ) + .await, + ); + + let (_, Json(_)) = + create_context(State(state_a.clone()), Json(create_request(context_name))) + .await + .unwrap(); + add_text_record(state_a, context_name, "from context-0").await; + + // A second server lazily opens the same dataset. Its instance id must + // select a different shard instead of fencing the first server's writer. + let state_b = Arc::new( + AppState::new_for_test_with_instance( + dir.path().to_path_buf(), + Some("context-1".to_string()), + ) + .await, + ); + add_text_record(state_b.clone(), context_name, "from context-1").await; + + let store = state_b + .get_or_open_context_store(context_name) + .await + .unwrap(); + assert_eq!(store.read().await.list(None, None).await.unwrap().len(), 2); + + let mem_wal = dir + .path() + .join(format!("{context_name}.lance")) + .join("_mem_wal"); + let shard_count = std::fs::read_dir(mem_wal) + .unwrap() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_ok_and(|kind| kind.is_dir())) + .count(); + assert_eq!(shard_count, 2, "each server instance must own one shard"); + } + + async fn add_text_record(state: Arc, context_name: &str, text: &str) { + let request = AddRecordsRequest { + records: vec![AddRecordRequest { + role: "user".to_string(), + content_type: "text/plain".to_string(), + text_payload: Some(text.to_string()), + ..Default::default() + }], + }; + let (_, Json(_)) = tokio::time::timeout( + Duration::from_secs(30), + add_records(State(state), Path(context_name.to_string()), Json(request)), + ) + .await + .expect("context write must not hang on shard contention") + .unwrap(); + } } diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 82cde80..b718589 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -54,8 +54,9 @@ pub struct AppState { pub rollout_registry: RwLock, pub base_uri: String, /// 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 - /// to a single shared shard (single-instance deployments only). + /// every server-managed store so each instance owns exactly one shard. + /// `None` falls back to a single shared shard (single-instance deployments + /// only). pub instance_id: Option, /// Count-triggered self-merge threshold for rollout MemWAL shards; `0` /// disables it. See `RolloutStoreOptions::merge_after_generations`. @@ -238,6 +239,14 @@ impl AppState { join_uri(&self.base_uri, &format!("{}.lance", name)) } + /// Options shared by every context-store open path in this server. + pub(crate) fn context_store_options(&self) -> ContextStoreOptions { + ContextStoreOptions { + shard_id: self.instance_id.clone(), + ..Default::default() + } + } + /// Build a default-configured `AppState` rooted at `base_path`, for tests. /// Opens a fresh registry under the directory; cleanup interval disabled. #[cfg(test)] @@ -426,7 +435,7 @@ impl AppState { } let uri = self.context_uri(name); - let opened = ContextStore::open_existing_with_options(&uri, ContextStoreOptions::default()) + let opened = ContextStore::open_existing_with_options(&uri, self.context_store_options()) .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened));