From fed5dceffa4b9253ed2e37dfda998754df831954 Mon Sep 17 00:00:00 2001 From: kid Date: Fri, 17 Jul 2026 13:48:46 +0800 Subject: [PATCH] fix(mem_wal): prevent cross-generation durability acknowledgements MemTable batch positions restart at zero after generation rotation, but the durable-write path waited on a writer-global watermark. A successful flush from generation N could therefore acknowledge a durable write at the same local position in generation N+1 before that write's own WAL append completed (#7760). Replace the watermark acknowledgement on the MemTable write path with a request-scoped completion: capture the accepting BatchStore and indexes while holding the writer lock (before a size-triggered freeze can rotate the generation), attach a per-request completion cell to the flush of that exact BatchStore range, and back the returned BatchDurableWatcher with that cell. Typed peer-fence and persistence-failure errors are preserved, and a flush handler that exits without reporting completion surfaces as an I/O error instead of blocking forever. The legacy watermark watcher remains only for the public BatchDurableWatcher::new compatibility path, with a caution on WalFlusher::track_batch. Also return an internal error instead of panicking if the terminal-error mutex is poisoned. Tests cover the coalesced/empty flush completion, flush-completion watcher success/typed-failure/closed-handler, a poisoned terminal-error mutex, and cross-generation aliasing regressions at both the lib level (injected WAL persistence failure) and the integration level (peer epoch claim across a generation rotation). --- rust/lance/src/dataset/mem_wal/wal.rs | 226 ++++++++++++++++++++---- rust/lance/src/dataset/mem_wal/write.rs | 158 ++++++++++++++--- rust/lance/tests/integration_tests.rs | 1 + rust/lance/tests/mem_wal.rs | 62 +++++++ 4 files changed, 387 insertions(+), 60 deletions(-) create mode 100644 rust/lance/tests/mem_wal.rs diff --git a/rust/lance/src/dataset/mem_wal/wal.rs b/rust/lance/src/dataset/mem_wal/wal.rs index b4caff180e1..74a1266037b 100644 --- a/rust/lance/src/dataset/mem_wal/wal.rs +++ b/rust/lance/src/dataset/mem_wal/wal.rs @@ -32,7 +32,8 @@ use uuid::Uuid; use super::manifest::ShardManifestStore; use super::util::{ - WatchableOnceCell, parse_bit_reversed_filename, shard_wal_path, wal_entry_filename, + WatchableOnceCell, WatchableOnceCellReader, parse_bit_reversed_filename, shard_wal_path, + wal_entry_filename, }; use super::index::IndexStore; @@ -87,20 +88,33 @@ impl WalFlushFailure { } } -/// Watcher for batch durability using watermark-based tracking. +#[derive(Clone)] +enum BatchDurabilityState { + /// Legacy writer-global watermark tracking: durable once the shared + /// watermark reaches `target_batch_position`. Used only by publicly + /// constructed watchers for compatibility; it is cross-generation and + /// therefore unsafe for MemTable durable writes (#7760). + Watermark { + rx: watch::Receiver, + target_batch_position: usize, + terminal_error: Arc>>, + }, + /// Request-scoped tracking: durable only when the one flush tied to the + /// exact accepting `BatchStore` generation reports completion. A watermark + /// from a different generation cannot satisfy it. + FlushCompletion { + reader: WatchableOnceCellReader>, + }, +} + +/// Watcher for batch durability. /// -/// Uses a shared watch channel that broadcasts the durable watermark. -/// The watcher waits until the watermark reaches or exceeds its target batch ID. +/// Publicly constructed watchers retain watermark-based tracking for +/// compatibility. MemTable writes use a request-scoped flush completion so a +/// watermark from another generation cannot acknowledge their data. #[derive(Clone)] pub struct BatchDurableWatcher { - /// Watch receiver for the durable watermark. - rx: watch::Receiver, - /// Target batch ID to wait for. - target_batch_position: usize, - /// Terminal flush failure shared with the flusher. When set, the watermark - /// can never reach the target, so `wait` returns this typed error instead of - /// blocking forever. - terminal_error: Arc>>, + state: BatchDurabilityState, } impl BatchDurableWatcher { @@ -111,45 +125,110 @@ impl BatchDurableWatcher { terminal_error: Arc>>, ) -> Self { Self { - rx, - target_batch_position, - terminal_error, + state: BatchDurabilityState::Watermark { + rx, + target_batch_position, + terminal_error, + }, + } + } + + pub(crate) fn from_flush_completion( + reader: WatchableOnceCellReader>, + ) -> Self { + Self { + state: BatchDurabilityState::FlushCompletion { reader }, } } - /// Wait until the batch is durable. + /// Wait until the watched batch is durable. + /// + /// Watermark-backed watchers return success when the durable watermark + /// reaches the target, and return an error for a terminal flush failure or + /// closed tracking channel. Request-scoped watchers return success only + /// after their flush completes, preserve typed flush failures, and return an + /// I/O error if the flush handler exits without reporting completion. /// - /// Returns Ok(()) when `durable_watermark >= target_batch_position`, or - /// Err if a terminal flush failure (e.g. a fence) means the watermark can - /// never reach the target. + /// For a non-blocking check see [`Self::is_durable`]. pub async fn wait(&mut self) -> Result<()> { - loop { - if let Some(failure) = self.terminal_error.lock().unwrap().clone() { - return Err(failure.into_error()); - } - let current = *self.rx.borrow(); - if current >= self.target_batch_position { - return Ok(()); - } - self.rx - .changed() - .await - .map_err(|_| Error::io("Durable watermark channel closed"))?; + match &mut self.state { + BatchDurabilityState::Watermark { + rx, + target_batch_position, + terminal_error, + } => loop { + let terminal_failure = terminal_error + .lock() + .map_err(|_| { + Error::internal( + "Batch durability terminal error mutex was poisoned".to_string(), + ) + })? + .clone(); + if let Some(failure) = terminal_failure { + return Err(failure.into_error()); + } + if *rx.borrow() >= *target_batch_position { + return Ok(()); + } + rx.changed() + .await + .map_err(|_| Error::io("Durable watermark channel closed"))?; + }, + BatchDurabilityState::FlushCompletion { reader } => match reader.await_value().await { + Some(Ok(_)) => Ok(()), + Some(Err(failure)) => Err(failure.into_error()), + None => Err(Error::io( + "WAL flush handler exited before reporting durability", + )), + }, } } /// Check if the batch is already durable (non-blocking). + /// + /// Cheaper, non-awaiting counterpart to [`Self::wait`]: returns the + /// current state without blocking for the flush to complete. pub fn is_durable(&self) -> bool { - *self.rx.borrow() >= self.target_batch_position + match &self.state { + BatchDurabilityState::Watermark { + rx, + target_batch_position, + .. + } => *rx.borrow() >= *target_batch_position, + BatchDurabilityState::FlushCompletion { reader } => { + matches!(reader.read(), Some(Ok(_))) + } + } } } impl std::fmt::Debug for BatchDurableWatcher { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("BatchDurableWatcher") - .field("target_batch_position", &self.target_batch_position) - .field("current_watermark", &*self.rx.borrow()) - .finish() + let mut debug = f.debug_struct("BatchDurableWatcher"); + match &self.state { + BatchDurabilityState::Watermark { + rx, + target_batch_position, + .. + } => { + debug + .field("backend", &"watermark") + .field("target_batch_position", target_batch_position) + .field("current_watermark", &*rx.borrow()); + } + BatchDurabilityState::FlushCompletion { reader } => { + let status = match reader.read() { + None => "pending", + Some(Ok(_)) => "durable", + Some(Err(_)) => "failed", + }; + debug + .field("backend", &"flush_completion") + .field("status", &status); + } + } + debug.finish() } } @@ -420,6 +499,11 @@ impl WalFlusher { /// /// Returns a `BatchDurableWatcher` that can be awaited for durability. /// The actual batch data is stored in the BatchStore. + /// + /// The returned watcher is watermark-based: it resolves against the + /// writer-global watermark, which is not scoped to a MemTable generation. + /// Do not use it for MemTable durable writes (#7760) — the write path uses + /// `BatchDurableWatcher::from_flush_completion` instead. pub fn track_batch(&self, batch_position: usize) -> BatchDurableWatcher { // Return a watcher that waits for this batch to become durable // batch_position is 0-indexed, so we wait for watermark > batch_position (i.e., >= batch_position + 1) @@ -1582,6 +1666,76 @@ mod tests { assert!(!watcher.is_durable()); } + #[tokio::test] + async fn test_flush_completion_watcher_reports_success() { + let completion = WatchableOnceCell::new(); + let mut watcher = BatchDurableWatcher::from_flush_completion(completion.reader()); + + assert!(!watcher.is_durable()); + completion.write(Ok(empty_flush_result())); + + watcher.wait().await.unwrap(); + assert!(watcher.is_durable()); + } + + #[tokio::test] + async fn test_flush_completion_watcher_preserves_typed_failure() { + let completion = WatchableOnceCell::new(); + let mut watcher = BatchDurableWatcher::from_flush_completion(completion.reader()); + completion.write(Err(WalFlushFailure { + fence_reason: Some(FenceReason::PersistenceFailure), + message: "injected persistence failure".to_string(), + })); + + let error = watcher.wait().await.unwrap_err(); + + assert_eq!(error.fence_reason(), Some(FenceReason::PersistenceFailure)); + assert!(error.to_string().contains("injected persistence failure")); + assert!(!watcher.is_durable()); + } + + #[tokio::test] + async fn test_flush_completion_watcher_reports_closed_handler() { + let completion = + WatchableOnceCell::>::new(); + let mut watcher = BatchDurableWatcher::from_flush_completion(completion.reader()); + drop(completion); + + let error = watcher.wait().await.unwrap_err(); + + assert!(matches!(error, Error::IO { .. })); + assert!( + error + .to_string() + .contains("WAL flush handler exited before reporting durability") + ); + assert!(!watcher.is_durable()); + } + + #[tokio::test] + async fn test_watermark_watcher_reports_poisoned_terminal_error() { + let terminal_error = Arc::new(StdMutex::new(None)); + let terminal_error_to_poison = terminal_error.clone(); + let poison_result = std::thread::spawn(move || { + let _guard = terminal_error_to_poison.lock().unwrap(); + panic!("poison terminal error mutex"); + }) + .join(); + assert!(poison_result.is_err()); + + let (_tx, rx) = watch::channel(0); + let mut watcher = BatchDurableWatcher::new(rx, 1, terminal_error); + + let error = watcher.wait().await.unwrap_err(); + + assert!(matches!(error, Error::Internal { .. })); + assert!( + error + .to_string() + .contains("Batch durability terminal error mutex was poisoned") + ); + } + // Regression test: track_batch must return a watcher wired to the real // WAL watermark, NOT a pre-resolved watcher. A pre-resolved watcher would // cause durable writes to return before the WAL is actually flushed. diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 3cb4586ebb4..d7e000c7bf2 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -1145,11 +1145,6 @@ impl SharedWriterState { Ok(next_generation) } - /// Track batch for WAL durability. - fn track_batch_for_wal(&self, batch_position: usize) -> super::wal::BatchDurableWatcher { - self.wal_flusher.track_batch(batch_position) - } - /// Check if memtable flush is needed and trigger if so. /// /// Takes `&mut WriterState` directly since caller already holds the lock. @@ -1859,49 +1854,46 @@ impl ShardWriter { let start = std::time::Instant::now(); // Acquire write lock for entire operation (atomic approach) - let (batch_positions, durable_watcher, batch_store, indexes) = { + let (batch_positions, batch_store, indexes) = { let mut state = state_lock.write().await; - // 1. Insert all batches into memtable atomically let results = state.memtable.insert_batches_only(batches).await?; - - // Get batch position range let start_pos = results.first().map(|(pos, _, _)| *pos).unwrap_or(0); let end_pos = results.last().map(|(pos, _, _)| pos + 1).unwrap_or(0); let batch_positions = start_pos..end_pos; - // 2. Track last batch for WAL durability - let durable_watcher = writer_state.track_batch_for_wal(end_pos.saturating_sub(1)); + // INVARIANT: capture the accepting BatchStore/indexes BEFORE calling + // maybe_trigger_* below. Those may freeze the active MemTable and + // rotate to a new generation, so any capture after them would bind + // the request-scoped flush completion to the wrong (new) generation. + // Moving these lines down reintroduces cross-generation aliasing (#7760). + let batch_store = state.memtable.batch_store(); + let indexes = state.memtable.indexes_arc(); - // 3. Check if WAL flush should be triggered writer_state.maybe_trigger_wal_flush(&mut state); - - // 4. Check if memtable flush is needed - if let Err(e) = writer_state.maybe_trigger_memtable_flush(&mut state) { - warn!("Failed to trigger memtable flush: {}", e); + if let Err(error) = writer_state.maybe_trigger_memtable_flush(&mut state) { + warn!("Failed to trigger memtable flush: {}", error); } - // Get batch_store and indexes while we have the lock (for durable_write case) - let batch_store = state.memtable.batch_store(); - let indexes = state.memtable.indexes_arc(); - - (batch_positions, durable_watcher, batch_store, indexes) - }; // Lock released here + (batch_positions, batch_store, indexes) + }; self.stats.record_put(start.elapsed()); // Trigger the flush here (outside the lock) so the watcher can resolve; // only the `wait()` is the caller's to schedule. let watcher = if self.config.durable_write { + let completion = WatchableOnceCell::new(); + let watcher = BatchDurableWatcher::from_flush_completion(completion.reader()); self.wal_flusher.trigger_flush( WalFlushSource::BatchStore { batch_store, indexes, }, batch_positions.end, - None, + Some(completion), )?; - Some(durable_watcher) + Some(watcher) } else { None }; @@ -3811,6 +3803,80 @@ mod tests { .unwrap() } + #[tokio::test] + async fn test_coalesced_empty_flush_completes_request_watcher() { + let (store, base_path, _base_uri, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let wal_appender = Arc::new( + WalAppender::open(store, base_path, shard_id, 0) + .await + .unwrap(), + ); + let wal_flusher = Arc::new(WalFlusher::new(wal_appender)); + + let schema = create_test_schema(); + let mut memtable = MemTable::new(schema.clone(), 1, vec![]).unwrap(); + memtable + .insert_batches_only(vec![create_test_batch(&schema, 0, 1)]) + .await + .unwrap(); + let batch_store = memtable.batch_store(); + let end_batch_position = batch_store.len(); + let state = Arc::new(RwLock::new(WriterState { + memtable, + last_flushed_wal_entry_position: 0, + frozen_memtable_bytes: 0, + frozen_flush_watchers: VecDeque::new(), + frozen_memtables: VecDeque::new(), + flush_requested: false, + wal_flush_trigger_count: 0, + last_wal_flush_trigger_time: 0, + })); + let mut handler = WalFlushHandler::new(wal_flusher, Some(state), new_shared_stats()); + + handler + .handle(TriggerWalFlush { + source: WalFlushSource::BatchStore { + batch_store: batch_store.clone(), + indexes: None, + }, + end_batch_position, + done: None, + }) + .await + .unwrap(); + assert_eq!( + batch_store.max_flushed_batch_position(), + Some(end_batch_position - 1), + "first flush must cover the request range" + ); + + let completion = WatchableOnceCell::new(); + let completion_reader = completion.reader(); + let mut watcher = BatchDurableWatcher::from_flush_completion(completion.reader()); + handler + .handle(TriggerWalFlush { + source: WalFlushSource::BatchStore { + batch_store, + indexes: None, + }, + end_batch_position, + done: Some(completion), + }) + .await + .unwrap(); + + let flush_result = completion_reader + .read() + .expect("coalesced flush must report request completion") + .expect("coalesced flush must complete successfully"); + assert!( + flush_result.entry.is_none(), + "already-covered flush must not append another WAL entry" + ); + watcher.wait().await.unwrap(); + } + #[tokio::test] async fn test_shard_writer_basic_write() { let (store, base_path, base_uri, _temp_dir) = create_local_store().await; @@ -4971,6 +5037,50 @@ mod tests { .unwrap(); } + /// Regression for #7760 (cross-generation aliasing): gen-1's flush + /// publishes the writer-global watermark, which must NOT satisfy the + /// gen-2 durable put. With the bug, this `put` returned `Ok` because the + /// gen-1 watermark already covered gen-2's local position 0; here gen-2's + /// own WAL append is forced to fail, so a correctly request-scoped watcher + /// must surface `PersistenceFailure` rather than return success. + #[tokio::test] + async fn test_durable_watcher_does_not_alias_across_memtable_generations() { + let (store, base_path, controls) = failing_memory_store().await; + let base_uri = "memory:///"; + let schema = create_test_schema(); + let config = ShardWriterConfig { + shard_id: Uuid::new_v4(), + durable_write: true, + sync_indexed_write: false, + max_wal_buffer_size: 64 * 1024 * 1024, + max_wal_flush_interval: None, + max_wal_persist_retries: 0, + wal_persist_retry_base_delay: Duration::from_millis(1), + max_memtable_size: 64 * 1024 * 1024, + manifest_scan_batch_size: 2, + ..Default::default() + }; + let writer = ShardWriter::open(store, base_path, base_uri, config, schema.clone(), vec![]) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 1)]) + .await + .unwrap(); + writer.force_seal_active().await.unwrap(); + writer.wait_for_flush_drain().await.unwrap(); + + controls.fail_wal_puts(usize::MAX); + let error = writer + .put(vec![create_test_batch(&schema, 1, 1)]) + .await + .expect_err("new generation must wait for its failing WAL flush"); + + assert_eq!(error.fence_reason(), Some(FenceReason::PersistenceFailure)); + assert!(error.to_string().contains("WAL persistence failed")); + } + /// Replay-on-open recovers durable WAL entries that were never flushed /// to a Lance generation. Setup: writer A durably writes batches, drops /// without close (so MemTable freeze never runs); writer B reopens and diff --git a/rust/lance/tests/integration_tests.rs b/rust/lance/tests/integration_tests.rs index 7a6d3e71ca4..e6ea4e321ae 100644 --- a/rust/lance/tests/integration_tests.rs +++ b/rust/lance/tests/integration_tests.rs @@ -4,6 +4,7 @@ // NOTE: we only create one integration test binary, to keep compilation overhead down. mod count_pushdown; +mod mem_wal; #[cfg(feature = "slow_tests")] mod query; #[cfg(feature = "slow_tests")] diff --git a/rust/lance/tests/mem_wal.rs b/rust/lance/tests/mem_wal.rs new file mode 100644 index 00000000000..856b91fea8c --- /dev/null +++ b/rust/lance/tests/mem_wal.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_array::record_batch; +use lance::dataset::mem_wal::{ShardWriter, ShardWriterConfig}; +use lance_core::FenceReason; +use lance_io::object_store::ObjectStore; +use uuid::Uuid; + +#[tokio::test] +async fn durable_put_does_not_alias_across_memtable_generations() { + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("file://{}", temp_dir.path().display()); + let (store, base_path) = ObjectStore::from_uri(&base_uri).await.unwrap(); + let shard_id = Uuid::new_v4(); + let first_batch = record_batch!(("id", Int32, [1])).unwrap(); + let schema = first_batch.schema(); + let config = ShardWriterConfig { + shard_id, + durable_write: true, + max_wal_buffer_size: 64 * 1024 * 1024, + max_wal_flush_interval: None, + max_memtable_size: 64 * 1024 * 1024, + manifest_scan_batch_size: 2, + ..Default::default() + }; + + let writer_a = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + config.clone(), + schema.clone(), + vec![], + ) + .await + .unwrap(); + + writer_a.put(vec![first_batch]).await.unwrap(); + let first_generation = writer_a.memtable_stats().await.unwrap().generation; + writer_a.force_seal_active().await.unwrap(); + writer_a.wait_for_flush_drain().await.unwrap(); + assert_eq!( + writer_a.memtable_stats().await.unwrap().generation, + first_generation + 1 + ); + + let writer_b = ShardWriter::open(store, base_path, base_uri, config, schema, vec![]) + .await + .unwrap(); + assert!(writer_b.epoch() > writer_a.epoch()); + + let second_batch = record_batch!(("id", Int32, [2])).unwrap(); + let error = writer_a + .put(vec![second_batch]) + .await + .expect_err("generation-2 durable put must wait for its own WAL flush"); + assert_eq!(error.fence_reason(), Some(FenceReason::PeerClaimedEpoch)); + assert!(error.to_string().contains("Writer fenced")); + + writer_b.close().await.unwrap(); +}