diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index f64f9810f0..400f7db1fc 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -199,7 +199,6 @@ impl RegionWriteCtx { &self.version } - /// Returns the version control of the region. #[cfg(test)] pub(crate) fn version_control(&self) -> &VersionControlRef { &self.version_control @@ -327,7 +326,7 @@ impl RegionWriteCtx { return; } #[cfg(test)] - test_hooks::pause_before_bulk_install(self.region_id).await; + test_hooks::pause_before_bulk_install(self.region_id, &self.version_control).await; let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["write_bulk"]) .start_timer(); @@ -375,20 +374,10 @@ impl RegionWriteCtx { } } - /// Publishes the sequences and entry id assigned by this context to the - /// region's committed watermark. - /// - /// Must be called after both [`write_memtable`](Self::write_memtable) and - /// [`write_bulk`](Self::write_bulk) have completed, so the committed - /// sequence never covers rows that are not yet physically installed in the - /// memtable (bulk part sequences are assigned in [`push_bulk`](Self::push_bulk) - /// before installation). Since we store the last sequence and entry id in - /// the region, we decrease `next_sequence` and `next_entry_id` by 1. - /// - /// If the write operation failed (e.g. the WAL entry could not be built), - /// no rows were installed and nothing must be published: advancing the - /// committed watermark here would make it cover rows that were never - /// written. + /// Publishes the assigned sequences and entry id to the region's committed + /// watermark. Only call after both [`write_memtable`](Self::write_memtable) + /// and [`write_bulk`](Self::write_bulk) have completed; a failed context + /// must not publish. pub(crate) fn publish_sequence_and_entry_id(&self) { if self.failed { return; @@ -407,25 +396,20 @@ pub(crate) mod test_hooks { use store_api::storage::RegionId; use tokio::sync::watch; - /// Channels of an armed bulk-install barrier. - /// - /// `reached` signals that a bulk write paused between ordinary-memtable - /// handling and bulk installation; `release` unblocks it. Dropping the - /// senders (by disarming the barrier) also unblocks paused writes because - /// their `wait_for` fails on a closed channel. + use crate::region::version::VersionControlRef; + + /// Channels of an armed bulk-install barrier; dropping the senders (by + /// disarming) unblocks writes paused on it. struct ActiveBarrier { id: u64, - /// Only bulk writes for this region pause at the barrier; writes for - /// other regions pass through untouched, so concurrently running tests - /// can never satisfy (or be blocked by) each other's barrier. + /// Only bulk writes for this region on this version control (Arc + /// identity) pause at the barrier. target_region_id: RegionId, + target_version_control: VersionControlRef, reached: watch::Sender, release: watch::Sender, } - /// The currently armed barrier, if any. Wrapped in a `Mutex` so a test can - /// arm a fresh barrier after a previous one was released (or dropped), and - /// so a barrier can never leak past the test that owns it. static ACTIVE_BARRIER: Mutex> = Mutex::new(None); static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); @@ -437,12 +421,8 @@ pub(crate) mod test_hooks { .unwrap_or_else(|poisoned| poisoned.into_inner()) } - /// RAII guard for the bulk-install barrier. - /// - /// The guard owns the armed barrier. Releasing it — or dropping the guard, - /// e.g. when a test panics — unblocks any write paused at the barrier and - /// disarms it, so a paused write can never hang and later tests can arm a - /// fresh barrier. + /// RAII guard: releasing (or dropping) it unblocks a paused write and + /// disarms the barrier. pub(crate) struct BulkInstallBarrier { id: u64, reached_rx: watch::Receiver, @@ -451,20 +431,17 @@ pub(crate) mod test_hooks { } impl BulkInstallBarrier { - /// Waits until a bulk write paused at the barrier. pub(crate) async fn wait_until_reached(&mut self) { if !*self.reached_rx.borrow() { let _ = self.reached_rx.wait_for(|reached| *reached).await; } } - /// Unblocks the paused write and disarms the barrier. pub(crate) fn release(&mut self) { if self.released { return; } self.released = true; - // Flip the release value so writes paused on this barrier proceed. let _ = self.release_tx.send(true); disarm_barrier(self.id); } @@ -476,12 +453,12 @@ pub(crate) mod test_hooks { } } - /// Arms the bulk-install barrier for `target_region_id` and returns a - /// guard that owns it. - /// - /// Any previously armed barrier is replaced: its senders are dropped, which - /// unblocks writes that were paused on it instead of letting them hang. - pub(crate) fn arm_bulk_install_barrier(target_region_id: RegionId) -> BulkInstallBarrier { + /// Arms the bulk-install barrier and returns the owning guard; any + /// previously armed barrier is replaced. + pub(crate) fn arm_bulk_install_barrier( + target_region_id: RegionId, + target_version_control: VersionControlRef, + ) -> BulkInstallBarrier { let (reached_tx, reached_rx) = watch::channel(false); let (release_tx, _release_rx) = watch::channel(false); let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed); @@ -489,6 +466,7 @@ pub(crate) mod test_hooks { *active = Some(ActiveBarrier { id, target_region_id, + target_version_control, reached: reached_tx, release: release_tx.clone(), }); @@ -500,7 +478,6 @@ pub(crate) mod test_hooks { } } - /// Removes the barrier with `id` from the statics if it is still active. fn disarm_barrier(id: u64) { let mut active = lock_active_barrier(); if active.as_ref().is_some_and(|barrier| barrier.id == id) { @@ -508,20 +485,27 @@ pub(crate) mod test_hooks { } } - /// Pauses a bulk write for `region_id` before installing its parts until - /// the test releases the barrier (or it is disarmed). Bulk writes for - /// other regions return immediately. - pub(crate) async fn pause_before_bulk_install(region_id: RegionId) { + /// Pauses a bulk write before installing its parts until the barrier is + /// released or disarmed. + pub(crate) async fn pause_before_bulk_install( + region_id: RegionId, + version_control: &VersionControlRef, + ) { let (reached_tx, release_rx) = { let active = lock_active_barrier(); match active.as_ref() { - Some(barrier) if barrier.target_region_id == region_id => { + Some(barrier) + if barrier.target_region_id == region_id + && std::sync::Arc::ptr_eq( + &barrier.target_version_control, + version_control, + ) => + { (barrier.reached.clone(), barrier.release.subscribe()) } _ => return, } }; - // Signal that a bulk write reached the pause point. let _ = reached_tx.send(true); let mut release_rx = release_rx; if !*release_rx.borrow() { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 0c2eecdea2..9694075624 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -128,9 +128,6 @@ impl RegionWorkerLoop { let mut region_ctx = region_ctxs.into_values().next().unwrap(); region_ctx.write_memtable().await; region_ctx.write_bulk().await; - // Publish only after all rows (including bulk parts) are - // physically installed, so a scan opening on the committed - // sequence can never bind H to invisible rows. region_ctx.publish_sequence_and_entry_id(); put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; @@ -142,8 +139,6 @@ impl RegionWorkerLoop { common_runtime::spawn_global(async move { region_ctx.write_memtable().await; region_ctx.write_bulk().await; - // The spawned task owns the moved ctx, so publish - // inside the task after the memtable writes. region_ctx.publish_sequence_and_entry_id(); (region_ctx.put_num, region_ctx.delete_num) }) @@ -763,10 +758,6 @@ mod tests { use crate::test_util::ts_ms_value; use crate::test_util::version_util::VersionControlBuilder; - /// Creates a bulk part with `num_rows` rows. The schema carries the - /// builder metadata's columns (`tag_0` primary key + `ts` time index) so - /// the bulk-install conversion succeeds; `timestamp_index` points at the - /// `ts` column. fn new_bulk_part(num_rows: i64) -> BulkPart { let schema = Arc::new(Schema::new(vec![ Field::new("tag_0", DataType::Utf8, true), @@ -887,7 +878,6 @@ mod tests { } } - /// Creates a write context for `region_id` with one pending mutation of one row. fn new_region_ctx( region_id: RegionId, ) -> (RegionWriteCtx, oneshot::Receiver>) { @@ -940,16 +930,13 @@ mod tests { region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); - // The failed region must not fail the batch or panic the worker. + // The failed region must not fail the batch. assert!(write_wal(&wal, &mut region_ctxs).await); assert!(region_ctxs[&failing_region].is_failed()); assert!(!region_ctxs[&ok_region].is_failed()); assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); - // Run the memtable and publication phase the worker runs after `write_wal`. - // The failed region installed no rows, so its committed sequence must not - // advance; the successful sibling region advances normally. for region_ctx in region_ctxs.values_mut() { region_ctx.write_memtable().await; region_ctx.write_bulk().await; @@ -969,30 +956,15 @@ mod tests { .committed_sequence() ); - // Waiters of the failed region get the error while others get the result. drop(region_ctxs); assert!(failing_rx.await.unwrap().is_err()); assert_eq!(1, ok_rx.await.unwrap().unwrap()); } - // The committed sequence must not be published before the bulk part's rows - // are physically installed in the memtable: publication happens strictly - // after `write_bulk` returns. Armed with the bulk-install test barrier, - // this deterministically pauses the worker's memtable phase between the - // ordinary-memtable handling and the bulk installation and asserts the - // region's committed sequence is still its initial value (0) — read via - // the `version_control()` accessor, never through a scanner API. After the - // barrier is released, the committed sequence must advance to cover the - // bulk rows. #[tokio::test] async fn test_bulk_write_sequence_not_committed_before_install_worker_level() { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); - assert_eq!( - 0, - version_control.committed_sequence(), - "the builder must start at sequence 0" - ); let mut region_ctxs = HashMap::new(); let mut ctx = RegionWriteCtx::new( @@ -1002,21 +974,18 @@ mod tests { None, ); let (tx, rx) = oneshot::channel(); - // Push a 3-row bulk part without an explicit sequence, exactly like the - // worker's `process_bulk_requests`. assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None)); region_ctxs.insert(region_id, ctx); - // Run the WAL phase the worker runs before the memtable phase. let wal = Wal::new(Arc::new(MockLogStore::default())); assert!(write_wal(&wal, &mut region_ctxs).await); assert!(!region_ctxs[®ion_id].is_failed()); - // Arm the bulk-install barrier: `write_bulk` pauses right before the - // parts are physically installed into the memtable. - let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier(region_id); + let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier( + region_id, + version_control.clone(), + ); - // Run the worker's memtable and publication phases in the background. let write_handle = tokio::spawn(async move { let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); region_ctx.write_memtable().await; @@ -1024,7 +993,6 @@ mod tests { region_ctx.publish_sequence_and_entry_id(); }); - // Wait until the write paused at the barrier: deterministic, no sleeps. tokio::time::timeout( std::time::Duration::from_secs(10), barrier.wait_until_reached(), @@ -1032,16 +1000,12 @@ mod tests { .await .expect("bulk write never reached the install barrier"); - // The committed sequence must not have advanced yet: publication - // happens strictly after the bulk part is installed. assert_eq!( 0, version_control.committed_sequence(), "committed sequence leaked before the bulk part was installed" ); - // Release the barrier: the bulk part installs and the committed - // sequence advances to cover all 3 bulk rows. barrier.release(); write_handle.await.expect("bulk write should complete"); assert_eq!( @@ -1050,7 +1014,6 @@ mod tests { "committed sequence must cover the installed bulk rows" ); - // The bulk waiter is notified with the number of installed rows. assert_eq!(3, rx.await.unwrap().unwrap()); }