Cancel disconnected rollout reads and add explicit own-WAL recovery

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<tokio::sync::RwLock<Self>>) -> Option<JoinHandle<()>> {

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

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::<std::collections::BTreeSet<_>>(),
+                        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();
+            });
+    }
 }

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<()> {

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

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

@@ -61,6 +61,13 @@ pub struct ServerConfig {
     #[arg(long, env = "INSTANCE_ID")]
     pub instance_id: Option<String>,
 
+    /// 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<crate::replay::ReplayTarget>,
+
     /// 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

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<Arc<Drain>>,
-    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;