fix(server): shard context writers by instance

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<String>,
@@ -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<String> {
         self.instance_id

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<AppState>, 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();
+    }
 }

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

@@ -54,8 +54,9 @@ pub struct AppState {
     pub rollout_registry: RwLock<RolloutRegistry>,
     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<String>,
     /// 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));