diff --git a/crates/lance-context-master/src/demand_publish.rs b/crates/lance-context-master/src/demand_publish.rs new file mode 100644 index 0000000..6591103 --- /dev/null +++ b/crates/lance-context-master/src/demand_publish.rs @@ -0,0 +1,69 @@ +//! Detached, bounded publication of demand watermarks from master-side +//! sources (`docs/design/scheduler-p1-data-model.md` §1.3). +//! +//! Same discipline as the writer-side publish: never awaited by the caller, +//! one hard timeout, a process-wide concurrency cap, drop-and-count when the +//! cap is full. Every publish is a per-shard monotonic join, so a dropped or +//! late publish cannot misstate demand; the next scan reconciles it. + +use std::sync::Arc; +use std::time::Duration; + +use lance_context_merge::demand::ShardMerged; +use lance_context_merge::Coordinator; +use tokio::sync::Semaphore; + +/// Base-table index read for the merged watermarks plus one etcd put per shard. +const MERGED_PUBLISH_TIMEOUT: Duration = Duration::from_secs(15); + +static MERGED_PUBLISH_SLOTS: std::sync::LazyLock> = + std::sync::LazyLock::new(|| Arc::new(Semaphore::new(4))); + +/// After a successful WAL catch-up has released ownership, publish every +/// shard's merged watermark. Detached and bounded; never on the release +/// path. Each shard is a separate monotonic put, so a table with hundreds of +/// shards never approaches the etcd transaction op limit and a concurrent +/// writer flush on one shard cannot fail the others. +pub(crate) fn spawn_merged_demand(coordinator: Coordinator, target: String, uri: String) { + let Ok(permit) = MERGED_PUBLISH_SLOTS.clone().try_acquire_owned() else { + metrics::counter!("master_demand_merged_watermarks_total", "result" => "dropped") + .increment(1); + return; + }; + tokio::spawn(async move { + let _permit = permit; + let result = tokio::time::timeout(MERGED_PUBLISH_TIMEOUT, async { + let marks = lance_context_core::rollout_append::merged_watermarks(&uri) + .await + .map_err(|e| e.to_string())?; + let merged: Vec = marks + .into_iter() + .map(|(shard, generation)| ShardMerged { + shard: shard.to_string(), + merged_through: generation, + }) + .collect(); + let mut written = 0usize; + for shard in &merged { + if coordinator.publish_merged_watermark(&target, shard).await? { + written += 1; + } + } + Ok::<(usize, usize), String>((written, merged.len())) + }) + .await; + let (label, count) = match result { + Ok(Ok((written, total))) => (if written > 0 { "written" } else { "unchanged" }, total), + Ok(Err(error)) => { + tracing::warn!(target = %target, %error, "merged watermark publish failed"); + ("failed", 0) + } + Err(_) => { + tracing::warn!(target = %target, "merged watermark publish timed out"); + ("timeout", 0) + } + }; + metrics::counter!("master_demand_merged_watermarks_total", "result" => label) + .increment(count.max(1) as u64); + }); +} diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs index f7adb65..cbbc866 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/crates/lance-context-master/src/lib.rs @@ -3,6 +3,7 @@ pub mod admission; pub mod catchup; pub mod config; +pub mod demand_publish; pub use lance_context_merge::demand; pub mod discovery; pub mod eligibility; diff --git a/crates/lance-context-master/src/maintenance_execution.rs b/crates/lance-context-master/src/maintenance_execution.rs index 28e132e..2ac2a79 100644 --- a/crates/lance-context-master/src/maintenance_execution.rs +++ b/crates/lance-context-master/src/maintenance_execution.rs @@ -321,39 +321,23 @@ where { return Err("maintenance ownership changed before completion".into()); } - // 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? - { + if !coordinator.release(&proof, &terminal).await? { return Err("maintenance completion claim lost".into()); } - if !merged.is_empty() { - metrics::counter!("master_demand_merged_watermarks_total").increment(merged.len() as u64); + // A successful WAL catch-up reclaimed generations. Publish the merged + // watermark per shard *after* release, detached and bounded, so the + // release path never waits on a storage read, never contends with + // writer flushes on the demand keys, and never exceeds the etcd + // transaction size for a table with many shards. Each shard is its own + // monotonic per-dimension join, so a lost or late publish cannot + // misstate demand; the next scan reconciles it. Other maintenance kinds + // do not move WAL watermarks. + if maintenance == MaintenanceKind::Catchup && outcome.is_ok() { + crate::demand_publish::spawn_merged_demand( + coordinator.clone(), + running.target.clone(), + uri.clone(), + ); } outcome } diff --git a/crates/lance-context-master/src/rollout_append.rs b/crates/lance-context-master/src/rollout_append.rs index c9acb71..c7014c5 100644 --- a/crates/lance-context-master/src/rollout_append.rs +++ b/crates/lance-context-master/src/rollout_append.rs @@ -634,19 +634,43 @@ mod tests { .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(); + let events_prefix = format!( + "{}/demand-events/{encoded}/", + first.config.etcd.etcd_prefix.trim_end_matches('/') + ); + // The merged-watermark publish is detached from release; wait for + // both shards' records to carry a merged mark, with a bound. + let deadline = std::time::Instant::now() + Duration::from_secs(10); + let events = loop { + let events = first + .task_store + .etcd_client() + .clone() + .get( + events_prefix.as_str(), + Some(etcd_client::GetOptions::new().with_prefix()), + ) + .await + .unwrap(); + let merged_count = events + .kvs() + .iter() + .filter(|kv| { + serde_json::from_slice::( + kv.value(), + ) + .is_ok_and(|e| e.merged_through.is_some()) + }) + .count(); + if merged_count == 2 { + break events; + } + assert!( + std::time::Instant::now() < deadline, + "{target}: merged watermarks never published ({merged_count}/2)" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + }; assert_eq!(events.kvs().len(), 2, "{target}"); let mut table = lance_context_merge::demand::TableDemand::default(); for kv in events.kvs() { diff --git a/crates/lance-context-merge/src/demand.rs b/crates/lance-context-merge/src/demand.rs index 25147f7..143d4a9 100644 --- a/crates/lance-context-merge/src/demand.rs +++ b/crates/lance-context-merge/src/demand.rs @@ -218,25 +218,25 @@ pub struct ShardMerged { } 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( + /// Publish one shard's merged watermark as its own monotonic put. + /// + /// Reads the stored record, advances only `merged_through` (never lowers + /// anything the writer set), and writes under a mod-revision compare. A + /// writer flush racing on the same key fails the compare; the put is + /// retried a bounded number of times against the fresh record so a + /// normal flush cannot make merged progress disappear. A shard with no + /// stored event gets a fresh executor-sourced record with + /// `sealed_through = merged_through` at epoch 0, which any real writer + /// epoch supersedes. Returns `Ok(false)` when the stored merged mark is + /// already at or past the new one. + pub async fn publish_merged_watermark( &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)); - } + shard: &ShardMerged, + ) -> Result { + let key = self.demand_event_key(target, &shard.shard); 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); + for _ in 0..5 { let current = client .get(key.as_str(), None) .await @@ -245,42 +245,49 @@ impl Coordinator { .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, - }, + // The executor merged generations numbered under the stored + // writer's epoch, so report under that epoch. + let epoch = stored.as_ref().map_or(0, |s| s.writer_epoch); + let event = DemandEvent { + v: SCHEMA_VERSION, + shard: shard.shard.clone(), + sealed_through: stored + .as_ref() + .map_or(shard.merged_through, |s| s.sealed_through), + sealed_bytes_through: stored.as_ref().map_or(0, |s| s.sealed_bytes_through), + merged_through: Some(shard.merged_through), + flushed_at_ms: stored + .as_ref() + .map_or_else(chrono_now_ms, |s| s.flushed_at_ms), + writer_epoch: epoch, + source: EventSource::Executor, + }; + let Some(next) = join_events(stored.as_ref(), &event) else { + return Ok(false); }; let value = serde_json::to_vec(&next).map_err(|e| e.to_string())?; - compares.push(match current.kvs().first() { + let compare = 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)); + }; + let response = client + .txn( + etcd_client::Txn::new() + .when(vec![compare]) + .and_then(vec![TxnOp::put(key.as_str(), value, None)]), + ) + .await + .map_err(|e| e.to_string())?; + if response.succeeded() { + return Ok(true); + } } - Ok((compares, operations)) + Err(format!( + "merged watermark for {target}/{} kept losing to concurrent writes", + shard.shard + )) } fn demand_event_key(&self, target: &str, shard: &str) -> String { @@ -291,68 +298,100 @@ impl Coordinator { ) } - /// Publish a shard watermark. Monotonic: if the stored event already has a - /// `(writer_epoch, sealed_through)` at least as new, nothing is written and - /// `Ok(false)` is returned. Never deletes. Safe to call on every flush. + /// Publish a shard watermark. The stored record is the per-dimension join + /// of everything published so far: `sealed` and `merged` are each the max + /// of `(writer_epoch, value)` across stored and new, independently, so a + /// writer event with `merged_through: None` keeps the executor's merged + /// mark and an executor event with an older sealed still advances merged. + /// Writes under a mod-revision compare and retries a bounded number of + /// times against the fresh record when a concurrent writer wins the race, + /// so no publisher's progress is lost to another's. Returns `Ok(false)` + /// when nothing advanced. Never deletes. Safe on every flush. pub async fn publish_demand_event(&self, target: &str, event: &DemandEvent) -> Result { let key = self.demand_event_key(target, &event.shard); - let value = serde_json::to_vec(event).map_err(|e| e.to_string())?; let mut client = self.client.clone(); - let current = client - .get(key.as_str(), None) - .await - .map_err(|e| e.to_string())?; - let Some(kv) = current.kvs().first() else { + for _ in 0..5 { + 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 Some(next) = join_events(stored.as_ref(), event) else { + return Ok(false); + }; + let value = serde_json::to_vec(&next).map_err(|e| e.to_string())?; + let compare = 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), + }; let response = client .txn( etcd_client::Txn::new() - .when(vec![Compare::version(key.as_str(), CompareOp::Equal, 0)]) + .when(vec![compare]) .and_then(vec![TxnOp::put(key.as_str(), value, None)]), ) .await .map_err(|e| e.to_string())?; - return Ok(response.succeeded()); - }; - let stored: DemandEvent = match serde_json::from_slice(kv.value()) { - Ok(stored) => stored, - // An unreadable record is replaced rather than left poisoning the key. - Err(_) => DemandEvent { - v: 0, - shard: event.shard.clone(), - sealed_through: 0, - sealed_bytes_through: 0, - merged_through: None, - flushed_at_ms: 0, - writer_epoch: 0, - source: event.source, - }, - }; - let newer = ( - event.writer_epoch, - event.sealed_through, - event.merged_through, - ) > ( - stored.writer_epoch, - stored.sealed_through, - stored.merged_through, - ); - if !newer { - return Ok(false); + if response.succeeded() { + return Ok(true); + } } - let response = client - .txn( - etcd_client::Txn::new() - .when(vec![Compare::mod_revision( - key.as_str(), - CompareOp::Equal, - kv.mod_revision(), - )]) - .and_then(vec![TxnOp::put(key.as_str(), value, None)]), - ) - .await - .map_err(|e| e.to_string())?; - Ok(response.succeeded()) + Err(format!( + "demand event for {target}/{} kept losing to concurrent writes", + event.shard + )) + } +} + +/// Per-dimension join of a stored event with a new one. Returns `None` when +/// the stored record already dominates on every dimension. `sealed` is +/// ordered by `(writer_epoch, sealed_through)` then bytes; `merged` by +/// `(writer_epoch, merged_through)`; each taken independently. The result +/// keeps the stored record's writer fields unless the new event's sealed +/// mark wins. A merged mark whose epoch differs from the resulting sealed +/// epoch is dropped: the wire record stores merged as a bare value under the +/// record's epoch, and a cross-epoch merged mark says nothing about this +/// epoch's generations anyway. +pub fn join_events(stored: Option<&DemandEvent>, event: &DemandEvent) -> Option { + let Some(stored) = stored else { + return Some(event.clone()); + }; + let stored_sealed = (stored.writer_epoch, stored.sealed_through); + let event_sealed = (event.writer_epoch, event.sealed_through); + let stored_merged = stored.merged_through.map(|v| (stored.writer_epoch, v)); + let event_merged = event.merged_through.map(|v| (event.writer_epoch, v)); + let sealed_wins = event_sealed > stored_sealed + || (event_sealed == stored_sealed + && event.sealed_bytes_through > stored.sealed_bytes_through); + let merged_wins = match (stored_merged, event_merged) { + (_, None) => false, + (None, Some(_)) => true, + (Some(s), Some(e)) => e > s, + }; + if !sealed_wins && !merged_wins { + return None; + } + let best_merged = match (stored_merged, event_merged) { + (Some(s), Some(e)) => Some(s.max(e)), + (s, e) => s.or(e), + }; + let mut next = if sealed_wins { + event.clone() + } else { + stored.clone() + }; + next.merged_through = best_merged + .filter(|(epoch, _)| *epoch == next.writer_epoch) + .map(|(_, value)| value); + if merged_wins && !sealed_wins { + next.source = event.source; } + Some(next) } #[cfg(test)] @@ -555,76 +594,85 @@ mod tests { #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] - async fn release_changes_only_advance_merged_and_keep_writer_fields() { + async fn merged_publish_advances_only_merged_and_survives_writer_races() { 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()), + format!("/demand-merged/{}", 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, - }], - ) + let shard = |m: u64| ShardMerged { + shard: "s".into(), + merged_through: m, + }; + assert!(coordinator + .publish_merged_watermark("t", &shard(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() + .unwrap()); + let key = coordinator.demand_event_key("t", "s"); + let read = |client: etcd_client::Client, key: String| async move { + serde_json::from_slice::( + client.clone().get(key, None).await.unwrap().kvs()[0].value(), + ) + .unwrap() }; - let after = read(client.clone(), coordinator.clone()).await; + let after = read(client.clone(), key.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( + assert!( + !coordinator + .publish_merged_watermark("t", &shard(12)) + .await + .unwrap(), + "lower merged is a no-op" + ); + // Race: a writer flush lands between the executor's read and its CAS. + // Simulate by interleaving: publish writer sealed=25 then merged=18; + // both must survive regardless of which the stored record saw first. + let c1 = coordinator.clone(); + let c2 = coordinator.clone(); + let (w, m) = tokio::join!( + async move { + c1.publish_demand_event("t", &ev("s", 25, None, 3, EventSource::Writer)) + .await + .unwrap() + }, + async move { c2.publish_merged_watermark("t", &shard(18)).await.unwrap() } + ); + assert!(w && m, "both writes must succeed (writer={w}, merged={m})"); + let stored = read(client.clone(), key.clone()).await; + assert_eq!( + (stored.sealed_through, stored.merged_through), + (25, Some(18)) + ); + // Unknown shard: fresh executor record, sealed == merged, epoch 0. + assert!(coordinator + .publish_merged_watermark( "t", - &[ShardMerged { + &ShardMerged { shard: "new".into(), - merged_through: 7, - }], + 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); + .unwrap()); + let fresh = read(client.clone(), coordinator.demand_event_key("t", "new")).await; + assert_eq!( + ( + fresh.sealed_through, + fresh.merged_through, + fresh.writer_epoch + ), + (7, Some(7), 0) + ); } #[test] diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index bc9f015..e89ec21 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -290,32 +290,11 @@ 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)),