fix: join admitted master work during executor drain

crates/lance-context-master/src/admission.rs

@@ -0,0 +1,113 @@
+//! Executor-local admission and joined shutdown, independent of table ownership.
+use serde::Serialize;
+use std::sync::Arc;
+use tokio::sync::watch;
+
+#[derive(Clone, Debug, Serialize)]
+pub struct Status {
+    pub executor_id: String,
+    pub accepting: bool,
+    /// Includes claim RPCs in flight, admitted tasks, and scanner passes.
+    pub active_operations: usize,
+}
+
+impl Status {
+    pub fn drained(&self) -> bool {
+        !self.accepting && self.active_operations == 0
+    }
+}
+
+pub struct Admission {
+    state: watch::Sender<Status>,
+}
+
+impl Default for Admission {
+    fn default() -> Self {
+        let (state, _) = watch::channel(Status {
+            executor_id: lance_context_core::generate_id(),
+            accepting: true,
+            active_operations: 0,
+        });
+        Self { state }
+    }
+}
+
+impl Admission {
+    pub fn status(&self) -> Status {
+        self.state.borrow().clone()
+    }
+
+    /// The same lock closes admission and reserves an operation before its
+    /// first ownership RPC. Drain cannot miss an accepted but delayed claim.
+    pub fn try_admit(self: &Arc<Self>) -> Option<Operation> {
+        let mut admitted = false;
+        self.state.send_if_modified(|state| {
+            if state.accepting {
+                state.active_operations += 1;
+                admitted = true;
+            }
+            admitted
+        });
+        admitted.then(|| Operation(self.clone()))
+    }
+
+    pub fn begin_drain(&self) -> Status {
+        self.state.send_modify(|state| state.accepting = false);
+        self.status()
+    }
+
+    /// No elapsed-time cancellation: admitted writers retain their existing
+    /// progress watchdog, lease and fencing protocol until they return.
+    pub async fn wait_drained(&self) {
+        let mut updates = self.state.subscribe();
+        updates
+            .wait_for(Status::drained)
+            .await
+            .expect("admission sender remains alive");
+    }
+}
+
+pub struct Operation(Arc<Admission>);
+
+impl Drop for Operation {
+    fn drop(&mut self) {
+        self.0
+            .state
+            .send_modify(|state| state.active_operations -= 1);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn drain_waits_for_admitted_claim_and_rejects_new_work() {
+        let admission = Arc::new(Admission::default());
+        let claim_in_flight = admission.try_admit().unwrap();
+        let running_writer = admission.try_admit().unwrap();
+        let snapshot = admission.begin_drain();
+        assert!(!snapshot.accepting);
+        assert_eq!(snapshot.active_operations, 2);
+        assert!(admission.try_admit().is_none());
+        let waiter = tokio::spawn({
+            let admission = admission.clone();
+            async move { admission.wait_drained().await }
+        });
+        drop(claim_in_flight);
+        tokio::task::yield_now().await;
+        assert!(!waiter.is_finished());
+        drop(running_writer);
+        waiter.await.unwrap();
+        assert!(admission.status().drained());
+        assert_eq!(admission.begin_drain().executor_id, snapshot.executor_id);
+        assert!(admission.try_admit().is_none());
+    }
+
+    #[tokio::test]
+    async fn drained_notification_before_subscription_is_not_lost() {
+        let admission = Arc::new(Admission::default());
+        admission.begin_drain();
+        admission.wait_drained().await;
+    }
+}

crates/lance-context-master/src/catchup/mod.rs

@@ -236,6 +236,9 @@ async fn admit(
     row: Option<&StatRow>,
     request: &Trigger,
 ) -> Result<Decision> {
+    let Some(_operation) = state.admission.try_admit() else {
+        return Ok(Decision::new(&request.target, "executor_draining"));
+    };
     let now = chrono::Utc::now().timestamp_millis();
     let target = &request.target;
     if let Some(reason) = eligibility(&state.config, row, target, now) {
@@ -327,6 +330,9 @@ pub fn spawn(state: &Arc<MasterState>) -> Option<tokio::task::JoinHandle<()>> {
     }))
 }
 async fn tick(state: &Arc<MasterState>, cursor: &mut usize) -> Result<()> {
+    let Some(_operation) = state.admission.try_admit() else {
+        return Ok(());
+    };
     let inventory = Inventory::new(state);
     inventory.ensure_policy(&state.config).await?;
     let kube = Kubernetes::in_cluster(&state.config.catchup)?;

crates/lance-context-master/src/lib.rs

@@ -1,5 +1,6 @@
 //! Control-plane (master) library surface.
 
+pub mod admission;
 pub mod catchup;
 pub mod config;
 pub mod discovery;

crates/lance-context-master/src/main.rs

@@ -68,6 +68,7 @@ async fn main() {
         app = app.fallback_service(serve);
     }
 
+    let admission = state.admission.clone();
     let app = app
         .with_state(state)
         .merge(lance_context_metrics::metrics_router(metrics_handle))
@@ -80,7 +81,14 @@ async fn main() {
 
     let listener = TcpListener::bind(&addr).await.unwrap();
     axum::serve(listener, app)
-        .with_graceful_shutdown(wait_for_shutdown_signal())
+        .with_graceful_shutdown(async move {
+            wait_for_shutdown_signal().await;
+            let status = admission.begin_drain();
+            tracing::info!(executor_id = %status.executor_id, active = status.active_operations,
+                "stopped local admission; waiting for admitted work");
+            admission.wait_drained().await;
+            tracing::info!("local executor drained");
+        })
         .await
         .unwrap();
 }

crates/lance-context-master/src/routes.rs

@@ -695,9 +695,32 @@ pub async fn get_task(
     }
 }
 
+/// Call through the exact Pod endpoint. The process identity prevents a load
+/// balancer or replacement Pod from draining a different executor by mistake.
+#[derive(Deserialize)]
+pub struct DrainRequest {
+    pub executor_id: String,
+}
+
+async fn executor_status(State(state): State<Arc<MasterState>>) -> Json<crate::admission::Status> {
+    Json(state.admission.status())
+}
+
+async fn drain_executor(
+    State(state): State<Arc<MasterState>>,
+    Json(request): Json<DrainRequest>,
+) -> Result<Json<crate::admission::Status>, StatusCode> {
+    if state.admission.status().executor_id != request.executor_id {
+        return Err(StatusCode::CONFLICT);
+    }
+    Ok(Json(state.admission.begin_drain()))
+}
+
 /// Build the admin API router (mounted under `/api/v1`).
 pub fn api_router() -> Router<Arc<MasterState>> {
     Router::new()
+        .route("/executor", get(executor_status))
+        .route("/executor/drain", post(drain_executor))
         .route("/experiments", get(list_experiments))
         .route("/experiments/{name}", get(get_experiment))
         .route("/experiments/{name}/records", get(list_experiment_records))

crates/lance-context-master/src/scanner.rs

@@ -208,6 +208,10 @@ async fn maintain_registry(
 /// for experiments no longer in the registry. Returns the number of
 /// experiments successfully observed.
 pub async fn scan_once(state: &Arc<MasterState>) -> lance::Result<usize> {
+    let _operation = state
+        .admission
+        .try_admit()
+        .ok_or_else(|| lance::Error::io("master executor is draining; stats scan not admitted"))?;
     let guard = state.task_store.coordination_lock("stats-writer").await?;
     let result = scan_once_inner(state).await;
     let release = state.task_store.release_coordination_lock(guard).await;
@@ -219,6 +223,9 @@ pub async fn scan_once(state: &Arc<MasterState>) -> lance::Result<usize> {
 }
 
 async fn try_scan_once(state: &Arc<MasterState>, maintain: bool) -> lance::Result<Option<usize>> {
+    let Some(_operation) = state.admission.try_admit() else {
+        return Ok(None);
+    };
     let Some(guard) = state
         .task_store
         .try_coordination_lock("stats-writer")

crates/lance-context-master/src/scheduler.rs

@@ -1084,6 +1084,9 @@ fn spawn_pool_poller(
                 let Ok(permit) = pool.clone().try_acquire_owned() else {
                     break;
                 };
+                let Some(operation) = state.admission.try_admit() else {
+                    break;
+                };
                 let claim_start = std::time::Instant::now();
                 match state.task_store.claim_next_of_kinds(kinds).await {
                     Ok(Some(claim)) => {
@@ -1099,6 +1102,7 @@ fn spawn_pool_poller(
                             run_task(&st, claim, timing).await;
                             drop(permit);
                             drop(compact_permit);
+                            drop(operation);
                         });
                     }
                     Ok(None) => break,
@@ -2102,6 +2106,100 @@ mod tests {
             )
     }
 
+    #[tokio::test]
+    #[ignore = "requires ETCD_TEST_ENDPOINTS"]
+    async fn executor_drain_finishes_current_merge_and_leaves_queue_for_other_master() {
+        use axum::{routing::post, Json, Router};
+        let started = Arc::new(tokio::sync::Notify::new());
+        let release = Arc::new(Semaphore::new(0));
+        let app = Router::new().route(
+            "/api/v1/internal/merge-wal/{name}",
+            post({
+                let started = started.clone();
+                let release = release.clone();
+                move || {
+                    let started = started.clone();
+                    let release = release.clone();
+                    async move {
+                        started.notify_one();
+                        release.acquire().await.unwrap().forget();
+                        Json(serde_json::json!({"reclaimed": 3}))
+                    }
+                }
+            }),
+        );
+        let dir = TempDir::new().unwrap();
+        let mut cfg = config(&dir);
+        cfg.merge_wal_concurrency = 1;
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        cfg.worker_endpoints = vec![format!("http://{}", listener.local_addr().unwrap())];
+        let app = owned_stub(app, &cfg).await;
+        let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+        let state = MasterState::new(cfg.clone()).await.unwrap();
+        let worker = spawn_scheduler(&state);
+        let current = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap();
+        tokio::time::timeout(Duration::from_secs(10), started.notified())
+            .await
+            .unwrap();
+        use tower::ServiceExt;
+        let api = crate::routes::api_router().with_state(state.clone());
+        let request = |id: &str| {
+            axum::http::Request::builder()
+                .method("POST")
+                .uri("/executor/drain")
+                .header("content-type", "application/json")
+                .body(axum::body::Body::from(
+                    serde_json::json!({"executor_id":id}).to_string(),
+                ))
+                .unwrap()
+        };
+        let wrong = api
+            .clone()
+            .oneshot(request("previous-process"))
+            .await
+            .unwrap();
+        assert_eq!(wrong.status(), axum::http::StatusCode::CONFLICT);
+        assert!(state.admission.status().accepting);
+        let id = state.admission.status().executor_id;
+        let response = api.oneshot(request(&id)).await.unwrap();
+        assert_eq!(response.status(), axum::http::StatusCode::OK);
+        let draining = state.admission.status();
+        assert!(!draining.accepting);
+        assert!(draining.active_operations > 0);
+        let next = enqueue(&state, TaskKind::MergeWal, "exp-1").await.unwrap();
+        release.add_permits(2);
+        tokio::time::timeout(Duration::from_secs(10), state.admission.wait_drained())
+            .await
+            .unwrap();
+        assert_eq!(
+            state
+                .task_store
+                .get(&current.id)
+                .await
+                .unwrap()
+                .unwrap()
+                .state,
+            TaskState::Done
+        );
+        assert_eq!(
+            state.task_store.get(&next.id).await.unwrap().unwrap().state,
+            TaskState::Queued
+        );
+        assert!(state.admission.try_admit().is_none());
+        assert!(crate::scanner::scan_once(&state).await.is_err());
+
+        let other = MasterState::new(cfg).await.unwrap();
+        let other_worker = spawn_scheduler(&other);
+        assert_eq!(
+            await_terminal(&other, &next.id).await.state,
+            TaskState::Done
+        );
+        assert!(state.admission.status().drained());
+        worker.abort();
+        other_worker.abort();
+        server.abort();
+    }
+
     /// MergeWal fans out to every configured worker endpoint and sums the
     /// reclaimed counts. Uses a tiny in-process stub server per "worker".
     #[tokio::test]

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

@@ -103,6 +103,8 @@ pub struct MasterState {
     /// Durable scheduler queue in etcd, providing shared CAS/lease semantics
     /// for stateless HA master replicas.
     pub task_store: TaskStore,
+    /// Local claim/scan admission, including in-flight ownership RPCs.
+    pub admission: Arc<crate::admission::Admission>,
     /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints.
     pub http: reqwest::Client,
     /// Process-wide compaction permits shared by scheduler and retirement work.
@@ -212,6 +214,7 @@ impl MasterState {
             base_uri,
             config,
             task_store,
+            admission: Arc::new(crate::admission::Admission::default()),
             http: reqwest::Client::new(),
             compaction_permits: Arc::new(Semaphore::new(compaction_concurrency)),
             stats_maintenance_failures: std::sync::atomic::AtomicU64::new(0),

docs/operations/master-drain.md

@@ -0,0 +1,32 @@
+# Drain one master executor
+
+New masters expose `GET /api/v1/executor` and
+`POST /api/v1/executor/drain` on the admin API. Use the exact Pod endpoint,
+read its `executor_id`, and include that identity in the POST JSON. A request
+for a different process returns HTTP 409. Draining is monotonic for that
+process; restarting creates a fresh identity and opens local admission.
+
+After POST, poll that same Pod and identity until `accepting` is false and
+`active_operations` is zero. Only then terminate the Pod. Operations are
+reserved before ownership RPCs, so an in-flight accepted claim is included.
+Existing tasks, scanner passes (including cold retirement) and catch-up
+controller operations finish normally. New scheduler claims, scanner passes
+and catch-up provisioning on that process are refused. Queued tasks remain
+available to other masters; table owners and fences are not reset.
+
+SIGTERM/SIGINT closes the same gate and waits for admitted work before HTTP
+shutdown and runtime exit. Existing progress watchdogs remain responsible for
+stalled operations; drain adds no elapsed-time writer cancellation. Kubernetes
+can still SIGKILL after its termination grace period: explicitly drain through
+the API before deletion instead of relying on a short grace period.
+
+This interface is process-local. Enqueue/read requests can still be served,
+and other masters and already launched dedicated Jobs continue. Dedicated
+one-table executor mode has no admin listener; use its supervisor's joined
+pass lifecycle. The endpoint is part of the existing administrative API and
+requires the same network access controls.
+
+Older images do not implement this protocol. Observing zero tasks and then
+deleting an old Pod is not an atomic drain, and installing this release does
+not retroactively add the endpoint to running old processes. Do not claim an
+old executor is drained merely from an idle metrics sample or SIGTERM receipt.