diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index 678f47c7..24c4b4c9 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -171,6 +171,8 @@ jobs: cargo clippy -p lance-context-ingest --features http --all-targets -- -D warnings cargo clippy -p lance-context-ingest --features http-client,materializer --all-targets -- -D warnings cargo test -p lance-context-ingest --features http-client,materializer --lib + - name: Materializer worker feature checks + run: cargo clippy -p lance-context-ingest --features worker --all-targets -- -D warnings - name: Start etcd for HA scheduler test run: | ETCD_VERSION=3.7.0 @@ -198,6 +200,7 @@ jobs: run: | cargo test -p lance-context-master -p lance-context-server -p lance-context-merge -- --ignored cargo test -p lance-context-core registry --lib -- --ignored + cargo test -p lance-context-ingest --features worker --lib tests::worker -- --ignored coverage: runs-on: ubuntu-24.04 @@ -263,6 +266,8 @@ jobs: -p lance-context-master -p lance-context-server -p lance-context-merge -- --ignored cargo llvm-cov --no-report --all-features \ -p lance-context-core --lib -- registry --ignored + cargo llvm-cov --no-report --all-features \ + -p lance-context-ingest --lib -- tests::worker --ignored - name: Merge coverage reports run: cargo llvm-cov report --lcov --output-path lcov.info - name: Save coverage report diff --git a/Cargo.lock b/Cargo.lock index 804a7bef..c7ac9567 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4185,6 +4185,7 @@ dependencies = [ "axum", "base64 0.22.1", "bytes", + "etcd-client", "futures", "http-body-util", "lance", diff --git a/crates/lance-context-ingest/Cargo.toml b/crates/lance-context-ingest/Cargo.toml index 1c7e6eae..bc54f5e3 100644 --- a/crates/lance-context-ingest/Cargo.toml +++ b/crates/lance-context-ingest/Cargo.toml @@ -7,11 +7,13 @@ publish = false description = "Durable rollout ingress over the independent lance-context WAL" [features] +worker = ["materializer", "dep:etcd-client"] http = ["dep:axum", "dep:http-body-util"] http-client = ["http", "dep:reqwest", "lance-context-wal/client-journal"] materializer = ["lance-context-wal/lance-integration", "dep:lance", "dep:lance-table", "dep:object_store", "dep:futures"] [dependencies] +etcd-client = { version = "0.19", features = ["tls"], optional = true } axum = { version = "0.8", optional = true, default-features = false, features = ["json", "http1", "tokio"] } http-body-util = { version = "0.1", optional = true } reqwest = { version = "0.12", optional = true, default-features = false, features = ["json", "rustls-tls"] } diff --git a/crates/lance-context-ingest/src/lib.rs b/crates/lance-context-ingest/src/lib.rs index 9b3e9246..6e48498f 100644 --- a/crates/lance-context-ingest/src/lib.rs +++ b/crates/lance-context-ingest/src/lib.rs @@ -34,6 +34,8 @@ pub mod http; #[cfg(feature = "materializer")] pub mod materializer; pub mod receipt; +#[cfg(feature = "worker")] +pub mod worker; /// Resolved by the application's authorized catalog, not an uploaded URI/UUID. #[derive(Clone, Debug)] diff --git a/crates/lance-context-ingest/src/materializer.rs b/crates/lance-context-ingest/src/materializer.rs index 488aa429..3df480ac 100644 --- a/crates/lance-context-ingest/src/materializer.rs +++ b/crates/lance-context-ingest/src/materializer.rs @@ -54,7 +54,7 @@ impl Options { }) .context("materializer workspace overflow") } - fn validate(&self, ingress: &Ingress, budget: &MergeMemoryBudget) -> Result<()> { + pub(super) fn validate(&self, ingress: &Ingress, budget: &MergeMemoryBudget) -> Result<()> { ensure!( self.max_source_bytes >= ingress.scope.contract.max_record_bytes as usize && self.max_source_bytes as u64 <= ingress.client.max_plan_source_bytes(), @@ -183,31 +183,9 @@ impl Publisher { ); Ok(view) } - fn reader(&self) -> Result { - p::lance::ManagedManifest::reader( - Arc::new(self.ingress.client.manifest_catalog()), - self.storage.base.clone(), - &self.ingress.layout, - self.authority.table, - self.ingress.scope.incarnation, - ) - } /// Readers of this managed table must use its catalog-selected manifest. pub async fn read_dataset(&self, version: Option) -> Result { - let mut builder = DatasetBuilder::from_uri(&self.storage.uri) - .with_storage_options(self.storage.options.clone()) - .with_session(self.storage.session.clone()) - .with_commit_handler(Arc::new(self.reader()?)); - if let Some(version) = version { - builder = builder.with_version(version); - } - let dataset = builder.load().await?; - let schema: arrow_schema::Schema = dataset.schema().into(); - ensure!( - schema == rollout_schema(), - "managed dataset is not the bound rollout schema" - ); - Ok(dataset) + read_dataset(&self.ingress, &self.storage, version).await } /// Select a complete bounded prefix, pin every source, write immutable Lance @@ -627,7 +605,41 @@ impl Publisher { } } -async fn query(ingress: &Ingress, group: u64, value: query::Query) -> Result { +/// Open a catalog-selected dataset without acquiring publisher ownership. +/// The application resolves and authorizes both the binding and storage first. +pub async fn read_dataset( + ingress: &Ingress, + storage: &Storage, + version: Option, +) -> Result { + let reader = p::lance::ManagedManifest::reader( + Arc::new(ingress.client.manifest_catalog()), + storage.base.clone(), + &ingress.layout, + ingress.scope.table, + ingress.scope.incarnation, + )?; + let mut builder = DatasetBuilder::from_uri(&storage.uri) + .with_storage_options(storage.options.clone()) + .with_session(storage.session.clone()) + .with_commit_handler(Arc::new(reader)); + if let Some(version) = version { + builder = builder.with_version(version); + } + let dataset = builder.load().await?; + let schema: arrow_schema::Schema = dataset.schema().into(); + ensure!( + schema == rollout_schema(), + "managed dataset is not the bound rollout schema" + ); + Ok(dataset) +} + +pub(super) async fn query( + ingress: &Ingress, + group: u64, + value: query::Query, +) -> Result { let response = ingress .client .request_routed(Request::Query { diff --git a/crates/lance-context-ingest/src/tests.rs b/crates/lance-context-ingest/src/tests.rs index 9825c6a5..d069b698 100644 --- a/crates/lance-context-ingest/src/tests.rs +++ b/crates/lance-context-ingest/src/tests.rs @@ -15,6 +15,8 @@ mod http; #[cfg(feature = "materializer")] mod materializer; mod network; +#[cfg(feature = "worker")] +mod worker; fn layout() -> Arc { Arc::new(Layout { diff --git a/crates/lance-context-ingest/src/tests/worker.rs b/crates/lance-context-ingest/src/tests/worker.rs new file mode 100644 index 00000000..bb4c80c2 --- /dev/null +++ b/crates/lance-context-ingest/src/tests/worker.rs @@ -0,0 +1,233 @@ +use super::*; +use crate::{ + materializer::{self as publication, Storage}, + worker::{Assignment, Claims, Options, Worker}, +}; +use futures::TryStreamExt; +use std::{collections::HashMap, path::Path}; + +async fn fixture() -> Result<(etcd_client::Client, Claims, String)> { + let endpoints = std::env::var("ETCD_TEST_ENDPOINTS").context("ETCD_TEST_ENDPOINTS required")?; + let client = + etcd_client::Client::connect(endpoints.split(',').collect::>(), None).await?; + let prefix = format!("/lance-context-test/ingest-worker/{}", Uuid::new_v4()); + let claims = Claims::new(client.clone(), prefix.clone(), 10, Duration::from_secs(3))?; + Ok((client, claims, prefix)) +} +async fn storage(path: &Path) -> Result { + let dataset = Dataset::open(path.to_str().context("test URI")?).await?; + let base = dataset + .manifest_location() + .path + .as_ref() + .rsplit_once("/_versions/") + .context("manifest base")? + .0 + .to_owned(); + Ok(Storage { + uri: path.to_str().unwrap().into(), + base: object_store::path::Path::parse(base)?, + options: HashMap::new(), + session: Arc::new(lance::session::Session::default()), + }) +} +fn options() -> Options { + Options { + parallel_tables: 1, + publications_per_claim: 4, + poll_interval: Duration::from_millis(50), + publication: publication::Options { + max_source_bytes: 1 << 20, + max_records: 8, + retained_manifest_versions: 2, + ..Default::default() + }, + } +} +async fn append( + ingress: &Ingress, + producer: &Producer, + sequence: u64, + id: &str, +) -> Result<(Bytes, AcceptedAppend)> { + let session = session(0); + let body = ipc_file(&[batch(&[id], &[&session])?])?; + let prepared = ingress + .prepare( + ingress.reserve_upload(body.len()).await?, + body.clone(), + producer, + sequence, + ) + .await?; + Ok((body, ingress.append(&prepared).await?)) +} +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires ETCD_TEST_ENDPOINTS; executed by the etcd CI suite"] +async fn expired_claim_cannot_steal_successor_epoch_and_workers_resume_prepared_lance() -> Result<()> +{ + let (mut etcd, claims, prefix) = fixture().await?; + let root = tempfile::tempdir()?; + let path = root.path().join("rollouts.lance"); + let c = network::Cluster::open_signed(root.path(), layout()).await?; + bootstrap(&c, &path, &binding()).await?; + let ingress = connect(&c, binding()).await?; + let storage = storage(&path).await?; + let assignment = Assignment { + ingress: ingress.clone(), + storage: storage.clone(), + }; + let budget = MergeMemoryBudget::new(32 << 20); + let producer = ingress + .open_producer(ProducerOpen { + bucket: 0, + slot: 0, + expected_generation: 0, + operation: Uuid::new_v4(), + }) + .await?; + let (first, receipt) = append(&ingress, &producer, 1, "before-takeover").await?; + let mut expected = BTreeMap::new(); + materializer::remember(&mut expected, first)?; + let old = claims + .take(&ingress) + .await? + .context("first claim unavailable")?; + assert!( + claims.take(&ingress).await?.is_none(), + "overlapping table claim granted" + ); + let mut old_publisher = old + .acquire(assignment.clone(), options().publication, budget.clone()) + .await? + .context("first claim invalid")?; + let prepared = old_publisher + .prepare_next() + .await? + .context("prepared result missing")?; + let old_epoch = old_publisher.authority().epoch; + let key = format!("{prefix}/{}/{}", layout().cluster, binding().table); + let lease = etcd.get(key, None).await?.kvs()[0].lease(); + etcd.lease_revoke(lease).await?; + let successor = claims + .take(&ingress) + .await? + .context("successor claim unavailable")?; + let mut publisher = successor + .acquire(assignment.clone(), options().publication, budget.clone()) + .await? + .context("successor claim invalid")?; + assert_eq!(publisher.authority().epoch, old_epoch + 1); + // This deliberately calls the OLD claim after the successor has committed + // its newer epoch. It must fail the lease check and leave that epoch intact. + assert!(old + .acquire(assignment.clone(), options().publication, budget.clone()) + .await? + .is_none()); + assert_eq!(publisher.state().await?.authority, *publisher.authority()); + assert!(old_publisher.run_once().await.is_err()); + assert_eq!(publisher.prepare_next().await?.unwrap(), prepared); + let _ = old.release().await; + assert!( + successor.valid().await?, + "old lease cleanup removed successor" + ); + assert!(publisher.run_once().await?.is_some()); + assert_eq!( + ingress + .wait_published(&receipt, Duration::from_secs(5)) + .await?, + 2 + ); + materializer::verify(&publisher, &expected, 2).await?; + successor.release().await?; + drop((old_publisher, publisher)); + + let (second, receipt) = append(&ingress, &producer, 2, "automatic-worker").await?; + materializer::remember(&mut expected, second)?; + let workers = [ + Worker::start( + claims.clone(), + vec![assignment.clone()], + options(), + budget.clone(), + )?, + Worker::start(claims.clone(), vec![assignment], options(), budget.clone())?, + ]; + assert_eq!( + ingress + .wait_published(&receipt, Duration::from_secs(15)) + .await?, + 3 + ); + assert!(workers.iter().all(|w| w.observation().active <= 1)); + for worker in workers { + worker.shutdown().await?; + } + assert_eq!(budget.reserved(), 0); + // Read independently of either stopped publisher; no acquisition for reads. + let dataset = publication::read_dataset(&ingress, &storage, None).await?; + assert_eq!(dataset.version().version, 3); + let mut found = BTreeMap::new(); + for batch in dataset + .scan() + .try_into_stream() + .await? + .try_collect::>() + .await? + { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + let id = ids.value(row); + assert!(found.insert(id.to_owned(), ()).is_none()); + assert_eq!(batch.slice(row, 1), expected[id]); + } + } + assert_eq!(found.len(), expected.len()); + etcd.delete( + prefix, + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await?; + drop((dataset, ingress)); + c.shutdown().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires ETCD_TEST_ENDPOINTS; executed by the etcd CI suite"] +async fn abandoned_claim_expires_without_a_detached_delete() -> Result<()> { + let (mut etcd, claims, prefix) = fixture().await?; + let root = tempfile::tempdir()?; + let c = network::Cluster::open_signed(root.path(), layout()).await?; + bootstrap(&c, &root.path().join("rollouts.lance"), &binding()).await?; + let ingress = connect(&c, binding()).await?; + let first = claims.take(&ingress).await?.context("claim unavailable")?; + drop(first); + assert!( + claims.take(&ingress).await?.is_none(), + "drop deleted the claim instead of letting its lease expire" + ); + let successor = tokio::time::timeout(Duration::from_secs(20), async { + loop { + if let Some(claim) = claims.take(&ingress).await? { + return Ok::<_, anyhow::Error>(claim); + } + tokio::time::sleep(Duration::from_millis(250)).await; + } + }) + .await??; + assert!(successor.valid().await?); + successor.release().await?; + etcd.delete( + prefix, + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await?; + drop(ingress); + c.shutdown().await +} diff --git a/crates/lance-context-ingest/src/worker.rs b/crates/lance-context-ingest/src/worker.rs new file mode 100644 index 00000000..8a2d850b --- /dev/null +++ b/crates/lance-context-ingest/src/worker.rs @@ -0,0 +1,238 @@ +//! Bounded materializer scheduling over etcd leases. Leases distribute work; +//! durable WAL epochs and manifest selection fence stale workers. The input +//! assignments must come from the application's authorized table registry. +use super::*; +use crate::materializer::{Acquisition, Options as PublishOptions, Publisher, Storage}; +use std::{ + collections::BTreeSet, + sync::atomic::{AtomicU64, AtomicUsize, Ordering}, +}; +use tokio::{ + sync::watch, + task::{JoinHandle, JoinSet}, +}; + +mod lease; +pub use lease::Claims; + +#[derive(Clone)] +pub struct Assignment { + pub ingress: Ingress, + pub storage: Storage, +} +#[derive(Clone, Debug)] +pub struct Options { + pub parallel_tables: usize, + /// Release a healthy claim periodically so another worker can take a turn. + pub publications_per_claim: usize, + pub poll_interval: Duration, + pub publication: PublishOptions, +} +impl Default for Options { + fn default() -> Self { + Self { + parallel_tables: 4, + publications_per_claim: 16, + poll_interval: Duration::from_millis(500), + publication: PublishOptions::default(), + } + } +} +#[derive(Default)] +struct Counters { + active: AtomicUsize, + claimed: AtomicU64, + published: AtomicU64, + errors: AtomicU64, +} +#[derive(Debug, Clone, Serialize)] +pub struct Observation { + pub running: bool, + pub active: usize, + pub claims: u64, + pub publications: u64, + pub errors: u64, + pub draining: bool, +} +pub struct Worker { + stop: watch::Sender, + task: Option>>, + counters: Arc, +} +impl Drop for Worker { + fn drop(&mut self) { + // Dropping a JoinHandle detaches the supervisor. It still owns and joins + // its in-flight publication tasks; cancellation must not abort file work. + let _ = self.stop.send(true); + } +} +impl Worker { + /// Start a fixed, authorized assignment set. Separate workers can receive + /// overlapping sets; one expiring scheduling claim coordinates each table. + pub fn start( + claims: Claims, + assignments: Vec, + options: Options, + budget: Arc, + ) -> Result { + ensure!( + (1..=4096).contains(&assignments.len()) + && (1..=64).contains(&options.parallel_tables) + && (1..=1024).contains(&options.publications_per_claim) + && (Duration::from_millis(10)..=Duration::from_secs(60)) + .contains(&options.poll_interval), + "invalid worker bounds" + ); + let mut tables = BTreeSet::new(); + for assignment in &assignments { + ensure!( + tables.insert(assignment.ingress.scope.table), + "duplicate worker table" + ); + options.publication.validate(&assignment.ingress, &budget)?; + } + let counters = Arc::new(Counters::default()); + let (stop, stopped) = watch::channel(false); + let task = tokio::spawn(run( + claims, + assignments, + options, + budget, + stopped, + counters.clone(), + )); + Ok(Self { + stop, + task: Some(task), + counters, + }) + } + pub fn observation(&self) -> Observation { + Observation { + running: self.task.as_ref().is_some_and(|task| !task.is_finished()), + active: self.counters.active.load(Ordering::Acquire), + claims: self.counters.claimed.load(Ordering::Relaxed), + publications: self.counters.published.load(Ordering::Relaxed), + errors: self.counters.errors.load(Ordering::Relaxed), + draining: *self.stop.borrow(), + } + } + /// Stop new scheduling, finish admitted publication work, stop renewal, and + /// revoke only each owned lease. A cancelled waiter does not abort the owner. + pub async fn shutdown(mut self) -> Result<()> { + let _ = self.stop.send(true); + self.task.take().context("worker already joined")?.await? + } +} + +async fn run( + claims: Claims, + assignments: Vec, + options: Options, + budget: Arc, + mut stop: watch::Receiver, + counters: Arc, +) -> Result<()> { + let mut jobs: JoinSet<(usize, Result<()>)> = JoinSet::new(); + let mut active = BTreeSet::new(); + let mut cursor = Uuid::new_v4().as_u128() as usize % assignments.len(); + let mut tick = tokio::time::interval(options.poll_interval); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut failure = None; + loop { + if *stop.borrow() { + break; + } + tokio::select! { + _ = stop.changed() => break, + done = jobs.join_next(), if !jobs.is_empty() => { + counters.active.fetch_sub(1, Ordering::Release); + match done { + Some(Ok((index, result))) => { + active.remove(&index); + if result.is_err() { counters.errors.fetch_add(1, Ordering::Relaxed); } + } + Some(Err(error)) => { + failure = Some(error); + break; + } + None => unreachable!("nonempty join set ended"), + } + } + _ = tick.tick() => { + // At most one pass and at most the configured number of jobs. + // Failed/busy claims wait for the next tick, avoiding a hot loop. + for _ in 0..assignments.len() { + if active.len() == options.parallel_tables { break; } + let index = cursor; + cursor = (cursor + 1) % assignments.len(); + if !active.insert(index) { continue; } + counters.active.fetch_add(1, Ordering::Release); + let (claims, assignment, options, budget, stop, counters) = ( + claims.clone(), assignments[index].clone(), options.clone(), + budget.clone(), stop.clone(), counters.clone(), + ); + jobs.spawn(async move { + let result = publish_claim(claims, assignment, options, budget, stop, counters).await; + (index, result) + }); + } + } + } + } + // Do not abort a submitted blob write or lose its memory reservation while + // it is completing. The deployer budgets termination grace for this drain. + while let Some(done) = jobs.join_next().await { + counters.active.fetch_sub(1, Ordering::Release); + match done { + Ok((_, Ok(()))) => {} + Ok((_, Err(_))) => { + counters.errors.fetch_add(1, Ordering::Relaxed); + } + Err(error) => { + failure = Some(error); + } + } + } + if let Some(error) = failure { + return Err(error.into()); + } + Ok(()) +} +async fn publish_claim( + claims: Claims, + assignment: Assignment, + options: Options, + budget: Arc, + stop: watch::Receiver, + counters: Arc, +) -> Result<()> { + let Some(claim) = claims.take(&assignment.ingress).await? else { + return Ok(()); + }; + counters.claimed.fetch_add(1, Ordering::Relaxed); + let result = async { + if *stop.borrow() { + return Ok(()); + } + let Some(mut publisher) = claim + .acquire(assignment, options.publication, budget) + .await? + else { + return Ok(()); + }; + for _ in 0..options.publications_per_claim { + if *stop.borrow() || !claim.valid().await? { + break; + } + if publisher.run_once().await?.is_none() { + break; + } + counters.published.fetch_add(1, Ordering::Relaxed); + } + Ok::<_, anyhow::Error>(()) + } + .await; + let release = claim.release().await; + result.and(release) +} diff --git a/crates/lance-context-ingest/src/worker/lease.rs b/crates/lance-context-ingest/src/worker/lease.rs new file mode 100644 index 00000000..4120fcb5 --- /dev/null +++ b/crates/lance-context-ingest/src/worker/lease.rs @@ -0,0 +1,238 @@ +use super::*; +use etcd_client::{Client as Etcd, Compare, CompareOp, GetOptions, PutOptions, Txn, TxnOp}; +use std::{future::Future, sync::atomic::AtomicBool}; + +/// Supply the same authorized etcd prefix to every worker serving these tables. +/// Transport credentials/TLS belong to the application's etcd connection. +#[derive(Clone)] +pub struct Claims { + client: Etcd, + prefix: String, + ttl: i64, + timeout: Duration, +} +impl Claims { + pub fn new(client: Etcd, prefix: String, ttl: i64, timeout: Duration) -> Result { + ensure!( + prefix.starts_with('/') + && !prefix.trim_matches('/').is_empty() + && prefix.len() <= 1024 + && !prefix.contains('\0') + && (10..=120).contains(&ttl) + && (Duration::from_millis(100)..=Duration::from_secs(ttl as u64 / 3)) + .contains(&timeout), + "invalid materializer scheduling scope" + ); + Ok(Self { + client, + prefix: prefix.trim_end_matches('/').into(), + ttl, + timeout, + }) + } + async fn rpc( + &self, + operation: impl Future>, + ) -> Result { + Ok(tokio::time::timeout(self.timeout, operation) + .await + .context("etcd scheduling deadline")??) + } + pub(crate) async fn take(&self, ingress: &Ingress) -> Result> { + let scope = super::super::receipt::scope_hash(&ingress.scope)?; + let owner = Uuid::new_v4(); + let request = Uuid::new_v4(); + let key = format!( + "{}/{}/{}", + self.prefix, ingress.scope.cluster, ingress.scope.table + ); + // Busy-table probes must not grant/revoke leases on every worker poll: + // those are replicated writes and can stall unrelated lease renewals. + // This is only a hint: a stale occupied value delays one attempt, while + // a stale empty value still must pass the transaction below. The local + // read avoids a quorum barrier on every occupied-table probe. The CAS + // and the later claim.valid() check remain linearizable. + if !self + .rpc(self.client.clone().get( + key.clone(), + Some(GetOptions::new().with_keys_only().with_serializable()), + )) + .await + .context("checking existing materializer claim")? + .kvs() + .is_empty() + { + return Ok(None); + } + let value = serde_json::to_vec(&(scope, owner, request))?; + let lease = self + .rpc(self.client.clone().lease_grant(self.ttl, None)) + .await + .context("granting materializer lease")? + .id(); + ensure!(lease != 0, "etcd returned an invalid lease"); + let outcome = self + .rpc( + self.client.clone().txn( + Txn::new() + .when([Compare::version(key.clone(), CompareOp::Equal, 0)]) + .and_then([TxnOp::put( + key.clone(), + value.clone(), + Some(PutOptions::new().with_lease(lease)), + )]), + ), + ) + .await + .context("claiming materializer table")?; + // Unknown outcomes leave only a bounded-TTL lease. Never assume the + // transaction failed and never delete the key to make acquisition pass. + if !outcome.succeeded() { + self.rpc(self.client.clone().lease_revoke(lease)) + .await + .context("revoking unused materializer lease")?; + return Ok(None); + } + let (mut keeper, mut stream) = self + .rpc(self.client.clone().lease_keep_alive(lease)) + .await + .context("starting materializer lease renewal")?; + let (stop, mut stopped) = watch::channel(false); + let live = Arc::new(AtomicBool::new(true)); + let state = live.clone(); + let renewal = Duration::from_secs(self.ttl as u64 / 3); + let timeout = self.timeout; + let task = tokio::spawn(async move { + let mut tick = tokio::time::interval(renewal); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = stopped.changed() => break, + _ = tick.tick() => { + let result = tokio::time::timeout(timeout, async { + keeper.keep_alive().await?; + let response = stream.message().await?.context("lease renewal stream closed")?; + ensure!(response.id() == lease && response.ttl() > 0, "lease expired"); + Ok::<_, anyhow::Error>(()) + }).await; + if !matches!(result, Ok(Ok(()))) { break; } + } + } + } + state.store(false, Ordering::Release); + }); + Ok(Some(Claim { + store: self.clone(), + key, + value, + lease, + scope, + owner, + request, + live, + stop, + task: Some(task), + })) + } +} +pub(crate) struct Claim { + store: Claims, + key: String, + value: Vec, + lease: i64, + scope: [u8; 32], + owner: Uuid, + request: Uuid, + live: Arc, + stop: watch::Sender, + task: Option>, +} +impl Drop for Claim { + fn drop(&mut self) { + let _ = self.stop.send(true); + // The bounded renewal finishes and the unrevoked lease expires if the + // owning task panics. No detached delete can remove a successor's claim. + } +} +impl Claim { + pub(crate) async fn valid(&self) -> Result { + if !self.live.load(Ordering::Acquire) { + return Ok(false); + } + let response = self + .store + .rpc(self.store.client.clone().get(self.key.clone(), None)) + .await + .context("validating materializer claim")?; + Ok(response + .kvs() + .first() + .is_some_and(|kv| kv.value() == self.value && kv.lease() == self.lease)) + } + pub(crate) async fn acquire( + &self, + assignment: Assignment, + options: PublishOptions, + budget: Arc, + ) -> Result> { + let ingress = &assignment.ingress; + ensure!( + super::super::receipt::scope_hash(&ingress.scope)? == self.scope, + "claim belongs to another binding" + ); + // Read the WAL epoch BEFORE the linearizable lease check. A paused old + // owner must not check its lease first, read the successor epoch later, + // and use that newer epoch to steal ownership back. + let query::Value::Publisher(Some(view)) = crate::materializer::query( + ingress, + ingress.client.catalog_group(), + query::Query::Publisher { + table: ingress.scope.table, + }, + ) + .await? + else { + bail!("publisher binding missing"); + }; + ensure!( + view.destination == ingress.scope.destination + && view.ingest == ingress.scope.contract + && view + .manifests + .as_ref() + .is_some_and(|m| m.incarnation == ingress.scope.incarnation), + "claim binding changed" + ); + if !self.valid().await? { + return Ok(None); + } + Ok(Some( + Publisher::acquire( + assignment.ingress, + assignment.storage, + options, + budget, + Acquisition { + owner: self.owner, + request: self.request, + expected_epoch: view.authority.epoch, + }, + ) + .await?, + )) + } + pub(crate) async fn release(mut self) -> Result<()> { + let _ = self.stop.send(true); + self.task + .take() + .context("lease keeper already joined")? + .await?; + // A lease belongs to exactly this claim. Revoke cannot delete a newer + // value attached to a successor's different lease, even after expiry. + self.store + .rpc(self.store.client.clone().lease_revoke(self.lease)) + .await + .context("releasing materializer lease")?; + Ok(()) + } +} diff --git a/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs b/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs index 6af19a6e..c7863553 100644 --- a/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs +++ b/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs @@ -117,11 +117,13 @@ async fn tls_peer_process_fixture() -> Result<()> { let listener = TcpListener::bind( std::env::var("CONTEXT_TLS_PEER_BIND").unwrap_or_else(|_| "127.0.0.1:0".into()), ) - .await?; + .await + .context("bind peer listener")?; let image_listener = TcpListener::bind( std::env::var("CONTEXT_TLS_IMAGE_BIND").unwrap_or_else(|_| "127.0.0.1:0".into()), ) - .await?; + .await + .context("bind image listener")?; let control = TcpListener::bind("127.0.0.1:0").await?; let rpc_data = TcpListener::bind("127.0.0.1:0").await?; let rpc_control = TcpListener::bind("127.0.0.1:0").await?; @@ -358,7 +360,15 @@ impl Process { Err(e) => { let _ = child.kill(); let _ = child.wait(); - Err(e) + let stdout = fs::read_to_string(root.join(format!("stdout-{id}-{incarnation}"))) + .unwrap_or_else(|error| format!("cannot read child stdout: {error}")); + let stderr = fs::read_to_string(root.join(format!("stderr-{id}-{incarnation}"))) + .unwrap_or_else(|error| format!("cannot read child stderr: {error}")); + Err(e).with_context(|| { + format!( + "TLS peer {id} incarnation {incarnation} startup failed; child stdout:\n{stdout}\nchild stderr:\n{stderr}" + ) + }) } } } diff --git a/crates/lance-context-wal/tests/raft_log.rs b/crates/lance-context-wal/tests/raft_log.rs index c755af17..2af88eb5 100644 --- a/crates/lance-context-wal/tests/raft_log.rs +++ b/crates/lance-context-wal/tests/raft_log.rs @@ -10,6 +10,11 @@ use raft::{GetEntriesContext, Storage}; use raft_proto::eraftpb::{Entry, EntryType, HardState}; use uuid::Uuid; +// A child spawned by another case can briefly inherit this case's flock until +// exec. Exclude that window around open/drop/reopen sequences. Concurrency and +// real child crashes within a case remain part of the test. +static TEST_CASE: std::sync::Mutex<()> = std::sync::Mutex::new(()); + fn layout() -> Layout { Layout { cluster: Uuid::from_u128(1), @@ -45,6 +50,7 @@ fn read(log: &GroupLog, low: u64, high: u64, limit: Option) -> raft::Result #[test] fn recovers_two_groups_and_rejects_layout_changes() -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); let directory = tempfile::tempdir()?; let path = directory.path().join("log"); let entries = vec![ @@ -98,6 +104,7 @@ fn recovers_two_groups_and_rejects_layout_changes() -> Result<()> { #[test] fn foreign_engine_and_invalid_metadata_are_rejected_without_rebinding() -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); fn engine(path: &std::path::Path) -> Result { Ok(raft_engine::Engine::open(raft_engine::Config { dir: path.to_str().unwrap().to_owned(), @@ -143,6 +150,7 @@ fn foreign_engine_and_invalid_metadata_are_rejected_without_rebinding() -> Resul #[test] fn replaces_only_uncommitted_suffix_and_preserves_vote() -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); let directory = tempfile::tempdir()?; let path = directory.path().join("log"); let expected = vec![ @@ -193,6 +201,7 @@ fn replaces_only_uncommitted_suffix_and_preserves_vote() -> Result<()> { #[test] fn read_size_and_chunk_boundaries_obey_raft_storage_contract() -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); let directory = tempfile::tempdir()?; let log = PhysicalLog::open(&directory.path().join("log"), &layout(), 1)?; let group = log.group(10)?; @@ -227,6 +236,7 @@ fn read_size_and_chunk_boundaries_obey_raft_storage_contract() -> Result<()> { #[test] fn concurrent_groups_share_engine_without_sharing_positions() -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); let directory = tempfile::tempdir()?; let path = directory.path().join("log"); { @@ -265,6 +275,7 @@ fn concurrent_groups_share_engine_without_sharing_positions() -> Result<()> { #[test] fn raft_ready_persistence_preserves_committed_data_through_election_and_restart() -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); use raft::{Config, RawNode, StateRole}; use std::path::Path; @@ -541,6 +552,7 @@ mod process { use std::path::Path; fn run_child(mode: &str) -> Result<()> { + let _case = TEST_CASE.lock().unwrap(); let directory = tempfile::tempdir()?; let path = directory.path().join("log"); let stdout = directory.path().join("worker.stdout"); diff --git a/docs/design/wal-materializer-worker.md b/docs/design/wal-materializer-worker.md new file mode 100644 index 00000000..75780e62 --- /dev/null +++ b/docs/design/wal-materializer-worker.md @@ -0,0 +1,55 @@ +# Materializer scheduling over durable WAL publication + +The optional `lance-context-ingest/worker` feature runs bounded publication tasks +over authorized `Assignment` values. Each assignment supplies an existing ingress +binding and its resolved Lance storage. Multiple processes may receive overlapping +assignments and share the same etcd claim prefix. The assignments are currently a +fixed set supplied by the application; dynamic namespace discovery and bootstrap +are separate integration work. + +`Claims` uses unique, expiring etcd leases to distribute tables. A local +keys-only hint avoids granting and revoking a lease on every busy-table poll; +a stale occupied hint delays an attempt, while a stale empty hint must still +pass the linearizable create-if-absent transaction. Claim validation before WAL +acquisition remains linearizable. Each successful +claim records a new owner/request identity and the complete binding fingerprint. +Renewal and scheduling RPCs have explicit deadlines. A lost claim transaction +response leaves an expiring lease; it never authorizes deleting or stealing an +existing claim. Releasing a claim joins its renewal task and revokes only its own +lease. A cancelled owner stops renewing and allows TTL expiry. + +The lease is not the publication fence. Before acquisition the worker reads the +WAL publisher epoch, then checks the claim with a linearizable etcd read, and +finally submits an epoch-conditional WAL acquisition. Reversing those first two +steps would let a paused old owner observe a successor's newer epoch and steal it +back. An old authority cannot prepare, select or confirm the successor's manifests. +Already prepared files and pending publication state survive a change of owner. + +`Worker::start` bounds active table jobs, publication rounds per claim, scanning, +and polling. Assignments share a supplied memory budget and the ingresses' WAL +connection pool. Workers release claims after the configured publication quantum +or when no work is found. Busy/failed claims wait for a later poll; they do not spin +in an immediate retry loop. Fixed assignment polling can produce idle acquisition +traffic and is not yet an active-table discovery index or a scaling result. + +`Worker::shutdown` stops new scheduling and joins in-flight publication and lease +cleanup. Cancelling a shutdown waiter does not cancel its owning supervisor or +release memory while file work still runs. A task panic drains the other jobs +before reporting failure. `observation` reports local job, claim, publication and +error counts plus supervisor/drain state; it does not claim quorum availability. +Deployment termination grace must cover admitted blob work and lease cleanup. + +Readers can use `materializer::read_dataset` with the authorized ingress/storage +binding without acquiring a publisher epoch. It opens only the manifest selected +by the WAL catalog and checks the rollout schema. + +The new etcd-backed tests revoke a claim after real Lance file preparation, allow +a successor to acquire, attempt to reacquire through the old claim, verify old +authority rejection and exact prepared-file reuse, then run two automatic workers +over the same table and read every resulting row through an independent reader. +A second test abandons a claim and requires actual TTL expiry before takeover. +Both etcd checks passed, along with the five HTTP/ingress/materializer regressions +and worker feature Clippy checks. They are correctness checks, not a process-kill +or sustained throughput benchmark. + +Tracking: #378.