diff --git a/crates/lance-context-core/src/rollout_append.rs b/crates/lance-context-core/src/rollout_append.rs index 235e1225..42a42fac 100644 --- a/crates/lance-context-core/src/rollout_append.rs +++ b/crates/lance-context-core/src/rollout_append.rs @@ -273,6 +273,15 @@ pub enum StageEvent { Failed(String), } +/// Merged generation watermark per shard, read from the base table's MemWAL +/// index at its latest version. Public so executors can report what they +/// merged alongside releasing ownership. Empty when the table has no WAL +/// index yet. +pub async fn merged_watermarks(uri: &str) -> Result> { + let dataset = Dataset::open(uri).await?; + watermarks(&dataset).await +} + pub(crate) async fn watermarks(dataset: &Dataset) -> Result> { let indices = dataset.load_indices().await?; let Some(index) = indices.iter().find(|i| i.name == MEM_WAL_INDEX_NAME) else { diff --git a/crates/lance-context-master/src/maintenance_execution.rs b/crates/lance-context-master/src/maintenance_execution.rs index fc92960c..28e132ee 100644 --- a/crates/lance-context-master/src/maintenance_execution.rs +++ b/crates/lance-context-master/src/maintenance_execution.rs @@ -261,7 +261,7 @@ where let scope = MergeWriteScope::with_pinned_authorizer(Arc::new(Guard { coordinator: coordinator.clone(), execution: running.clone(), - uri, + uri: uri.clone(), })); let executor_id = state.admission.status().executor_id; let outcome = std::panic::AssertUnwindSafe(async { @@ -321,9 +321,40 @@ where { return Err("maintenance ownership changed before completion".into()); } - if !coordinator.release(&proof, &terminal).await? { + // A successful WAL catch-up reclaimed generations: publish the merged + // watermark per shard in the release transaction so demand and ownership + // change together. Other maintenance kinds do not move WAL watermarks. A + // watermark read failure must not strand a completed merge; release + // without demand and let the next scan correct it. + let merged = if maintenance == MaintenanceKind::Catchup && outcome.is_ok() { + match lance_context_core::rollout_append::merged_watermarks(&uri).await { + Ok(marks) => marks + .into_iter() + .map( + |(shard, generation)| lance_context_merge::demand::ShardMerged { + shard: shard.to_string(), + merged_through: generation, + }, + ) + .collect(), + Err(error) => { + tracing::warn!(target = %running.target, %error, + "merged watermarks unavailable; releasing without demand update"); + Vec::new() + } + } + } else { + Vec::new() + }; + if !coordinator + .release_with_demand(&proof, &terminal, &merged) + .await? + { return Err("maintenance completion claim lost".into()); } + if !merged.is_empty() { + metrics::counter!("master_demand_merged_watermarks_total").increment(merged.len() as u64); + } outcome } diff --git a/crates/lance-context-master/src/rollout_append.rs b/crates/lance-context-master/src/rollout_append.rs index 54c7c5b9..c9acb711 100644 --- a/crates/lance-context-master/src/rollout_append.rs +++ b/crates/lance-context-master/src/rollout_append.rs @@ -626,6 +626,44 @@ mod tests { format!("{target}-{id}") ); } + // The release published merged watermarks for both shards in + // the same transaction that gave up ownership, so a fold over + // the events shows no pending generations. + let marks = lance_context_core::rollout_append::merged_watermarks(&uri) + .await + .unwrap(); + assert_eq!(marks.len(), 2, "{target}: {marks:?}"); + let encoded: String = target.bytes().map(|b| format!("{b:02x}")).collect(); + let events = first + .task_store + .etcd_client() + .clone() + .get( + format!( + "{}/demand-events/{encoded}/", + first.config.etcd.etcd_prefix.trim_end_matches('/') + ), + Some(etcd_client::GetOptions::new().with_prefix()), + ) + .await + .unwrap(); + assert_eq!(events.kvs().len(), 2, "{target}"); + let mut table = lance_context_merge::demand::TableDemand::default(); + for kv in events.kvs() { + let event: lance_context_merge::demand::DemandEvent = + serde_json::from_slice(kv.value()).unwrap(); + assert_eq!( + event.source, + lance_context_merge::demand::EventSource::Executor + ); + let expected = marks + .iter() + .find(|(shard, _)| shard.to_string() == event.shard) + .map(|(_, generation)| *generation); + assert_eq!(event.merged_through, expected); + table.fold(&event, kv.mod_revision(), 0).unwrap(); + } + assert_eq!(table.pending_generations(), 0, "{target}"); } } diff --git a/crates/lance-context-merge/src/demand.rs b/crates/lance-context-merge/src/demand.rs index 6857baed..25147f7e 100644 --- a/crates/lance-context-merge/src/demand.rs +++ b/crates/lance-context-merge/src/demand.rs @@ -14,6 +14,13 @@ use serde::{Deserialize, Serialize}; use crate::{Coordinator, Result}; +fn chrono_now_ms() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or_default() +} + pub const SCHEMA_VERSION: u32 = 1; /// Who produced a watermark event. @@ -203,7 +210,79 @@ impl TableDemand { } } +/// A shard's merged watermark as observed by an executor after its commit. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ShardMerged { + pub shard: String, + pub merged_through: u64, +} + impl Coordinator { + /// Compare/op pairs that publish `merged_through` for each shard, for + /// inclusion in a larger transaction. Reads current events first; a + /// shard with no stored event gets a fresh executor-sourced record with + /// `sealed_through = merged_through` (the writer's next flush raises it). + /// Skips shards whose stored `merged_through` is already ≥ the new one. + pub(crate) async fn demand_release_changes( + &self, + target: &str, + merged: &[ShardMerged], + ) -> Result<(Vec, Vec)> { + let mut compares = Vec::new(); + let mut operations = Vec::new(); + if merged.is_empty() { + return Ok((compares, operations)); + } + let mut client = self.client.clone(); + let now_ms = chrono_now_ms(); + for shard in merged { + let key = self.demand_event_key(target, &shard.shard); + let current = client + .get(key.as_str(), None) + .await + .map_err(|e| e.to_string())?; + let stored: Option = current + .kvs() + .first() + .and_then(|kv| serde_json::from_slice(kv.value()).ok()); + let next = match &stored { + Some(stored) => { + if stored + .merged_through + .is_some_and(|m| m >= shard.merged_through) + { + continue; + } + DemandEvent { + merged_through: Some(shard.merged_through), + sealed_through: stored.sealed_through.max(shard.merged_through), + source: EventSource::Executor, + ..stored.clone() + } + } + None => DemandEvent { + v: SCHEMA_VERSION, + shard: shard.shard.clone(), + sealed_through: shard.merged_through, + sealed_bytes_through: 0, + merged_through: Some(shard.merged_through), + flushed_at_ms: now_ms, + writer_epoch: 0, + source: EventSource::Executor, + }, + }; + let value = serde_json::to_vec(&next).map_err(|e| e.to_string())?; + compares.push(match current.kvs().first() { + Some(kv) => { + Compare::mod_revision(key.as_str(), CompareOp::Equal, kv.mod_revision()) + } + None => Compare::version(key.as_str(), CompareOp::Equal, 0), + }); + operations.push(TxnOp::put(key, value, None)); + } + Ok((compares, operations)) + } + fn demand_event_key(&self, target: &str, shard: &str) -> String { let encoded: String = target.bytes().map(|b| format!("{b:02x}")).collect(); format!( @@ -474,6 +553,80 @@ mod tests { ); } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn release_changes_only_advance_merged_and_keep_writer_fields() { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let client = etcd_client::Client::connect([endpoint], None) + .await + .unwrap(); + let coordinator = Coordinator::new( + client.clone(), + format!("/demand-release/{}", uuid::Uuid::new_v4()), + ); + // Writer has sealed 20 at epoch 3 with bytes; executor merged 15 then 12. + assert!(coordinator + .publish_demand_event("t", &ev("s", 20, None, 3, EventSource::Writer)) + .await + .unwrap()); + let (c, o) = coordinator + .demand_release_changes( + "t", + &[ShardMerged { + shard: "s".into(), + merged_through: 15, + }], + ) + .await + .unwrap(); + assert_eq!((c.len(), o.len()), (1, 1)); + assert!(coordinator.transact(c, o).await.unwrap()); + let read = |c: etcd_client::Client, coordinator: Coordinator| async move { + let key = coordinator.demand_event_key("t", "s"); + let kv = c.clone().get(key, None).await.unwrap(); + serde_json::from_slice::(kv.kvs()[0].value()).unwrap() + }; + let after = read(client.clone(), coordinator.clone()).await; + assert_eq!(after.merged_through, Some(15)); + assert_eq!(after.sealed_through, 20, "writer's sealed kept"); + assert_eq!(after.writer_epoch, 3, "writer's epoch kept"); + assert_eq!(after.sealed_bytes_through, 20_000, "writer's bytes kept"); + assert_eq!(after.source, EventSource::Executor); + // A lower merged is skipped entirely: no compare, no op. + let (c, o) = coordinator + .demand_release_changes( + "t", + &[ShardMerged { + shard: "s".into(), + merged_through: 12, + }], + ) + .await + .unwrap(); + assert!(c.is_empty() && o.is_empty()); + // Unknown shard: fresh executor record, sealed == merged. + let (c, o) = coordinator + .demand_release_changes( + "t", + &[ShardMerged { + shard: "new".into(), + merged_through: 7, + }], + ) + .await + .unwrap(); + assert!(coordinator.transact(c, o).await.unwrap()); + let key = coordinator.demand_event_key("t", "new"); + let fresh: DemandEvent = + serde_json::from_slice(client.clone().get(key, None).await.unwrap().kvs()[0].value()) + .unwrap(); + assert_eq!((fresh.sealed_through, fresh.merged_through), (7, Some(7))); + let mut table = TableDemand::default(); + table.fold(&after, 1, 0).unwrap(); + table.fold(&fresh, 2, 0).unwrap(); + assert_eq!(table.pending_generations(), 5); + } + #[test] fn wire_format_is_stable() { let event = ev("s", 1, Some(1), 1, EventSource::Writer); diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index 6690d919..bc9f015f 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -290,11 +290,32 @@ impl Coordinator { /// Never delete on elapsed time, an HTTP error, or claim expiry. Both the /// terminal result and the current task claim must still match atomically. pub async fn release(&self, claim: &ClaimProof, execution: &Execution) -> Result { + self.release_with_demand(claim, execution, &[]).await + } + + /// Release ownership and, in the same transaction, publish the merged + /// watermark for every shard in `merged` as a `demand-events` record with + /// source `executor`. Demand and ownership therefore change together: a + /// planner never sees a released table whose pending count still counts + /// the generations this execution merged. Each watermark write is + /// monotonic by the same rule as `publish_demand_event`; a shard whose + /// stored event is already at least as new is skipped, never lowered. + pub async fn release_with_demand( + &self, + claim: &ClaimProof, + execution: &Execution, + merged: &[demand::ShardMerged], + ) -> Result { if !matches!(execution.phase, Phase::Finished | Phase::Recovered) { return Err("cannot release a live execution".into()); } let key = execution_key(&self.prefix, &execution.target); let (mut compares, mut operations) = self.completion_changes(execution).await?; + let (demand_compares, demand_ops) = self + .demand_release_changes(&execution.target, merged) + .await?; + compares.extend(demand_compares); + operations.extend(demand_ops); compares.extend([ Compare::value(claim.key.as_str(), CompareOp::Equal, claim.token.as_bytes()), Compare::value(key.as_str(), CompareOp::Equal, encode(execution)), @@ -327,7 +348,11 @@ impl Coordinator { .await } - async fn transact(&self, compares: Vec, operations: Vec) -> Result { + pub(crate) async fn transact( + &self, + compares: Vec, + operations: Vec, + ) -> Result { self.client .clone() .txn(Txn::new().when(compares).and_then(operations))