fix(master): publish merged watermarks after release, one shard per put

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<Arc<Semaphore>> =
+    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<ShardMerged> = 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);
+    });
+}

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;

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
 }

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::<lance_context_merge::demand::DemandEvent>(
+                            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() {