From c36b83f5abede75a73bdf59325f58360138a7dea Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Mon, 5 Oct 2026 20:41:13 +0000 Subject: [PATCH] fix: drain local master admission before stopping admitted work --- crates/lance-context-master/src/admission.rs | 113 ++++++++++++++++++ .../lance-context-master/src/catchup/mod.rs | 6 + crates/lance-context-master/src/lib.rs | 1 + crates/lance-context-master/src/main.rs | 10 +- crates/lance-context-master/src/routes.rs | 23 ++++ crates/lance-context-master/src/scanner.rs | 7 ++ crates/lance-context-master/src/scheduler.rs | 98 +++++++++++++++ crates/lance-context-master/src/state.rs | 3 + docs/operations/master-drain.md | 32 +++++ 9 files changed, 292 insertions(+), 1 deletion(-) create mode 100644 crates/lance-context-master/src/admission.rs create mode 100644 docs/operations/master-drain.md diff --git a/crates/lance-context-master/src/admission.rs b/crates/lance-context-master/src/admission.rs new file mode 100644 index 00000000..4380038c --- /dev/null +++ b/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, +} + +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) -> Option { + 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); + +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; + } +} diff --git a/crates/lance-context-master/src/catchup/mod.rs b/crates/lance-context-master/src/catchup/mod.rs index 5bdeb880..71d0672a 100644 --- a/crates/lance-context-master/src/catchup/mod.rs +++ b/crates/lance-context-master/src/catchup/mod.rs @@ -236,6 +236,9 @@ async fn admit( row: Option<&StatRow>, request: &Trigger, ) -> Result { + 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) -> Option> { })) } async fn tick(state: &Arc, 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)?; diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs index 34603d89..38e9fede 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/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; diff --git a/crates/lance-context-master/src/main.rs b/crates/lance-context-master/src/main.rs index dc6d6c0d..26f1b397 100644 --- a/crates/lance-context-master/src/main.rs +++ b/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(); } diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 3a42c7af..58eb0c49 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/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>) -> Json { + Json(state.admission.status()) +} + +async fn drain_executor( + State(state): State>, + Json(request): Json, +) -> Result, 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> { 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)) diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index 0f7d5f97..d1d5b5f0 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/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) -> lance::Result { + 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) -> lance::Result { } async fn try_scan_once(state: &Arc, maintain: bool) -> lance::Result> { + let Some(_operation) = state.admission.try_admit() else { + return Ok(None); + }; let Some(guard) = state .task_store .try_coordination_lock("stats-writer") diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 56d995db..fa59a29b 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/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(¤t.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] diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index 0e09c3d1..ab654c94 100644 --- a/crates/lance-context-master/src/state.rs +++ b/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, /// 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), diff --git a/docs/operations/master-drain.md b/docs/operations/master-drain.md new file mode 100644 index 00000000..18efba05 --- /dev/null +++ b/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.