feat: publish merged watermarks in the maintenance release transaction

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<HashMap<Uuid, u64>> {
+    let dataset = Dataset::open(uri).await?;
+    watermarks(&dataset).await
+}
+
 pub(crate) async fn watermarks(dataset: &Dataset) -> Result<HashMap<Uuid, u64>> {
     let indices = dataset.load_indices().await?;
     let Some(index) = indices.iter().find(|i| i.name == MEM_WAL_INDEX_NAME) else {

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
 }
 

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}");
         }
     }
 

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<Compare>, Vec<TxnOp>)> {
+        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<DemandEvent> = 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::<DemandEvent>(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);

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<bool> {
+        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<bool> {
         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<Compare>, operations: Vec<TxnOp>) -> Result<bool> {
+    pub(crate) async fn transact(
+        &self,
+        compares: Vec<Compare>,
+        operations: Vec<TxnOp>,
+    ) -> Result<bool> {
         self.client
             .clone()
             .txn(Txn::new().when(compares).and_then(operations))