diff --git a/crates/lance-context-core/src/rollout_append.rs b/crates/lance-context-core/src/rollout_append.rs index 235e122..42a42fa 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 fc92960..28e132e 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 54c7c5b..c9acb71 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 6857bae..fd62006 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. @@ -52,26 +59,98 @@ fn schema_version() -> u32 { SCHEMA_VERSION } +/// A watermark with the epoch that produced it. Ordered lexicographically: +/// a higher epoch wins regardless of value; within an epoch the higher +/// value wins. This is the only ordering used anywhere in the fold. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default)] +pub struct Mark { + pub epoch: u64, + pub value: u64, +} + +/// How many `(generation, flushed_at)` pairs a shard keeps for age tracking. +/// Only the lowest unmerged generations matter for "oldest pending", so the +/// set is truncated to the smallest `SEALED_TIMES_CAP` generation numbers. +/// Union-then-truncate-to-smallest is order independent. +pub const SEALED_TIMES_CAP: usize = 64; + +/// One shard's folded state. `sealed`, `merged` and `sealed_bytes` are pure +/// joins (max), so `pending_generations` and `pending_bytes` are exactly +/// order independent (spec invariant 5a). `sealed` and `merged` are +/// independent: nothing about one lowers or replaces the other. +/// `sealed_times` is a join (union with min time) that is truncated to the +/// lowest `SEALED_TIMES_CAP` pending generations; `oldest_pending_ms` is +/// therefore exact while a shard has at most that many pending generations +/// and a best-known lower bound beyond it. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct ShardDemand { - pub sealed: u64, + /// Highest `(writer_epoch, sealed_through)` seen. + pub sealed: Mark, + /// Highest `(writer_epoch, merged_through)` seen; `None` until an + /// executor or scan reports one. + #[serde(default)] + pub merged: Option, + /// Max `sealed_bytes_through` seen at the current `sealed` epoch. + #[serde(default)] pub sealed_bytes: u64, - pub merged: u64, - pub epoch: u64, - pub flushed_at_ms: i64, + /// Flush times of sealed generations at the current `sealed` epoch, for + /// the lowest `SEALED_TIMES_CAP` generations above `merged`. Pruned as + /// `merged` advances. The oldest pending generation's time is the first + /// entry; a newer flush adds an entry but never changes the first. + #[serde(default)] + pub sealed_times: BTreeMap, } impl ShardDemand { + /// Generations sealed but not merged. Comparable only when both marks + /// are at the same epoch. A `merged` from an older epoch says nothing + /// about this epoch's generations; a `merged` from a newer epoch means + /// the shard was replaced and nothing from the old epoch is pending. pub fn pending_generations(&self) -> u64 { - self.sealed.saturating_sub(self.merged) + match self.merged { + None => self.sealed.value, + Some(merged) if merged.epoch == self.sealed.epoch => { + self.sealed.value.saturating_sub(merged.value) + } + Some(merged) if merged.epoch > self.sealed.epoch => 0, + Some(_) => self.sealed.value, + } + } + + /// Merged generation at the sealed epoch, if known. + fn merged_at_sealed_epoch(&self) -> Option { + match self.merged { + Some(m) if m.epoch == self.sealed.epoch => Some(m.value), + Some(m) if m.epoch > self.sealed.epoch => Some(u64::MAX), + _ => None, + } + } + + fn prune_sealed_times(&mut self) { + if let Some(merged) = self.merged_at_sealed_epoch() { + self.sealed_times + .retain(|generation, _| *generation > merged); + } + while self.sealed_times.len() > SEALED_TIMES_CAP { + self.sealed_times.pop_last(); + } + } + + /// Flush time of the oldest generation still unmerged, if known. + pub fn oldest_pending_ms(&self) -> Option { + if self.pending_generations() == 0 { + return None; + } + self.sealed_times.values().next().copied() } } /// The per-table record under `P/demand/`. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct TableDemand { #[serde(default = "schema_version")] pub v: u32, + #[serde(default)] pub shards: BTreeMap, #[serde(default)] pub fragment_count: u64, @@ -81,25 +160,48 @@ pub struct TableDemand { pub uncovered_fragments: u64, #[serde(default)] pub missing_fragments: bool, - /// etcd revision of the newest input folded in. + /// etcd revision of the newest input folded in. Freshness bookkeeping + /// only; never an ordering input. #[serde(default)] pub observed_revision: i64, #[serde(default)] pub updated_ms: i64, } +impl Default for TableDemand { + fn default() -> Self { + Self { + v: SCHEMA_VERSION, + shards: BTreeMap::new(), + fragment_count: 0, + index_stale: false, + uncovered_fragments: 0, + missing_fragments: false, + observed_revision: 0, + updated_ms: 0, + } + } +} + /// Why an event did not change the table record. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Ignored { - /// A lower writer epoch than the stored one. - StaleEpoch, - /// Same epoch, nothing newer than what is stored. + /// Neither `sealed` nor `merged` nor bytes advanced. NotNewer, - /// A scan tried to lower `sealed_through` without a higher epoch. - ScanCannotLower, + /// The event's schema major version is not understood. + UnknownVersion(u32), } impl TableDemand { + /// Reject records from a future schema. `v == 0` is a pre-versioning + /// record and is read as version 1. + pub fn check_version(v: u32) -> std::result::Result<(), Ignored> { + match v { + 0 | SCHEMA_VERSION => Ok(()), + other => Err(Ignored::UnknownVersion(other)), + } + } + pub fn pending_generations(&self) -> u64 { self.shards .values() @@ -118,92 +220,246 @@ impl TableDemand { .sum() } - /// Oldest flush time among shards with pending generations. + /// Flush time of the oldest generation still unmerged on any shard. + /// Does not move when newer generations arrive. pub fn oldest_pending_ms(&self) -> Option { self.shards .values() - .filter(|s| s.pending_generations() > 0) - .map(|s| s.flushed_at_ms) + .filter_map(ShardDemand::oldest_pending_ms) .min() } /// Fold one event. Returns `Err(reason)` when the record is unchanged. /// - /// Rules (spec §1.1): - /// * a strictly higher `writer_epoch` replaces the shard entirely; - /// * a lower epoch is ignored; - /// * same epoch: `sealed` and `merged` only move forward, except that a - /// `Scan` may never lower `sealed` and no source may lower `merged`; - /// * `flushed_at_ms` follows the newest `sealed`. + /// `sealed` and `merged` are joined independently by `(epoch, value)`. + /// Source is irrelevant to ordering: writer, executor and scan all report + /// absolute watermarks and the max wins, so a stale scan cannot lower + /// anything - a lower value simply loses the max. pub fn fold( &mut self, event: &DemandEvent, revision: i64, now_ms: i64, ) -> std::result::Result<(), Ignored> { - let entry = self.shards.entry(event.shard.clone()); - let changed = match entry { - std::collections::btree_map::Entry::Vacant(slot) => { - slot.insert(ShardDemand { - sealed: event.sealed_through, - sealed_bytes: event.sealed_bytes_through, - merged: event.merged_through.unwrap_or(0), - epoch: event.writer_epoch, - flushed_at_ms: event.flushed_at_ms, - }); - true + Self::check_version(event.v)?; + let shard = self.shards.entry(event.shard.clone()).or_default(); + let mut changed = false; + + let sealed = Mark { + epoch: event.writer_epoch, + value: event.sealed_through, + }; + // A time entry is only useful for a generation that is still + // pending; recording one for an already-merged generation would be + // pruned straight away and must not count as a change. + let records_time = event.sealed_through > 0 + && shard + .merged + .is_none_or(|m| m.epoch != event.writer_epoch || event.sealed_through > m.value); + match sealed.cmp(&shard.sealed) { + std::cmp::Ordering::Greater => { + if sealed.epoch > shard.sealed.epoch { + // New writer epoch: byte and age bookkeeping restart. + shard.sealed_bytes = event.sealed_bytes_through; + shard.sealed_times.clear(); + } else { + shard.sealed_bytes = shard.sealed_bytes.max(event.sealed_bytes_through); + } + shard.sealed = sealed; + if records_time { + shard + .sealed_times + .entry(event.sealed_through) + .and_modify(|t| *t = (*t).min(event.flushed_at_ms)) + .or_insert(event.flushed_at_ms); + } + changed = true; } - std::collections::btree_map::Entry::Occupied(mut slot) => { - let shard = slot.get_mut(); - if event.writer_epoch < shard.epoch { - return Err(Ignored::StaleEpoch); + std::cmp::Ordering::Equal => { + if event.sealed_bytes_through > shard.sealed_bytes { + shard.sealed_bytes = event.sealed_bytes_through; + changed = true; } - if event.writer_epoch > shard.epoch { - // A new writer epoch restarts the shard's numbering - // authority; take the event wholesale but never lose a - // merged watermark the executor already recorded. - *shard = ShardDemand { - sealed: event.sealed_through, - sealed_bytes: event.sealed_bytes_through, - merged: event.merged_through.unwrap_or(0).max(shard.merged), - epoch: event.writer_epoch, - flushed_at_ms: event.flushed_at_ms, - }; - true - } else { - let mut changed = false; - if event.sealed_through > shard.sealed { - shard.sealed = event.sealed_through; - shard.sealed_bytes = shard.sealed_bytes.max(event.sealed_bytes_through); - shard.flushed_at_ms = event.flushed_at_ms; + // Same generation seen again: keep the earliest time. + if records_time { + let entry = shard.sealed_times.entry(event.sealed_through); + if let std::collections::btree_map::Entry::Occupied(mut slot) = entry { + if event.flushed_at_ms < *slot.get() { + slot.insert(event.flushed_at_ms); + changed = true; + } + } else { + entry.or_insert(event.flushed_at_ms); changed = true; - } else if event.sealed_through < shard.sealed - && event.source == EventSource::Scan - { - return Err(Ignored::ScanCannotLower); } - if let Some(merged) = event.merged_through { - if merged > shard.merged { - shard.merged = merged; + } + } + std::cmp::Ordering::Less => { + // An older sealed watermark at the same epoch still tells us + // when that generation was flushed, which matters for age. + if sealed.epoch == shard.sealed.epoch && records_time { + let entry = shard.sealed_times.entry(event.sealed_through); + match entry { + std::collections::btree_map::Entry::Occupied(mut slot) => { + if event.flushed_at_ms < *slot.get() { + slot.insert(event.flushed_at_ms); + changed = true; + } + } + std::collections::btree_map::Entry::Vacant(slot) => { + slot.insert(event.flushed_at_ms); changed = true; } } - if !changed { - return Err(Ignored::NotNewer); - } - true } } - }; + } + + if let Some(merged_through) = event.merged_through { + let merged = Mark { + epoch: event.writer_epoch, + value: merged_through, + }; + if shard.merged.is_none_or(|current| merged > current) { + shard.merged = Some(merged); + changed = true; + } + } + if changed { + shard.prune_sealed_times(); self.observed_revision = self.observed_revision.max(revision); self.updated_ms = now_ms; + Ok(()) + } else { + Err(Ignored::NotNewer) } - Ok(()) } } +/// Per-dimension join of a stored event with a new one. Returns `None` when +/// the stored record already dominates on every dimension. The result keeps +/// the stored record's writer fields unless the new event's sealed mark +/// wins, in which case the new event's epoch, bytes and flush time apply. +pub fn join_events(stored: Option<&DemandEvent>, event: &DemandEvent) -> Option { + let Some(stored) = stored else { + return Some(event.clone()); + }; + let stored_sealed = Mark { + epoch: stored.writer_epoch, + value: stored.sealed_through, + }; + let event_sealed = Mark { + epoch: event.writer_epoch, + value: event.sealed_through, + }; + let stored_merged = stored.merged_through.map(|v| Mark { + epoch: stored.writer_epoch, + value: v, + }); + let event_merged = event.merged_through.map(|v| Mark { + epoch: event.writer_epoch, + value: 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() + }; + // The wire record stores merged as a bare value under the record's + // epoch. A merged mark from a different epoch than the record's sealed + // epoch cannot be expressed there, and says nothing about this epoch's + // generations anyway (see ShardDemand::pending_generations), so drop it. + next.merged_through = best_merged + .filter(|m| m.epoch == next.writer_epoch) + .map(|m| m.value); + if merged_wins && !sealed_wins { + next.source = event.source; + } + Some(next) +} + +/// 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. Uses the same per-dimension join as + /// `publish_demand_event`; shards that would not advance emit nothing. + /// A shard with no stored event gets a fresh executor-sourced record + /// with `sealed_through = merged_through` (the writer's next flush + /// raises it) at epoch 0, which any real writer epoch supersedes. + 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()); + // 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(now_ms, |s| s.flushed_at_ms), + writer_epoch: epoch, + source: EventSource::Executor, + }; + let Some(next) = join_events(stored.as_ref(), &event) else { + continue; + }; + 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!( @@ -212,18 +468,21 @@ 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 (`join_events`): a writer event + /// with `merged_through: None` keeps the executor's merged mark, and an + /// executor event with an older sealed still advances merged. Returns + /// `Ok(false)` when nothing advanced. Never deletes. Safe on every flush. pub async fn publish_demand_event(&self, target: &str, event: &DemandEvent) -> Result { + TableDemand::check_version(event.v).map_err(|e| format!("{e:?}"))?; 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 { + let value = serde_json::to_vec(event).map_err(|e| e.to_string())?; let response = client .txn( etcd_client::Txn::new() @@ -234,32 +493,12 @@ impl Coordinator { .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 { + // An unreadable record is treated as empty and replaced. + let stored: Option = 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 response = client .txn( etcd_client::Txn::new() @@ -325,81 +564,269 @@ mod tests { out } - #[test] - fn fold_is_order_independent_and_idempotent() { - let events = vec![ - ev("a", 10, None, 1, EventSource::Writer), - ev("a", 12, Some(8), 1, EventSource::Executor), - ev("a", 11, Some(11), 1, EventSource::Executor), - ev("b", 3, None, 2, EventSource::Writer), - ev("b", 5, None, 2, EventSource::Scan), - ]; - let canonical = fold_all(&events); - assert_eq!(canonical.shards["a"].sealed, 12); - assert_eq!(canonical.shards["a"].merged, 11); - assert_eq!(canonical.shards["b"].sealed, 5); - assert_eq!(canonical.pending_generations(), 1 + 5); - for perm in permutations(&events) { + fn assert_order_independent(events: &[DemandEvent]) -> TableDemand { + let canonical = fold_all(events); + for perm in permutations(events) { assert_eq!(fold_all(&perm), canonical, "order {perm:?}"); - // Duplicating every event changes nothing. let doubled: Vec<_> = perm.iter().flat_map(|e| [e.clone(), e.clone()]).collect(); - assert_eq!(fold_all(&doubled), canonical); + assert_eq!(fold_all(&doubled), canonical, "doubled {perm:?}"); } - // Replaying the final state onto itself is a no-op. let mut again = canonical.clone(); - for event in &events { + for event in events { let _ = again.fold(event, 99, 99); } again.observed_revision = 0; again.updated_ms = 0; - assert_eq!(again, canonical); + assert_eq!(again, canonical, "replay onto final state"); + canonical } #[test] - fn scan_may_raise_but_not_lower_without_higher_epoch() { + fn fold_is_order_independent_and_idempotent() { + let table = assert_order_independent(&[ + ev("a", 10, None, 1, EventSource::Writer), + ev("a", 12, Some(8), 1, EventSource::Executor), + ev("a", 11, Some(11), 1, EventSource::Executor), + ev("b", 3, None, 2, EventSource::Writer), + ev("b", 5, None, 2, EventSource::Scan), + ]); + assert_eq!(table.shards["a"].sealed.value, 12); + assert_eq!(table.shards["a"].merged.unwrap().value, 11); + assert_eq!(table.shards["b"].sealed.value, 5); + assert_eq!(table.pending_generations(), 1 + 5); + } + + /// Review counterexamples: an old-epoch merge beside a new-epoch flush, + /// and a scan carrying an older sealed but newer merged. Both must fold + /// to one answer regardless of order. + #[test] + fn review_counterexamples_fold_to_one_answer() { + let table = assert_order_independent(&[ + ev("s", 100, Some(40), 1, EventSource::Executor), + ev("s", 60, None, 2, EventSource::Writer), + ]); + // Epoch 2 sealed 60; merged is from epoch 1 and says nothing about + // epoch 2's generations: all 60 pending. + assert_eq!(table.pending_generations(), 60); + + let table = assert_order_independent(&[ + ev("s", 30, Some(10), 1, EventSource::Writer), + ev("s", 20, Some(20), 1, EventSource::Scan), + ]); + // sealed = max(30, 20) = 30, merged = max(10, 20) = 20: 10 pending. + assert_eq!(table.pending_generations(), 10); + } + + /// Review finding 3: a writer event without `merged_through` must not + /// discard a previously known merged watermark. + #[test] + fn writer_event_never_forgets_merged() { + let table = assert_order_independent(&[ + ev("s", 10, Some(8), 1, EventSource::Executor), + ev("s", 12, None, 1, EventSource::Writer), + ]); + assert_eq!(table.shards["s"].merged.unwrap().value, 8); + assert_eq!(table.pending_generations(), 4); + } + + #[test] + fn scan_cannot_lower_anything() { let mut table = TableDemand::default(); table - .fold(&ev("a", 20, None, 3, EventSource::Writer), 1, 0) + .fold(&ev("a", 20, Some(5), 3, EventSource::Writer), 1, 0) + .unwrap(); + // Lower sealed and lower merged from a scan at the same epoch: a + // time entry for generation 15 is learned (useful for age) but no + // watermark moves. + table + .fold(&ev("a", 15, Some(2), 3, EventSource::Scan), 2, 0) .unwrap(); + assert_eq!(table.shards["a"].sealed.value, 20); + assert_eq!(table.shards["a"].merged.unwrap().value, 5); assert_eq!( - table.fold(&ev("a", 15, None, 3, EventSource::Scan), 2, 0), - Err(Ignored::ScanCannotLower) + table.fold(&ev("a", 15, Some(2), 3, EventSource::Scan), 3, 0), + Err(Ignored::NotNewer), + "and the second time it is a pure no-op" ); - assert_eq!(table.shards["a"].sealed, 20); table - .fold(&ev("a", 25, None, 3, EventSource::Scan), 3, 0) + .fold(&ev("a", 25, None, 3, EventSource::Scan), 4, 0) .unwrap(); - assert_eq!(table.shards["a"].sealed, 25); - // A retired shard re-observed by a scan under a higher epoch may lower. + assert_eq!(table.shards["a"].sealed.value, 25); + // A higher epoch from a scan replaces sealed; the epoch-3 merged mark + // stays recorded but no longer applies to epoch 4 generations. table - .fold(&ev("a", 2, None, 4, EventSource::Scan), 4, 0) + .fold(&ev("a", 2, None, 4, EventSource::Scan), 5, 0) .unwrap(); - assert_eq!(table.shards["a"].sealed, 2); - assert_eq!(table.shards["a"].epoch, 4); + assert_eq!(table.shards["a"].sealed, Mark { epoch: 4, value: 2 }); + assert_eq!(table.shards["a"].pending_generations(), 2); } #[test] - fn epoch_rules_and_merged_watermark_preservation() { + fn lower_epoch_events_do_not_move_marks() { let mut table = TableDemand::default(); table .fold(&ev("a", 50, Some(40), 2, EventSource::Executor), 1, 0) .unwrap(); assert_eq!( - table.fold(&ev("a", 60, None, 1, EventSource::Writer), 2, 0), - Err(Ignored::StaleEpoch) + table.fold(&ev("a", 60, Some(55), 1, EventSource::Writer), 2, 0), + Err(Ignored::NotNewer) + ); + assert_eq!( + table.shards["a"].sealed, + Mark { + epoch: 2, + value: 50 + } ); assert_eq!( - table.fold(&ev("a", 50, Some(39), 2, EventSource::Executor), 3, 0), - Err(Ignored::NotNewer), - "merged never moves backwards" + table.shards["a"].merged, + Some(Mark { + epoch: 2, + value: 40 + }) ); - // New epoch takes the event but keeps the recorded merged watermark. + } + + /// Review finding 4: the age clock is set by the oldest pending + /// generation and does not reset when newer generations arrive. + #[test] + fn oldest_pending_does_not_reset_on_newer_flushes() { + let mut table = TableDemand::default(); + let flush = |g: u64, t: i64| DemandEvent { + flushed_at_ms: t, + ..ev("s", g, None, 1, EventSource::Writer) + }; + table.fold(&flush(1, 1000), 1, 0).unwrap(); + table.fold(&flush(2, 2000), 2, 0).unwrap(); + table.fold(&flush(3, 3000), 3, 0).unwrap(); + assert_eq!(table.oldest_pending_ms(), Some(1000)); + // Merging generation 1 moves the clock to generation 2's time. + table + .fold(&ev("s", 3, Some(1), 1, EventSource::Executor), 4, 0) + .unwrap(); + assert_eq!(table.oldest_pending_ms(), Some(2000)); table - .fold(&ev("a", 5, None, 3, EventSource::Writer), 4, 0) + .fold(&ev("s", 3, Some(3), 1, EventSource::Executor), 5, 0) .unwrap(); - assert_eq!(table.shards["a"].merged, 40); - assert_eq!(table.shards["a"].sealed, 5); - assert_eq!(table.shards["a"].pending_generations(), 0); + assert_eq!(table.oldest_pending_ms(), None); + // Independent of order below the cap. + assert_order_independent(&[ + flush(1, 1000), + flush(2, 2000), + flush(3, 3000), + ev("s", 3, Some(1), 1, EventSource::Executor), + ]); + } + + /// Above the cap, pending counts stay exact in every order; the age + /// clock is a lower bound and may be unknown, never wrong-direction. + #[test] + fn beyond_cap_counts_stay_exact_and_age_is_a_bound() { + let n = (SEALED_TIMES_CAP as u64) + 20; + let forward: Vec<_> = (1..=n) + .map(|g| DemandEvent { + flushed_at_ms: g as i64 * 10, + ..ev("s", g, None, 1, EventSource::Writer) + }) + .collect(); + let mut reverse = forward.clone(); + reverse.reverse(); + let f = fold_all(&forward); + let r = fold_all(&reverse); + assert_eq!(f.pending_generations(), r.pending_generations()); + assert_eq!(f.pending_generations(), n); + assert_eq!(f.pending_bytes(), r.pending_bytes()); + assert_eq!(f.oldest_pending_ms(), Some(10)); + assert_eq!(r.oldest_pending_ms(), Some(10)); + assert!(f.shards["s"].sealed_times.len() <= SEALED_TIMES_CAP); + } + + /// Deterministic pseudo-random streams over two shards, three epochs, + /// all sources, with and without merged. Small streams: every + /// permutation. Large streams: forward, reverse and several shuffles. + /// The exact fields (sealed/merged marks, pending counts, bytes) must + /// agree; the bounded age field must agree whenever it is exact. + #[test] + fn randomized_streams_fold_order_independently() { + let mut seed: u64 = 0x5eed_1234_abcd_0001; + let mut next = || { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + seed + }; + let gen_event = |r: &mut dyn FnMut() -> u64| { + let shard = if r().is_multiple_of(2) { "a" } else { "b" }; + let epoch = 1 + r() % 3; + let sealed = r() % 40; + let merged = if r().is_multiple_of(3) { + None + } else { + Some(r() % 40) + }; + let source = match r() % 3 { + 0 => EventSource::Writer, + 1 => EventSource::Executor, + _ => EventSource::Scan, + }; + DemandEvent { + v: 1, + shard: shard.into(), + sealed_through: sealed, + sealed_bytes_through: sealed * 7, + merged_through: merged, + flushed_at_ms: (r() % 1000) as i64, + writer_epoch: epoch, + source, + } + }; + let exact = |t: &TableDemand| { + t.shards + .iter() + .map(|(k, s)| { + ( + k.clone(), + s.sealed, + s.merged, + s.sealed_bytes, + s.pending_generations(), + ) + }) + .collect::>() + }; + for _ in 0..200 { + let events: Vec<_> = (0..5).map(|_| gen_event(&mut next)).collect(); + let canonical = fold_all(&events); + for perm in permutations(&events) { + let folded = fold_all(&perm); + assert_eq!(exact(&folded), exact(&canonical), "{perm:?}"); + assert_eq!( + folded.pending_generations(), + canonical.pending_generations() + ); + assert_eq!(folded.pending_bytes(), canonical.pending_bytes()); + // Below the cap, age is exact too. + assert_eq!(folded.oldest_pending_ms(), canonical.oldest_pending_ms()); + } + } + for _ in 0..50 { + let events: Vec<_> = (0..200).map(|_| gen_event(&mut next)).collect(); + let canonical = fold_all(&events); + let mut reversed = events.clone(); + reversed.reverse(); + let mut shuffled = events.clone(); + for i in (1..shuffled.len()).rev() { + shuffled.swap(i, (next() as usize) % (i + 1)); + } + for other in [reversed, shuffled] { + let folded = fold_all(&other); + assert_eq!(exact(&folded), exact(&canonical)); + assert_eq!( + folded.pending_generations(), + canonical.pending_generations() + ); + } + } } #[test] @@ -416,7 +843,6 @@ mod tests { .unwrap(); assert_eq!(table.pending_generations(), 14, "a: 6, b: 0, c: 8"); assert_eq!(table.pending_bytes(), 10_000 + 8_000); - assert_eq!(table.oldest_pending_ms(), Some(80)); assert_eq!(table.observed_revision, 9); assert_eq!(table.updated_ms, 300); assert_eq!( @@ -426,6 +852,29 @@ mod tests { assert_eq!(table.observed_revision, 9, "ignored events do not advance"); } + /// Review finding 5: default carries the schema version; unknown + /// versions are rejected rather than silently folded. + #[test] + fn schema_version_is_set_and_enforced() { + assert_eq!(TableDemand::default().v, SCHEMA_VERSION); + let mut table = TableDemand::default(); + let future = DemandEvent { + v: 999, + ..ev("s", 1, None, 1, EventSource::Writer) + }; + assert_eq!(table.fold(&future, 1, 0), Err(Ignored::UnknownVersion(999))); + assert!(table.shards.is_empty()); + let legacy = DemandEvent { + v: 0, + ..ev("s", 1, None, 1, EventSource::Writer) + }; + table.fold(&legacy, 1, 0).unwrap(); + let wire = serde_json::to_value(&table).unwrap(); + assert_eq!(wire["v"], SCHEMA_VERSION); + let back: TableDemand = serde_json::from_value(wire).unwrap(); + assert_eq!(back, table); + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn publish_is_monotonic_by_epoch_then_sealed_then_merged() { @@ -474,6 +923,120 @@ mod tests { ); } + /// Review finding 3 at the storage layer: a writer publish that carries + /// no `merged_through` must keep the stored merged watermark. + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn writer_publish_preserves_stored_merged() { + 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-preserve/{}", uuid::Uuid::new_v4()), + ); + assert!(coordinator + .publish_demand_event("t", &ev("s", 10, Some(8), 1, EventSource::Executor)) + .await + .unwrap()); + assert!(coordinator + .publish_demand_event("t", &ev("s", 12, None, 1, EventSource::Writer)) + .await + .unwrap()); + let key = coordinator.demand_event_key("t", "s"); + let stored: DemandEvent = + serde_json::from_slice(client.clone().get(key, None).await.unwrap().kvs()[0].value()) + .unwrap(); + assert_eq!(stored.sealed_through, 12); + assert_eq!(stored.merged_through, Some(8), "merged preserved"); + let mut table = TableDemand::default(); + table.fold(&stored, 1, 0).unwrap(); + assert_eq!(table.pending_generations(), 4); + // And a later executor publish with a lower sealed but higher + // merged is accepted on the merged dimension. + assert!(coordinator + .publish_demand_event("t", &ev("s", 10, Some(10), 1, EventSource::Executor)) + .await + .unwrap()); + let key = coordinator.demand_event_key("t", "s"); + let stored: DemandEvent = + serde_json::from_slice(client.clone().get(key, None).await.unwrap().kvs()[0].value()) + .unwrap(); + assert_eq!( + (stored.sealed_through, stored.merged_through), + (12, Some(10)) + ); + } + + #[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()), + ); + 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 key = coordinator.demand_event_key("t", "s"); + let after: DemandEvent = + serde_json::from_slice(client.clone().get(key, None).await.unwrap().kvs()[0].value()) + .unwrap(); + 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"); + let (c, o) = coordinator + .demand_release_changes( + "t", + &[ShardMerged { + shard: "s".into(), + merged_through: 12, + }], + ) + .await + .unwrap(); + assert!(c.is_empty() && o.is_empty(), "lower merged emits nothing"); + 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 6690d91..bc9f015 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)) diff --git a/docs/design/scheduler-p1-data-model.md b/docs/design/scheduler-p1-data-model.md index 882e01f..f72f440 100644 --- a/docs/design/scheduler-p1-data-model.md +++ b/docs/design/scheduler-p1-data-model.md @@ -31,21 +31,36 @@ P/demand-events// ``` Rules: -* `put` only; no deletes by writers. Each key is overwritten with a strictly - monotonic `(writer_epoch, sealed_through)`; a writer compares before writing - and skips if the stored pair is already ≥ its own. Loss or duplication of a - put cannot corrupt the count because the value is absolute. +* `put` only; no deletes by writers. The stored record is the **per‑dimension + join** of every event published so far: `sealed` and `merged` are each the + max of `(writer_epoch, value)` across stored and new, taken independently. + A writer event with `merged_through: None` therefore keeps the executor's + merged mark, and an executor event with an older `sealed_through` still + advances `merged`. Nothing about one dimension ever lowers or replaces the + other. Loss or duplication of a put cannot corrupt the count because every + value is absolute and the join is idempotent and commutative. +* Folding (`TableDemand::fold`) applies the same join. `pending` for a shard is + `sealed − merged` only when both marks are at the same epoch; a `merged` + from an older epoch says nothing about the current epoch's generations + (all pending), and a `merged` from a newer epoch means the shard was + replaced (none pending). +* `oldest_pending_ms` is the flush time of the lowest unmerged generation, + kept in a bounded per‑shard map (`SEALED_TIMES_CAP = 64`) that is pruned as + `merged` advances. A newer flush adds an entry and never resets the age + clock. Pending counts and bytes are exact in every order; the age field is + exact up to the cap and a lower bound beyond it. * Executors write `merged_through` for the shards they merged **in the same txn as the merge-crate `release`**, so demand and ownership agree. * The stats scan (source `scan`) writes the same record when it observes a - shard. Ordering is by the shard's own `(writer_epoch, sealed_through)` only; - etcd revisions and manifest versions are different clocks and are never - compared with each other. A scan may raise `sealed_through` at an equal - epoch and may replace the shard record only with a strictly greater - `writer_epoch` (a retired shard whose writer is gone). It may never lower a - watermark at an equal epoch, regardless of when the snapshot was written - (review point 3). `observed_revision` on the table record is bookkeeping for - "how fresh is this cache", not an ordering input. + shard. Source is irrelevant to ordering: writer, executor and scan all report + absolute watermarks and the per‑dimension max wins, so a stale scan cannot + lower anything — a lower value simply loses the max. etcd revisions and + manifest versions are different clocks and are never compared. + `observed_revision` on the table record is bookkeeping for "how fresh is + this cache", not an ordering input. +* Records carry `v`; `TableDemand::default()` sets it, a reader rejects any + version it does not understand (`Ignored::UnknownVersion`), and `v == 0` is + read as version 1. ### 1.2 Table record — maintained by the planner @@ -192,7 +207,12 @@ existing loop keeps running for every table throughout P1. ## 7. Tests P1 must land with 1. `fold` is idempotent and commutative over any permutation/duplication of - events (property test, no etcd). + events (property test, no etcd): every permutation of small streams and + forward/reverse/shuffled large randomized streams across shards, epochs, + sources and missing `merged` agree on sealed/merged marks, pending counts + and bytes; the age field agrees below the cap. Includes the review + counterexamples (old‑epoch merge beside new‑epoch flush; scan with older + sealed and newer merged) and "writer event never forgets merged". 2. A scan with an equal epoch cannot lower `sealed_through`; a higher epoch can. 3. Headroom is computed from assignments, not samples: an executor whose `bytes_free_sample` says 2 GiB but with 2 GiB reserved gets no placement.