Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions crates/lance-context-core/src/generic_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,11 @@ impl GenericStore {
self.base.count_merge_due().await
}

/// Check this shard's pending metadata without enabling local self-merge.
pub async fn count_merge_due_at(&self, threshold: usize) -> LanceResult<bool> {
self.base.count_merge_due_at(threshold).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
Expand Down
5 changes: 5 additions & 0 deletions crates/lance-context-core/src/rollout_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,11 @@ impl RolloutStore {
self.base.count_merge_due().await
}

/// Check this shard's pending metadata without enabling local self-merge.
pub async fn count_merge_due_at(&self, threshold: usize) -> LanceResult<bool> {
self.base.count_merge_due_at(threshold).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(
Expand Down
15 changes: 11 additions & 4 deletions crates/lance-context-core/src/store_base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -867,7 +867,13 @@ impl StorageBase {
/// 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<bool> {
if self.merge_after_generations == 0 || self.is_version_pinned() || self.deleted {
self.count_merge_due_at(self.merge_after_generations).await
}

/// Check this shard's manifest against an external scheduling threshold.
/// This does not enable inline merges or prepare any payload batches.
pub async fn count_merge_due_at(&self, threshold: usize) -> LanceResult<bool> {
if threshold == 0 || self.is_version_pinned() || self.deleted {
return Ok(false);
}
let manifest_store = ShardManifestStore::new(
Expand All @@ -876,9 +882,10 @@ impl StorageBase {
self.write_shard,
DEFAULT_MANIFEST_SCAN_BATCH_SIZE,
);
Ok(manifest_store.read_latest().await?.is_some_and(|manifest| {
manifest.flushed_generations.len() >= self.merge_after_generations
}))
Ok(manifest_store
.read_latest()
.await?
.is_some_and(|manifest| manifest.flushed_generations.len() >= threshold))
}

/// [`Self::prepare_merge_if_ready`], but seals the active memtable *before*
Expand Down
53 changes: 53 additions & 0 deletions crates/lance-context-master/src/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1126,6 +1126,59 @@ mod tests {
assert!(coordinator.request_page(None).await.unwrap().0.is_empty());
}

#[tokio::test]
#[ignore = "requires ETCD_TEST_ENDPOINTS"]
async fn historical_batches_continue_without_new_writes_and_yield_between_passes() {
let dir = TempDir::new().unwrap();
let mut cfg = config(&dir);
cfg.merge_rollout.owned_targets = vec!["hot".into()];
cfg.merge_wal_interval_secs = 0;
let state = MasterState::new(cfg).await.unwrap();
let coordinator = state.task_store.merge_coordinator();
enqueue(&state, TaskKind::MergeWal, "hot").await.unwrap();
// There are no worker threshold requests and no stats sweep. Only
// releasing successful batches can arrange the remaining work.
for reclaimed in [3, 2, 0] {
let claim = state.task_store.claim_next().await.unwrap().unwrap();
assert_eq!(claim.task.target, "hot");
let proof = state.task_store.merge_claim(&claim);
let execution = lance_context_merge::Execution::new("hot", "worker", "boot", 600);
assert!(coordinator.reserve(&proof, &execution).await.unwrap());
let running = coordinator.start(&execution).await.unwrap().unwrap();
assert!(coordinator.finish(&running, Ok(reclaimed)).await.unwrap());
let finished = coordinator.get("hot").await.unwrap().unwrap();
assert!(coordinator.release(&proof, &finished).await.unwrap());
if reclaimed > 0 {
let request = coordinator.request_page(None).await.unwrap().0.remove(0);
enqueue_merge_request(&state, &coordinator, &request)
.await
.unwrap();
assert_eq!(coordinator.request_page(None).await.unwrap().0.len(), 1);
}
// Another table already waiting must run before the continuation.
enqueue(&state, TaskKind::MergeWal, "other").await.unwrap();
state
.task_store
.finish(claim, Ok("batch complete".into()))
.await
.unwrap();
for request in coordinator.request_page(None).await.unwrap().0 {
enqueue_merge_request(&state, &coordinator, &request)
.await
.unwrap();
}
let other = state.task_store.claim_next().await.unwrap().unwrap();
assert_eq!(other.task.target, "other");
state
.task_store
.finish(other, Ok("done".into()))
.await
.unwrap();
}
assert!(state.task_store.claim_next().await.unwrap().is_none());
assert!(coordinator.request_page(None).await.unwrap().0.is_empty());
}

/// Wait until the task reaches a terminal state, returning its final record.
async fn await_terminal(state: &Arc<MasterState>, id: &str) -> TaskRecord {
for _ in 0..100 {
Expand Down
79 changes: 79 additions & 0 deletions crates/lance-context-merge/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,22 @@ impl Coordinator {
Some(etcd_client::PutOptions::new().with_lease(claim.lease_id)),
),
]);
// A byte-bounded batch can succeed while leaving historical WAL behind,
// even when no new writes arrive to trigger demand. Persist continuation
// with release so a master crash cannot lose it. The scheduler coalesces
// it behind the current task, then queues another fair, bounded pass.
// Zero progress, failures and local maintenance must not hot-loop.
if execution.maintenance.is_none()
&& execution.phase == Phase::Finished
&& execution.error.is_none()
&& execution.reclaimed > 0
{
operations.push(TxnOp::put(
self.request_key(&execution.target),
execution.target.clone(),
None,
));
}
self.transact(compares, operations).await
}

Expand Down Expand Up @@ -616,6 +632,69 @@ mod tests {
)
}

#[tokio::test]
#[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
async fn successful_release_durably_requests_another_pass_without_resetting_other_failures() {
let (coordinator, client, proof, _) = fixture().await;
let failure = coordinator
.record_failure(&proof, "table", "broken-worker", "schema mismatch")
.await
.unwrap();
coordinator.request_merge("table").await.unwrap();
let old_request = coordinator.request_page(None).await.unwrap().0.remove(0);
let execution = Execution::new("table", "worker", "boot", 600);
assert!(coordinator.reserve(&proof, &execution).await.unwrap());
let running = coordinator.start(&execution).await.unwrap().unwrap();
assert!(coordinator.finish(&running, Ok(3)).await.unwrap());
let finished = coordinator.get("table").await.unwrap().unwrap();
let stale_claim = ClaimProof {
token: "stale".into(),
..proof.clone()
};
assert!(!coordinator.release(&stale_claim, &finished).await.unwrap());
assert_eq!(
coordinator.request_page(None).await.unwrap().0[0].revision,
old_request.revision
);
assert!(coordinator.release(&proof, &finished).await.unwrap());
// Simulate restarting the master after releasing the execution, before
// it has finished its task. The continuation survives that crash window.
let reconnected = Coordinator::new(client, coordinator.prefix.clone());
assert!(reconnected.get("table").await.unwrap().is_none());
assert!(!reconnected.acknowledge_request(&old_request).await.unwrap());
let requests = reconnected.request_page(None).await.unwrap().0;
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].target, "table");
let after = reconnected
.failure("table", "broken-worker")
.await
.unwrap()
.unwrap();
assert_eq!(after.consecutive_attempts, failure.consecutive_attempts);
assert_eq!(after.next_retry_ms, failure.next_retry_ms);
}

#[tokio::test]
#[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
async fn empty_failed_and_maintenance_releases_do_not_request_more_merges() {
let (coordinator, _, proof, _) = fixture().await;
for (maintenance, outcome) in [
(None, Ok(0)),
(None, Err("storage timeout".into())),
(Some(MaintenanceKind::Compact), Ok(4)),
(Some(MaintenanceKind::IndexId), Ok(4)),
] {
let mut execution = Execution::new("table", "worker", "boot", 600);
execution.maintenance = maintenance;
assert!(coordinator.reserve(&proof, &execution).await.unwrap());
let running = coordinator.start(&execution).await.unwrap().unwrap();
assert!(coordinator.finish(&running, outcome).await.unwrap());
let finished = coordinator.get("table").await.unwrap().unwrap();
assert!(coordinator.release(&proof, &finished).await.unwrap());
assert!(coordinator.request_page(None).await.unwrap().0.is_empty());
}
}

#[tokio::test]
#[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
async fn worker_demand_coalesces_without_resetting_failure_backoff() {
Expand Down
2 changes: 1 addition & 1 deletion crates/lance-context-merge/src/rollout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ pub struct MergeRequest {
}

impl Coordinator {
fn request_key(&self, target: &str) -> String {
pub(crate) fn request_key(&self, target: &str) -> String {
execution_key(&self.prefix, target).replace("/merge-executions/", "/merge-requests/")
}

Expand Down
8 changes: 8 additions & 0 deletions crates/lance-context-server/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,14 @@ pub struct ServerConfig {
#[arg(long, env = "ROLLOUT_MERGE_AFTER_GENERATIONS", default_value = "0")]
pub rollout_merge_after_generations: usize,

/// Per-shard threshold for requesting a master-owned merge on explicitly
/// selected MERGE_OWNED_TARGETS. Uses only manifest metadata on the flush
/// timer and never enables legacy self-merges. 0 disables this trigger.
/// Unset inherits a positive ROLLOUT_MERGE_AFTER_GENERATIONS, or uses 32
/// when legacy count-triggered merging is disabled.
#[arg(long, env = "OWNED_MERGE_AFTER_GENERATIONS")]
pub owned_merge_after_generations: Option<usize>,

/// Maximum flushed MemWAL generations folded into the base table by one
/// merge pass. `0` disables this cap; the byte budget still applies.
/// Leftovers stay pending for the next pass. Both caps apply to count- and
Expand Down
4 changes: 4 additions & 0 deletions crates/lance-context-server/src/merge_execution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ pub struct Executions {
coordinator: OnceCell<Coordinator>,
etcd: Option<lance_context_merge::EtcdConfig>,
pub(crate) rollout: lance_context_merge::rollout::MergeRollout,
pub(crate) merge_after_generations: usize,
queue_timeout_secs: u64,
idle_timeout_secs: u64,
instance: String,
Expand All @@ -62,6 +63,7 @@ impl Executions {
coordinator: OnceCell::new_with(coordinator),
etcd: None,
rollout: Default::default(),
merge_after_generations: 32,
queue_timeout_secs: 600,
idle_timeout_secs: 600,
instance: uuid::Uuid::new_v4().to_string(),
Expand All @@ -75,11 +77,13 @@ impl Executions {
timeout_secs: u64,
queue_timeout_secs: u64,
idle_timeout_secs: u64,
merge_after_generations: usize,
) -> Self {
Self {
coordinator: OnceCell::new(),
etcd: Some(etcd),
rollout,
merge_after_generations,
timeout_secs,
queue_timeout_secs,
idle_timeout_secs,
Expand Down
7 changes: 7 additions & 0 deletions crates/lance-context-server/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,13 @@ impl AppState {
config.merge_execution_timeout_secs,
config.merge_queue_timeout_secs,
config.merge_idle_timeout_secs,
config.owned_merge_after_generations.unwrap_or({
if config.rollout_merge_after_generations > 0 {
config.rollout_merge_after_generations
} else {
32
}
}),
),
stores: RwLock::new(std::collections::HashMap::new()),
rollout_stores: Mutex::new(LruCache::new(capacity)),
Expand Down
26 changes: 18 additions & 8 deletions crates/lance-context-server/src/sweeper.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,10 @@ pub(crate) trait Sweepable: Send + Sync + 'static {
}

/// Cheap metadata-only predicate used when the master owns execution.
fn count_merge_due(&self) -> impl std::future::Future<Output = Result<bool, String>> + Send {
fn count_merge_due(
&self,
_threshold: usize,
) -> impl std::future::Future<Output = Result<bool, String>> + Send {
async { Ok(false) }
}

Expand All @@ -66,10 +69,10 @@ impl Sweepable for Arc<RwLock<RolloutStore>> {
guard.flush().await.map_err(|e| e.to_string())
}

async fn count_merge_due(&self) -> Result<bool, String> {
async fn count_merge_due(&self, threshold: usize) -> Result<bool, String> {
self.read()
.await
.count_merge_due()
.count_merge_due_at(threshold)
.await
.map_err(|e| e.to_string())
}
Expand Down Expand Up @@ -141,10 +144,10 @@ impl Sweepable for Arc<RwLock<GenericStore>> {
guard.flush().await.map_err(|e| e.to_string())
}

async fn count_merge_due(&self) -> Result<bool, String> {
async fn count_merge_due(&self, threshold: usize) -> Result<bool, String> {
self.read()
.await
.count_merge_due()
.count_merge_due_at(threshold)
.await
.map_err(|e| e.to_string())
}
Expand Down Expand Up @@ -339,7 +342,11 @@ async fn route_merge<S: Sweepable>(
if !state.merge_executions.owned(&target) {
return Ok(false);
}
if !count_trigger || store.count_merge_due().await? {
if !count_trigger
|| store
.count_merge_due(state.merge_executions.merge_after_generations)
.await?
{
state.merge_executions.request_merge(&target).await?;
metrics::counter!("rollout_wal_coordinated_merge_requests_total", "kind" => S::kind())
.increment(1);
Expand Down Expand Up @@ -403,7 +410,7 @@ mod tests {

#[tokio::test]
#[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"]
async fn owned_count_trigger_enqueues_without_bypassing_slots_or_preparing_payloads() {
async fn owned_count_trigger_works_with_legacy_trigger_disabled_and_slots_occupied() {
let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").unwrap();
let client = etcd_client::Client::connect([endpoint], None)
.await
Expand All @@ -421,9 +428,12 @@ mod tests {
.rollout
.owned_targets
.push("generic:s".into());
state.merge_executions.merge_after_generations = 2;
let state = Arc::new(state);
let dir = TempDir::new().unwrap();
let store = generic_with_pending(&dir, 2, 2).await;
let store = generic_with_pending(&dir, 0, 2).await;
assert!(!store.count_merge_due(3).await.unwrap());
assert!(!store.count_merge_due(0).await.unwrap());
let slots = Arc::new(Semaphore::new(1));
let _held = slots.clone().acquire_owned().await.unwrap();
flush_pass_coordinated(
Expand Down
34 changes: 34 additions & 0 deletions docs/owned-maintenance-recovery.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,3 +59,37 @@ this change cannot retroactively fence an old unguarded write. It does not enabl
owned targets, alter production pods, or complete registry migration. Recovery
continues on draining targets while an owned execution remains. Validate large
table timings and fault recovery before selecting tighter deadlines.

## Continuous WAL catch-up

A successful owned WAL execution that reclaims generations now writes a coalesced
merge request in the same etcd transaction that releases its execution. The
master's 15-second demand poll retains that request while the current task runs,
then queues another pass. Each pass returns to the normal task queue, preserving
serial worker fan-out and giving other tables and maintenance a scheduling
opportunity. A final zero-progress pass stops the continuation. Failed executions
still use their durable retry budgets; successful work on another shard cannot
reset those budgets. This also drains historical batches when no new writes or
stats sweeps arrive.

Workers use `OWNED_MERGE_AFTER_GENERATIONS` (0 disables) to request work
for explicitly selected owned targets from the existing flush timer. The check
reads only this worker's shard manifest, without loading WAL payloads or taking a
merge slot. Unset, it inherits a positive `ROLLOUT_MERGE_AFTER_GENERATIONS` or
uses 32 when legacy count-triggered merging is disabled. An explicit threshold
is independent of `ROLLOUT_MERGE_AFTER_GENERATIONS`,
which can remain 0 for legacy targets during a per-table rollout. The flush timer
must be enabled. Draining targets request no new work. Count thresholds bound
triggering per shard, not total table pending or merge memory.

Deployment is not activation: all regular and auxiliary masters must support
owned maintenance before selecting a table. Keep legacy writers excluded until
old work is drained or its admitted storage writes have been fenced. Transfer a
dedicated executor's claim only at a verified terminal boundary, then remove the
drain selector and select the table for owned execution on masters and workers.
Never remove a live execution record just because its request timed out.

Validate both backlog-only and sustained-write workloads with small byte-bounded
batches, plus worker/master interruption and repeated shard failures. Check that
ordinary services resume without a helper, another table gets service between
passes, empty tables stop producing continuations, and memory remains bounded.
Loading