From bebf68d5c1d580c58e7fd5e7cbb41f4afda9c973 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Mon, 19 Oct 2026 12:53:21 +0000 Subject: [PATCH 1/2] Cancel disconnected rollout reads after protected store opening --- crates/lance-context-server/src/drain.rs | 25 +- .../src/drain/cancellation_tests.rs | 242 ++++++++++++++++++ crates/lance-context-server/src/error.rs | 7 + .../src/routes/rollouts.rs | 80 +++++- 4 files changed, 347 insertions(+), 7 deletions(-) create mode 100644 crates/lance-context-server/src/drain/cancellation_tests.rs diff --git a/crates/lance-context-server/src/drain.rs b/crates/lance-context-server/src/drain.rs index b9981e93..ebe9a8e8 100644 --- a/crates/lance-context-server/src/drain.rs +++ b/crates/lance-context-server/src/drain.rs @@ -157,14 +157,34 @@ impl Drain { } } +/// Opt-in cancellation for a route's audited read-only phase. Accepted setup +/// and mutations retain their existing task ownership and shutdown joining. +#[derive(Clone)] +pub(crate) struct ClientCancellation(tokio::sync::watch::Receiver<()>); +impl ClientCancellation { + pub(crate) fn channel() -> (tokio::sync::watch::Sender<()>, Self) { + let (sender, receiver) = tokio::sync::watch::channel(()); + (sender, Self(receiver)) + } + + pub(crate) async fn cancelled(&self) { + let mut receiver = self.0.clone(); + // The only sender belongs to the outer HTTP future. It never sends; + // its destruction closes the channel, including before this await. + let _ = receiver.changed().await; + } +} + pub(crate) async fn track_handler( State(drain): State>, - request: Request, + mut request: Request, next: Next, ) -> Response { let Some(handler) = drain.admit() else { return (StatusCode::SERVICE_UNAVAILABLE, "worker is draining").into_response(); }; + let (_request_lifetime, cancellation) = ClientCancellation::channel(); + request.extensions_mut().insert(cancellation); // Dropping a JoinHandle detaches it. A client disconnect therefore cannot // drop an accepted mutation halfway through a storage commit. The guard // remains in the task until success, error or panic. @@ -293,3 +313,6 @@ impl AsyncWrite for ConnectionIo { #[cfg(test)] mod tests; + +#[cfg(test)] +mod cancellation_tests; diff --git a/crates/lance-context-server/src/drain/cancellation_tests.rs b/crates/lance-context-server/src/drain/cancellation_tests.rs new file mode 100644 index 00000000..68859761 --- /dev/null +++ b/crates/lance-context-server/src/drain/cancellation_tests.rs @@ -0,0 +1,242 @@ +use super::*; +use axum::{ + body::Body, + extract::Extension, + middleware, + routing::{get, post}, +}; +use tokio::sync::RwLock; +use tower::ServiceExt; + +fn shield(app: Router, drain: Arc) -> Router { + app.layer(middleware::from_fn_with_state(drain, track_handler)) +} +fn read_route(lock: Arc>, entered: Arc) -> Router { + Router::new().route( + "/read", + get(move |Extension(cancel): Extension| { + let (lock, entered) = (lock.clone(), entered.clone()); + async move { + // Real integration must place this AFTER any protected lazy open. + tokio::select! { + biased; + _ = cancel.cancelled() => StatusCode::REQUEST_TIMEOUT, + response = async { + let _guard = lock.read_owned().await; + entered.notify_one(); + std::future::pending::<()>().await; + StatusCode::OK + } => response, + } + } + }), + ) +} +async fn bounded(f: impl Future) { + tokio::time::timeout(Duration::from_secs(2), f) + .await + .unwrap(); +} +#[tokio::test] +async fn opted_in_read_cancellation_releases_commit_lock_and_joins() { + let lock = Arc::new(RwLock::new(())); + let entered = Arc::new(Notify::new()); + let drain = Arc::new(Drain::default()); + let request = tokio::spawn( + shield(read_route(lock.clone(), entered.clone()), drain.clone()) + .oneshot(Request::get("/read").body(Body::empty()).unwrap()), + ); + bounded(entered.notified()).await; + assert!(lock.try_write().is_err()); + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + drain.close_admission(); + bounded(drain.joined()).await; + assert!(lock.try_write().is_ok()); +} +#[tokio::test] +async fn cancellation_during_protected_setup_waits_for_setup_then_skips_read() { + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let queried = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let drain = Arc::new(Drain::default()); + let app = Router::new().route( + "/read", + get({ + let (entered, release, queried) = (entered.clone(), release.clone(), queried.clone()); + move |Extension(cancel): Extension| { + let (entered, release, queried) = + (entered.clone(), release.clone(), queried.clone()); + async move { + entered.notify_one(); + release.notified().await; + tokio::select! { + biased; + _ = cancel.cancelled() => StatusCode::REQUEST_TIMEOUT, + response = async { + queried.store(true, std::sync::atomic::Ordering::SeqCst); + StatusCode::OK + } => response, + } + } + } + }), + ); + let request = tokio::spawn( + shield(app, drain.clone()).oneshot(Request::get("/read").body(Body::empty()).unwrap()), + ); + bounded(entered.notified()).await; + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + drain.close_admission(); + assert!( + tokio::time::timeout(Duration::from_millis(100), drain.joined()) + .await + .is_err() + ); + release.notify_one(); + bounded(drain.joined()).await; + assert!(!queried.load(std::sync::atomic::Ordering::SeqCst)); +} +#[tokio::test] +async fn accepted_post_remains_joined_and_commits_after_client_cancellation() { + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let committed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let drain = Arc::new(Drain::default()); + let app = Router::new().route( + "/write", + post({ + let (entered, release, committed) = + (entered.clone(), release.clone(), committed.clone()); + move || { + let (entered, release, committed) = + (entered.clone(), release.clone(), committed.clone()); + async move { + entered.notify_one(); + release.notified().await; + committed.store(true, std::sync::atomic::Ordering::SeqCst); + StatusCode::OK + } + } + }), + ); + let request = tokio::spawn( + shield(app, drain.clone()).oneshot(Request::post("/write").body(Body::empty()).unwrap()), + ); + bounded(entered.notified()).await; + request.abort(); + assert!(request.await.unwrap_err().is_cancelled()); + drain.close_admission(); + assert!( + tokio::time::timeout(Duration::from_millis(100), drain.joined()) + .await + .is_err() + ); + assert!(!committed.load(std::sync::atomic::Ordering::SeqCst)); + release.notify_one(); + bounded(drain.joined()).await; + assert!(committed.load(std::sync::atomic::Ordering::SeqCst)); +} +#[tokio::test] +async fn connected_read_is_not_cancelled_by_closing_admission() { + let drain = Arc::new(Drain::default()); + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let app = Router::new().route( + "/read", + get({ + let (entered, release) = (entered.clone(), release.clone()); + move |Extension(cancel): Extension| { + let (entered, release) = (entered.clone(), release.clone()); + async move { + entered.notify_one(); + tokio::select! { + _ = cancel.cancelled() => StatusCode::REQUEST_TIMEOUT, + _ = release.notified() => StatusCode::OK, + } + } + } + }), + ); + let request = tokio::spawn( + shield(app, drain.clone()).oneshot(Request::get("/read").body(Body::empty()).unwrap()), + ); + bounded(entered.notified()).await; + drain.close_admission(); + assert!( + tokio::time::timeout(Duration::from_millis(100), drain.joined()) + .await + .is_err() + ); + release.notify_one(); + assert_eq!(request.await.unwrap().unwrap().status(), StatusCode::OK); + bounded(drain.joined()).await; +} +#[tokio::test] +async fn panic_still_returns_500_and_releases_drain_guard() { + let drain = Arc::new(Drain::default()); + let app = Router::new().route( + "/panic", + get(|| async { + panic!("injected panic"); + #[allow(unreachable_code)] + StatusCode::OK + }), + ); + let response = shield(app, drain.clone()) + .oneshot(Request::get("/panic").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + drain.close_admission(); + bounded(drain.joined()).await; +} +#[tokio::test] +async fn closed_admission_still_rejects_new_requests() { + let drain = Arc::new(Drain::default()); + drain.close_admission(); + let app = Router::new().route("/read", get(|| async { StatusCode::OK })); + let response = shield(app, drain.clone()) + .oneshot(Request::get("/read").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + bounded(drain.joined()).await; +} +#[tokio::test] +async fn actual_http_disconnect_cancels_opted_in_read() { + use tokio::io::AsyncWriteExt; + let lock = Arc::new(RwLock::new(())); + let entered = Arc::new(Notify::new()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (stop, stopping) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(serve( + listener, + read_route(lock.clone(), entered.clone()), + Duration::from_millis(10), + async { + let _ = stopping.await; + }, + )); + let mut client = TcpStream::connect(address).await.unwrap(); + client + .write_all(b"GET /read HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + bounded(entered.notified()).await; + let client = client.into_std().unwrap(); + client.shutdown(Shutdown::Both).unwrap(); + drop(client); + bounded(async { + let _writer = lock.write().await; + }) + .await; + stop.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(2), server) + .await + .unwrap() + .unwrap() + .unwrap(); +} diff --git a/crates/lance-context-server/src/error.rs b/crates/lance-context-server/src/error.rs index 2cd4809d..da73c71b 100644 --- a/crates/lance-context-server/src/error.rs +++ b/crates/lance-context-server/src/error.rs @@ -11,6 +11,8 @@ pub enum AppError { InvalidRequest(String), Internal(String), CompactionInProgress, + /// An opted-in read stopped after its outer HTTP future was dropped. + ClientDisconnected, /// The in-flight blob-byte budget is exhausted; the client should retry /// later. Maps to `503 Service Unavailable`. Overloaded(String), @@ -65,6 +67,11 @@ impl AppError { impl IntoResponse for AppError { fn into_response(self) -> Response { let (status, code, message) = match self { + AppError::ClientDisconnected => ( + StatusCode::from_u16(499).expect("499 is a valid HTTP status"), + "CLIENT_DISCONNECTED", + "client disconnected".to_string(), + ), AppError::NotFound(msg) => (StatusCode::NOT_FOUND, "NOT_FOUND", msg), AppError::AlreadyExists(msg) => (StatusCode::CONFLICT, "ALREADY_EXISTS", msg), AppError::InvalidRequest(msg) => (StatusCode::BAD_REQUEST, "INVALID_REQUEST", msg), diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 2b39f7ed..aabe57c8 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -5,7 +5,7 @@ use axum::body::Body; use axum::extract::{FromRequest, Multipart, Path, Query, Request, State}; use axum::http::{header, StatusCode}; use axum::response::Response; -use axum::Json; +use axum::{Extension, Json}; use lance_context_api::{ AddRolloutRequest, AddRolloutsRequest, AddRolloutsResponse, CheckoutRequest, CompactRequest, CompactResponse, CompactStatsResponse, CreateRolloutStoreRequest, GetRolloutResponse, @@ -17,6 +17,7 @@ use lance_context_core::{ }; use tokio::sync::RwLock; +use crate::drain::ClientCancellation; use crate::error::AppError; use crate::state::AppState; @@ -508,6 +509,7 @@ pub struct RolloutListParams { } pub async fn list_rollouts( + cancellation: Option>, State(state): State>, Path(name): Path, Query(params): Query, @@ -525,11 +527,25 @@ pub async fn list_rollouts( let store_lock = state.get_or_open_rollout_store(&name).await?; - let store = store_lock.read().await; - let records = store - .list_with_filters(params.limit, params.offset, filters.as_ref()) - .await - .map_err(AppError::from_lance)?; + // Opening a store stays inside the accepted handler's protected lifetime. + // Only this read-only phase, including waiting for the table lock, opts in + // to client cancellation. Dropping it releases the read guard so a merge + // can acquire its commit lock; accepted mutations remain shielded. + let read = async { + let store = store_lock.read().await; + store + .list_with_filters(params.limit, params.offset, filters.as_ref()) + .await + .map_err(AppError::from_lance) + }; + let records = match cancellation { + Some(Extension(cancellation)) => tokio::select! { + biased; + _ = cancellation.cancelled() => return Err(AppError::ClientDisconnected), + records = read => records?, + }, + None => read.await?, + }; Ok(Json(ListRolloutsResponse { records: records.into_iter().map(rollout_record_to_dto).collect(), @@ -1416,8 +1432,55 @@ mod tests { .unwrap() } + #[tokio::test] + async fn cancelled_list_stops_waiting_for_the_table_read_lock() { + let (state, _dir) = rollout_state().await; + let store = state.get_or_open_rollout_store("rl").await.unwrap(); + let write_guard = store.write().await; + let (lifetime, cancellation) = ClientCancellation::channel(); + let read = list_rollouts( + Some(Extension(cancellation)), + State(state.clone()), + Path("rl".to_string()), + Query(RolloutListParams::default()), + ); + tokio::pin!(read); + assert!(futures::poll!(&mut read).is_pending()); + drop(lifetime); + let result = tokio::time::timeout(std::time::Duration::from_secs(2), read) + .await + .unwrap(); + assert!(matches!(result, Err(AppError::ClientDisconnected))); + drop(write_guard); + assert!(store.try_write().is_ok()); + } + + #[tokio::test] + async fn cancelled_list_keeps_store_lookup_protected() { + let (state, _dir) = rollout_state().await; + let cache_guard = state.rollout_stores.lock().await; + let (lifetime, cancellation) = ClientCancellation::channel(); + let read = list_rollouts( + Some(Extension(cancellation)), + State(state.clone()), + Path("rl".to_string()), + Query(RolloutListParams::default()), + ); + tokio::pin!(read); + assert!(futures::poll!(&mut read).is_pending()); + drop(lifetime); + // Cancellation cannot drop lazy store opening or its protected setup. + assert!(futures::poll!(&mut read).is_pending()); + drop(cache_guard); + let result = tokio::time::timeout(std::time::Duration::from_secs(2), read) + .await + .unwrap(); + assert!(matches!(result, Err(AppError::ClientDisconnected))); + } + async fn count_rollouts(state: &Arc) -> usize { let Json(list) = list_rollouts( + None, State(state.clone()), Path("rl".to_string()), Query(RolloutListParams::default()), @@ -1674,6 +1737,7 @@ mod tests { // Read routed to B must lazily open the store, not 404. let Json(list) = list_rollouts( + None, State(state_b.clone()), Path("rl".to_string()), Query(RolloutListParams::default()), @@ -1706,6 +1770,7 @@ mod tests { async fn read_of_nonexistent_store_still_404s() { let (state, _dir) = rollout_state().await; let err = list_rollouts( + None, State(state.clone()), Path("no-such-store".to_string()), Query(RolloutListParams::default()), @@ -1814,6 +1879,7 @@ mod tests { flush_store(&state, "rl").await; let Json(response) = list_rollouts( + None, State(state), Path("rl".to_string()), Query(RolloutListParams { @@ -1838,6 +1904,7 @@ mod tests { async fn list_rollouts_rejects_invalid_filters() { let (state, _dir) = rollout_state().await; let invalid_json = list_rollouts( + None, State(state.clone()), Path("rl".to_string()), Query(RolloutListParams { @@ -1850,6 +1917,7 @@ mod tests { assert!(matches!(invalid_json, AppError::InvalidRequest(_))); let unsupported = list_rollouts( + None, State(state), Path("rl".to_string()), Query(RolloutListParams { From 2e239a4c1224cc0975fd1a500bf166dfbb6a7092 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Mon, 19 Oct 2026 16:46:58 +0000 Subject: [PATCH 2/2] Add explicit one-shot WAL replay for retired instance recovery --- .../lance-context-core/src/datagen_store.rs | 7 + .../lance-context-core/src/generic_store.rs | 7 + .../lance-context-core/src/rollout_store.rs | 127 +++++++++++++++++ crates/lance-context-core/src/store.rs | 7 + crates/lance-context-core/src/store_base.rs | 15 ++ crates/lance-context-server/src/config.rs | 7 + crates/lance-context-server/src/main.rs | 16 +++ crates/lance-context-server/src/replay.rs | 132 ++++++++++++++++++ 8 files changed, 318 insertions(+) create mode 100644 crates/lance-context-server/src/replay.rs diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index cea5156d..7aa68e31 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -438,6 +438,13 @@ impl DatagenStore { self.base.flush().await } + /// Initialize/replay and materialize this instance's WAL after the caller + /// has established ownership and retired the previous process with this + /// identity. Fencing and flush errors are returned to the caller. + pub async fn replay_and_flush_own_wal(&self) -> LanceResult<()> { + self.base.replay_and_flush_own_wal().await + } + /// Start periodic cleanup after the configured interval. The task keeps /// only a weak reference and exits when the store is dropped. pub fn spawn_periodic_cleanup(store: Arc>) -> Option> { diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 8798d075..497c426e 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -598,6 +598,13 @@ impl GenericStore { self.base.flush().await } + /// Initialize/replay and materialize this instance's WAL after the caller + /// has established ownership and retired the previous process with this + /// identity. Fencing and flush errors are returned to the caller. + pub async fn replay_and_flush_own_wal(&self) -> LanceResult<()> { + self.base.replay_and_flush_own_wal().await + } + /// Close the resident writer, draining its background tasks. Idempotent. pub async fn close(&mut self) -> LanceResult<()> { self.base.close().await diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 5564b3c2..5bda1c63 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -625,6 +625,15 @@ impl RolloutStore { self.base.flush().await } + /// Replay and materialize durable WAL for this instance without appending rows. + /// + /// This opens the instance's writer even on an idle table. The caller must + /// establish shard ownership and retire its previous writer process first; + /// this method is not an ownership or uncertain-operation retry protocol. + pub async fn replay_and_flush_own_wal(&self) -> LanceResult<()> { + self.base.replay_and_flush_own_wal().await + } + /// Gracefully close the resident writer, draining its background tasks. /// Idempotent. See `StorageBase::close`. pub async fn close(&mut self) -> LanceResult<()> { @@ -6194,4 +6203,122 @@ mod tests { } } } + + // Child process deliberately exits without dropping the store or runtime. + // The parent verifies acknowledged idle WAL after that exact process joins. + #[test] + #[ignore = "subprocess helper for explicit_replay_materializes_idle_acknowledged_wal"] + fn explicit_replay_crash_child() { + use std::io::Write; + let root = std::path::PathBuf::from(std::env::var("RK_REPLAY_CRASH_ROOT_20182").unwrap()); + assert!(root.is_dir()); + let uri = root.join("table").to_string_lossy().into_owned(); + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap() + .block_on(async { + let store = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("explicit-replay-crashed-owner".into()), + merge_after_generations: Some(0), + ..Default::default() + }, + ) + .await + .unwrap(); + for i in 0..20 { + store + .add(&[assistant_record(&format!("ack-{i:02}"))]) + .await + .unwrap(); + } + assert!(store.list(None, None).await.unwrap().is_empty()); + let mut receipt = std::fs::File::create(root.join("acknowledged.json")).unwrap(); + receipt + .write_all(b"{\"acknowledged\":20,\"flushed\":false}") + .unwrap(); + receipt.sync_all().unwrap(); + std::process::exit(0); + }); + } + + #[test] + fn explicit_replay_materializes_idle_acknowledged_wal() { + let dir = TempDir::new().unwrap(); + let mut child = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "rollout_store::tests::explicit_replay_crash_child", + "--ignored", + "--nocapture", + "--test-threads=1", + ]) + .env("RK_REPLAY_CRASH_ROOT_20182", dir.path()) + .spawn() + .unwrap(); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + let status = loop { + if let Some(status) = child.try_wait().unwrap() { + break status; + } + if std::time::Instant::now() >= deadline { + child.kill().unwrap(); + let _ = child.wait(); + panic!("owned replay-test child exceeded its deadline"); + } + std::thread::sleep(std::time::Duration::from_millis(10)); + }; + assert!(status.success()); + let receipt: serde_json::Value = + serde_json::from_slice(&std::fs::read(dir.path().join("acknowledged.json")).unwrap()) + .unwrap(); + assert_eq!(receipt["acknowledged"], 20); + assert_eq!(receipt["flushed"], false); + let uri = dir.path().join("table").to_string_lossy().into_owned(); + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap() + .block_on(async { + let mut store = RolloutStore::open_existing_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("explicit-replay-crashed-owner".into()), + merge_after_generations: Some(0), + ..Default::default() + }, + ) + .await + .unwrap(); + // Existing flush/merge do not instantiate the idle writer. + store.flush().await.unwrap(); + assert_eq!(store.cleanup_own_shard().await.unwrap(), 0); + assert!(store.list(None, None).await.unwrap().is_empty()); + // No new append is issued after the child exits. + store.replay_and_flush_own_wal().await.unwrap(); + let expected: std::collections::BTreeSet<_> = + (0..20).map(|i| format!("ack-{i:02}")).collect(); + for _ in 0..2 { + let rows = store.list(None, None).await.unwrap(); + assert_eq!(rows.len(), 20); + assert_eq!( + rows.iter() + .map(|r| r.id.clone()) + .collect::>(), + expected + ); + for row in &rows { + assert_eq!(row.content, assistant_record(&row.id).content); + } + store.replay_and_flush_own_wal().await.unwrap(); + } + assert!(store.cleanup_own_shard().await.unwrap() > 0); + assert_eq!(store.list(None, None).await.unwrap().len(), 20); + store.close().await.unwrap(); + }); + } } diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 184418a6..a4014368 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -2118,6 +2118,13 @@ impl ContextStore { self.base.flush().await } + /// Initialize/replay and materialize this instance's WAL after the caller + /// has established ownership and retired the previous process with this + /// identity. Fencing and flush errors are returned to the caller. + pub async fn replay_and_flush_own_wal(&self) -> LanceResult<()> { + self.base.replay_and_flush_own_wal().await + } + /// Gracefully close the resident MemWAL writer, draining its background /// tasks and sealing whatever it still buffers. Idempotent. pub async fn close(&mut self) -> LanceResult<()> { diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 06bfaf52..371697ed 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -732,6 +732,21 @@ impl StorageBase { } } + /// Explicitly initialize/replay this instance's writer and materialize its WAL. + /// + /// Unlike `flush`, this opens a writer when none is resident. Opening claims + /// the shard epoch: the caller must establish ownership and retire any + /// previous process using the same instance identity before recovery. + /// This operation does not itself establish remote retirement or authorize + /// retrying an uncertain mutation. Fencing and drain errors are propagated. + pub async fn replay_and_flush_own_wal(&self) -> LanceResult<()> { + self.ensure_writable()?; + let writer = self.resident_writer().await?; + writer.force_seal_active().await?; + writer.wait_for_flush_drain().await?; + Ok(()) + } + /// Materialize the active memtable into a flushed, queryable generation: /// `force_seal_active` freezes it and hands it to the background flusher, /// then `wait_for_flush_drain` blocks until every frozen memtable has landed diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 91622e87..c577d913 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -61,6 +61,13 @@ pub struct ServerConfig { #[arg(long, env = "INSTANCE_ID")] pub instance_id: Option, + /// Replay and flush one existing store's WAL, then exit without serving HTTP. + /// Requires the original instance identity and verified retirement of its + /// previous process. TARGET is [rollout:|generic:|context:|datagen:]NAME. + /// This mode does not establish ownership or resolve uncertain operations. + #[arg(long, requires = "instance_id")] + pub replay_own_wal: Option, + /// Count-triggered self-merge threshold for rollout MemWAL shards. When this /// instance's shard has at least this many un-merged flushed generations, the /// periodic sweeper folds them into the base table and drains the shard (see diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index b00c4f71..83a9a1bd 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -3,6 +3,7 @@ mod config; mod drain; mod error; mod merge_execution; +mod replay; mod routes; mod state; mod sweeper; @@ -26,6 +27,11 @@ async fn main() { .init(); let config = ServerConfig::parse(); + if let Err(e) = replay::validate_config(&config) { + tracing::error!("Invalid WAL replay configuration: {e}"); + std::process::exit(2); + } + let replay_target = config.replay_own_wal.clone(); let body_idle_timeout = std::time::Duration::from_secs(config.request_body_idle_timeout_secs); let response_idle = std::time::Duration::from_secs(config.http_shutdown_response_idle_secs); let addr = format!("{}:{}", config.host, config.port); @@ -47,6 +53,16 @@ async fn main() { } }; + // Recovery is a one-shot operation: never start listeners or sweepers. + if let Some(target) = replay_target { + if let Err(e) = replay::run(&state, &target).await { + tracing::error!(target = ?target, error = ?e, "WAL replay failed"); + std::process::exit(1); + } + tracing::info!(target = ?target, "WAL replay and writer close completed"); + return; + } + // 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. diff --git a/crates/lance-context-server/src/replay.rs b/crates/lance-context-server/src/replay.rs new file mode 100644 index 00000000..61ad0a53 --- /dev/null +++ b/crates/lance-context-server/src/replay.rs @@ -0,0 +1,132 @@ +//! Explicit materialization after the caller has retired the previous owner. +//! No ownership transfer, claim reset, or uncertain-operation retry happens here. + +use std::str::FromStr; + +use crate::{config::ServerConfig, error::AppError, state::AppState}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ReplayTarget { + Rollout(String), + Generic(String), + Context(String), + Datagen(String), +} + +impl FromStr for ReplayTarget { + type Err = String; + + fn from_str(value: &str) -> Result { + let (kind, name) = value.split_once(':').unwrap_or(("rollout", value)); + AppState::validate_name(name).map_err(|e| format!("invalid replay target: {e:?}"))?; + match kind { + "rollout" => Ok(Self::Rollout(name.to_owned())), + "generic" => Ok(Self::Generic(name.to_owned())), + "context" => Ok(Self::Context(name.to_owned())), + "datagen" => Ok(Self::Datagen(name.to_owned())), + _ => Err(format!("unsupported replay store kind: {kind}")), + } + } +} + +pub fn validate_config(config: &ServerConfig) -> Result<(), &'static str> { + if config.replay_own_wal.is_some() + && config + .instance_id + .as_deref() + .is_none_or(|id| id.trim().is_empty()) + { + return Err("--replay-own-wal requires a nonempty explicit --instance-id or INSTANCE_ID; HOSTNAME fallback is forbidden"); + } + Ok(()) +} + +pub async fn run(state: &AppState, target: &ReplayTarget) -> Result<(), AppError> { + // Hold the target exclusively through replay and close. Propagate close + // errors too: AppState::shutdown's log-only policy is unsuitable for a + // one-shot recovery command's success receipt. + macro_rules! replay { + ($open:ident, $name:expr) => {{ + let store = state.$open($name).await?; + let mut store = store.write().await; + store + .replay_and_flush_own_wal() + .await + .map_err(AppError::from_lance)?; + store.close().await.map_err(AppError::from_lance) + }}; + } + match target { + ReplayTarget::Rollout(name) => replay!(get_or_open_rollout_store, name), + ReplayTarget::Generic(name) => replay!(get_or_open_generic_store, name), + ReplayTarget::Context(name) => replay!(get_or_open_context_store, name), + ReplayTarget::Datagen(name) => replay!(get_or_open_datagen_store, name), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use clap::Parser; + + #[test] + fn replay_target_rejects_ambiguous_or_unsafe_names() { + assert_eq!("table".parse(), Ok(ReplayTarget::Rollout("table".into()))); + for kind in ["rollout", "generic", "context", "datagen"] { + assert!(format!("{kind}:table").parse::().is_ok()); + } + for invalid in [ + "", + "rollout:", + "unknown:table", + "generic:../x", + "rollout:a:b", + "/absolute", + ] { + assert!(invalid.parse::().is_err(), "{invalid}"); + } + } + + #[test] + fn replay_requires_explicit_nonempty_identity() { + let mut config = ServerConfig::try_parse_from([ + "server", + "--replay-own-wal", + "table", + "--instance-id", + "owner-15", + ]) + .unwrap(); + assert!(validate_config(&config).is_ok()); + // Direct validation also forbids HOSTNAME fallback regardless of the + // invoking test environment's INSTANCE_ID. + config.instance_id = None; + assert!(validate_config(&config).is_err()); + config.instance_id = Some(" ".into()); + assert!(validate_config(&config).is_err()); + config.replay_own_wal = None; + assert!(validate_config(&config).is_ok()); + } + + #[tokio::test] + async fn replay_missing_target_is_an_error_without_creating_it() { + let dir = tempfile::tempdir().unwrap(); + let config = ServerConfig::try_parse_from([ + "server", + "--data-dir", + dir.path().to_str().unwrap(), + "--instance-id", + "test-owner", + ]) + .unwrap(); + let state = AppState::new(config).await.unwrap(); + for kind in ["rollout", "generic", "context", "datagen"] { + let target = format!("{kind}:absent").parse().unwrap(); + assert!(run(&state, &target).await.is_err(), "{kind}"); + } + assert!(!dir.path().join("absent.lance").exists()); + assert!(!dir.path().join("absent.rollout.lance").exists()); + assert!(!dir.path().join("absent.generic.lance").exists()); + assert!(!dir.path().join("absent.datagen.lance").exists()); + } +}