diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 340d0b69..3cbc3ccc 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -634,6 +634,13 @@ impl GenericStore { self.base.count_merge_due().await } + /// This writer's shard watermark from its own manifest. Metadata only. + pub async fn own_shard_watermark( + &self, + ) -> LanceResult> { + self.base.own_shard_watermark().await + } + /// The shared-lock half of [`Self::cleanup_wal`]: seal, then read a /// budgeted prefix of flushed generations into memory. Callers holding a /// read lock run this while appends continue, then take the write lock diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 485f71e8..836d9b6b 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -28,7 +28,8 @@ mod storage; mod store; mod store_base; pub use store_base::{ - is_pending_generations_exceeded, DEFAULT_PENDING_GENERATIONS_MAX, PENDING_GENERATIONS_EXCEEDED, + is_pending_generations_exceeded, ShardWatermark, DEFAULT_PENDING_GENERATIONS_MAX, + PENDING_GENERATIONS_EXCEEDED, }; // Request/DTO conversions, exported so the server does not keep its own copies. diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index a1aedb90..4ef8893a 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -666,6 +666,13 @@ impl RolloutStore { self.base.count_merge_due().await } + /// This writer's shard watermark from its own manifest. Metadata only. + pub async fn own_shard_watermark( + &self, + ) -> LanceResult> { + self.base.own_shard_watermark().await + } + /// [`Self::prepare_merge_if_ready`], but seals the active memtable first — /// the time-triggered behavior of [`Self::cleanup_own_shard`]. pub async fn prepare_cleanup_merge( diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 3513cdf6..3d71bef6 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -215,6 +215,17 @@ pub const PENDING_GENERATIONS_EXCEEDED: &str = "too many pending WAL generations pub const DEFAULT_PENDING_GENERATIONS_MAX: usize = 4096; /// Whether `err` is a read refused by the pending-generations cap. +/// A writer shard's position, read from its own MemWAL manifest. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ShardWatermark { + pub shard_id: Uuid, + pub writer_epoch: u64, + /// Highest flushed generation number; 0 when none are flushed. + pub sealed_through: u64, + /// Flushed generations not yet merged into the base table. + pub pending_generations: usize, +} + pub fn is_pending_generations_exceeded(err: &LanceError) -> bool { err.to_string().contains(PENDING_GENERATIONS_EXCEEDED) } @@ -914,6 +925,32 @@ impl StorageBase { .await } + /// This writer's shard watermark from its own manifest: highest flushed + /// generation, writer epoch and flushed generation count. Metadata only. + /// `None` when the shard has no manifest yet. + pub async fn own_shard_watermark(&self) -> LanceResult> { + let manifest_store = ShardManifestStore::new( + self.dataset.object_store(None).await?, + &self.dataset.branch_location().path, + self.write_shard, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + Ok(manifest_store + .read_latest() + .await? + .map(|manifest| ShardWatermark { + shard_id: manifest.shard_id, + writer_epoch: manifest.writer_epoch, + sealed_through: manifest + .flushed_generations + .iter() + .map(|g| g.generation) + .max() + .unwrap_or(0), + pending_generations: manifest.flushed_generations.len(), + })) + } + /// Metadata-only count trigger for a coordinated sweeper. Never prepares /// batches or takes a merge-memory reservation just to enqueue work. pub async fn count_merge_due(&self) -> LanceResult { diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs index 779e4a8a..f7adb65f 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/crates/lance-context-master/src/lib.rs @@ -3,7 +3,7 @@ pub mod admission; pub mod catchup; pub mod config; -pub mod demand; +pub use lance_context_merge::demand; pub mod discovery; pub mod eligibility; pub mod error; diff --git a/crates/lance-context-master/src/demand.rs b/crates/lance-context-merge/src/demand.rs similarity index 71% rename from crates/lance-context-master/src/demand.rs rename to crates/lance-context-merge/src/demand.rs index 6aebeb3e..6857baed 100644 --- a/crates/lance-context-master/src/demand.rs +++ b/crates/lance-context-merge/src/demand.rs @@ -1,14 +1,19 @@ //! Per-table maintenance demand folded from per-shard watermarks. //! -//! Schema: `docs/design/scheduler-p1-data-model.md` §1. Everything here is a -//! pure function of its inputs; no etcd. The record under `P/demand/` -//! is a cache rebuilt from `demand-events/` plus the stats scan, so folding -//! must be idempotent and order-independent. +//! Schema: `docs/design/scheduler-p1-data-model.md` §1. The types and +//! `TableDemand::fold` are pure; `Coordinator::publish_demand_event` is the +//! one etcd write, a monotonic put keyed by `(writer_epoch, sealed_through)`. +//! The record under `P/demand/` is a cache rebuilt from +//! `demand-events/` plus the stats scan, so folding must be idempotent and +//! order-independent. use std::collections::BTreeMap; +use etcd_client::{Compare, CompareOp, TxnOp}; use serde::{Deserialize, Serialize}; +use crate::{Coordinator, Result}; + pub const SCHEMA_VERSION: u32 = 1; /// Who produced a watermark event. @@ -130,7 +135,12 @@ impl TableDemand { /// * 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`. - pub fn fold(&mut self, event: &DemandEvent, revision: i64, now_ms: i64) -> Result<(), Ignored> { + 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) => { @@ -193,6 +203,79 @@ impl TableDemand { } } +impl Coordinator { + fn demand_event_key(&self, target: &str, shard: &str) -> String { + let encoded: String = target.bytes().map(|b| format!("{b:02x}")).collect(); + format!( + "{}/demand-events/{encoded}/{shard}", + self.prefix.trim_end_matches('/') + ) + } + + /// 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. + 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 { + let response = client + .txn( + etcd_client::Txn::new() + .when(vec![Compare::version(key.as_str(), CompareOp::Equal, 0)]) + .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); + } + 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()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -343,6 +426,54 @@ mod tests { assert_eq!(table.observed_revision, 9, "ignored events do not advance"); } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn publish_is_monotonic_by_epoch_then_sealed_then_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, format!("/demand-publish/{}", uuid::Uuid::new_v4())); + let publish = |e: DemandEvent| { + let c = coordinator.clone(); + async move { c.publish_demand_event("t", &e).await.unwrap() } + }; + assert!(publish(ev("s", 10, None, 1, EventSource::Writer)).await); + assert!( + !publish(ev("s", 10, None, 1, EventSource::Writer)).await, + "duplicate" + ); + assert!( + !publish(ev("s", 9, None, 1, EventSource::Scan)).await, + "lower sealed" + ); + assert!( + !publish(ev("s", 50, None, 0, EventSource::Writer)).await, + "lower epoch" + ); + assert!( + publish(ev("s", 10, Some(4), 1, EventSource::Executor)).await, + "merged advances" + ); + assert!( + !publish(ev("s", 10, Some(3), 1, EventSource::Executor)).await, + "merged regress" + ); + assert!( + publish(ev("s", 12, None, 1, EventSource::Writer)).await, + "sealed advances" + ); + assert!( + publish(ev("s", 1, None, 2, EventSource::Scan)).await, + "higher epoch" + ); + assert!( + publish(ev("other", 1, None, 1, EventSource::Writer)).await, + "separate shard" + ); + } + #[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 35e782e4..6690d919 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -4,6 +4,7 @@ //! Only the executor can publish its terminal outcome. Recovery first cancels //! and reconciles the old execution; it never clears a running fence on timeout. +pub mod demand; pub mod failure; pub mod fencing; pub mod progress; diff --git a/crates/lance-context-server/src/merge_execution.rs b/crates/lance-context-server/src/merge_execution.rs index 3b620e5c..4a70abc0 100644 --- a/crates/lance-context-server/src/merge_execution.rs +++ b/crates/lance-context-server/src/merge_execution.rs @@ -101,6 +101,18 @@ impl Executions { Ok(()) } + pub(crate) async fn publish_demand_event( + &self, + target: &str, + event: &lance_context_merge::demand::DemandEvent, + ) -> Result { + self.coordinator() + .await + .map_err(|e| format!("{e:?}"))? + .publish_demand_event(target, event) + .await + } + pub(crate) async fn request_merge(&self, target: &str) -> Result<(), String> { self.coordinator() .await diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index 79aeeef2..8a2b56ae 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -50,6 +50,15 @@ pub(crate) trait Sweepable: Send + Sync + 'static { async { Ok(false) } } + /// This writer's shard watermark, metadata only. `None` when the store + /// has no shard manifest yet or does not publish demand. + fn own_shard_watermark( + &self, + ) -> impl std::future::Future, String>> + + Send { + async { Ok(None) } + } + /// Fold **every** pending flushed generation into the base table; returns /// how many were reclaimed. fn merge_wal(&self) -> impl std::future::Future> + Send; @@ -74,6 +83,16 @@ impl Sweepable for Arc> { .map_err(|e| e.to_string()) } + async fn own_shard_watermark( + &self, + ) -> Result, String> { + self.read() + .await + .own_shard_watermark() + .await + .map_err(|e| e.to_string()) + } + async fn merge_if_due(&self) -> Result { let prepared = { let guard = self.read().await; @@ -149,6 +168,16 @@ impl Sweepable for Arc> { .map_err(|e| e.to_string()) } + async fn own_shard_watermark( + &self, + ) -> Result, String> { + self.read() + .await + .own_shard_watermark() + .await + .map_err(|e| e.to_string()) + } + async fn merge_if_due(&self) -> Result { // Generic stores default to a deferred seal and take the same // one-row-per-append traffic as rollout, so the count trigger rides @@ -321,6 +350,48 @@ pub(crate) async fn merge_pass_coordinated( } } +async fn publish_demand( + state: &Arc, + target: &str, + store: &S, +) { + let watermark = match store.own_shard_watermark().await { + Ok(Some(watermark)) => watermark, + Ok(None) => return, + Err(error) => { + tracing::debug!(target, %error, "shard watermark unavailable"); + return; + } + }; + let event = lance_context_merge::demand::DemandEvent { + v: lance_context_merge::demand::SCHEMA_VERSION, + shard: watermark.shard_id.to_string(), + sealed_through: watermark.sealed_through, + sealed_bytes_through: 0, + merged_through: None, + flushed_at_ms: chrono::Utc::now().timestamp_millis(), + writer_epoch: watermark.writer_epoch, + source: lance_context_merge::demand::EventSource::Writer, + }; + match state + .merge_executions + .publish_demand_event(target, &event) + .await + { + Ok(written) => { + metrics::counter!( + "rollout_wal_demand_events_total", + "result" => if written { "written" } else { "unchanged" } + ) + .increment(1); + } + Err(error) => { + metrics::counter!("rollout_wal_demand_events_total", "result" => "failed").increment(1); + tracing::debug!(target, %error, "demand event publish failed"); + } + } +} + /// Return true when routing handled the target (including drain/no-op). async fn route_merge( state: &Arc, @@ -336,6 +407,10 @@ async fn route_merge( if state.merge_executions.rollout.draining(&target) { return Ok(true); } + // Demand is published on every flush, below any count threshold, so the + // planner sees low-count tails without a sweep. Failure to publish never + // blocks the merge request path. Additive: `merge-requests` is unchanged. + publish_demand(state, &target, store).await; if !state.merge_executions.owned(&target) { return Ok(false); } @@ -451,6 +526,98 @@ mod tests { assert_eq!(store.read().await.count_base_rows().await.unwrap(), 0); } + /// Demand is published on every flush, below the count threshold and for + /// targets this master does not own, and is monotonic across flushes. + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn flush_publishes_shard_demand_below_threshold_and_without_ownership() { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let client = etcd_client::Client::connect([endpoint], None) + .await + .unwrap(); + let prefix = format!("/demand-sweeper/{}", uuid::Uuid::new_v4()); + let coordinator = lance_context_merge::Coordinator::new(client.clone(), prefix.clone()); + let state_dir = TempDir::new().unwrap(); + let mut state = crate::state::AppState::new_for_test(state_dir.path().to_path_buf()).await; + state.merge_executions = + crate::merge_execution::Executions::new(Some(coordinator.clone()), 3600); + // Deliberately NOT owned: demand must still be visible to a planner. + let state = Arc::new(state); + let dir = TempDir::new().unwrap(); + // Threshold 8, only 2 pending: no merge request is expected. + let store = generic_with_pending(&dir, 8, 2).await; + flush_pass_coordinated( + vec![("s".into(), store.clone())], + Duration::from_secs(5), + None, + Some(state.clone()), + ) + .await; + assert!(coordinator.request_page(None).await.unwrap().0.is_empty()); + let encoded: String = "generic:s".bytes().map(|b| format!("{b:02x}")).collect(); + let events = client + .clone() + .get( + format!("{prefix}/demand-events/{encoded}/"), + Some(etcd_client::GetOptions::new().with_prefix()), + ) + .await + .unwrap(); + assert_eq!(events.kvs().len(), 1, "one event per shard"); + let first: lance_context_merge::demand::DemandEvent = + serde_json::from_slice(events.kvs()[0].value()).unwrap(); + assert_eq!( + first.source, + lance_context_merge::demand::EventSource::Writer + ); + assert!(first.sealed_through >= 2, "{first:?}"); + let mut table = lance_context_merge::demand::TableDemand::default(); + table.fold(&first, 1, 0).unwrap(); + assert_eq!(table.pending_generations(), first.sealed_through); + // A second flush with no new generation changes nothing. + let revision = events.kvs()[0].mod_revision(); + flush_pass_coordinated( + vec![("s".into(), store.clone())], + Duration::from_secs(5), + None, + Some(state.clone()), + ) + .await; + let again = client + .clone() + .get(events.kvs()[0].key(), None) + .await + .unwrap(); + assert_eq!( + again.kvs()[0].mod_revision(), + revision, + "unchanged watermark not rewritten" + ); + // New generation advances it. + store + .read() + .await + .add(&[json!({"id": "r9"}).as_object().unwrap().clone()]) + .await + .unwrap(); + flush_pass_coordinated( + vec![("s".into(), store.clone())], + Duration::from_secs(5), + None, + Some(state), + ) + .await; + let advanced = client + .clone() + .get(events.kvs()[0].key(), None) + .await + .unwrap(); + let second: lance_context_merge::demand::DemandEvent = + serde_json::from_slice(advanced.kvs()[0].value()).unwrap(); + assert!(second.sealed_through > first.sealed_through); + assert_eq!(second.writer_epoch, first.writer_epoch); + } + fn spec() -> SchemaSpec { SchemaSpec::new(vec![( ID_COLUMN.to_string(),