From d62e8243e64a931eb8f1869d7f925ef41e64ad2c Mon Sep 17 00:00:00 2001 From: Emil Ernerfeldt Date: Sat, 28 Mar 2026 13:11:15 +0100 Subject: [PATCH] fix(namespace): avoid write starvation under heavy reads `DatasetConsistencyWrapper::get_mut` reloaded the dataset and only then took the write lock, so during the reload it held only a read lock and no write request was queued. tokio's `RwLock` only holds back new readers once a writer is waiting, so every reader that arrived during the reload was admitted ahead of the writer, which then had to wait for all of them before it could take the write lock. `get_mut` now takes the write lock first and reloads while holding it, so write-fairness covers the whole operation. The reload body is factored into `reload_under_write_lock`, shared with the read path's `reload`, which still checks under a read lock and upgrades to a write lock only when a successor version exists. A unit test queues `get_mut` and then a reader behind an exclusively held lock and asserts that the writer is served before the reader. --- .../lance-namespace-impls/src/dir/manifest.rs | 88 +++++++++++++++++-- 1 file changed, 80 insertions(+), 8 deletions(-) diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index 123f7e314fe..6dce81ef50b 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -642,11 +642,13 @@ impl DatasetConsistencyWrapper { /// Get a mutable reference to the dataset. /// Always reloads to ensure strong consistency. + /// + /// Acquires the write lock before reloading so that tokio's write-fairness + /// prevents reader starvation of the writer. pub async fn get_mut(&self) -> Result> { - self.reload().await?; - let guard = DatasetWriteGuard { - guard: self.0.write().await, - }; + let mut write_guard = self.0.write().await; + Self::reload_under_write_lock(&mut write_guard).await?; + let guard = DatasetWriteGuard { guard: write_guard }; ensure_readable(guard.metadata())?; ensure_writable(guard.metadata())?; Ok(guard) @@ -663,7 +665,10 @@ impl DatasetConsistencyWrapper { } } - /// Reload the dataset to the latest version. + /// Reload the dataset to the latest version (for the read path). + /// + /// Takes a read lock first to check if a reload is needed, then upgrades + /// to a write lock only if necessary. async fn reload(&self) -> Result<()> { // First check if we need to reload (with read lock) let read_guard = self.0.read().await; @@ -699,16 +704,29 @@ impl DatasetConsistencyWrapper { // Need to reload, acquire write lock let mut write_guard = self.0.write().await; + Self::reload_under_write_lock(&mut write_guard).await + } - // Double-check after acquiring write lock (someone else might have reloaded) - let has_successor_version = write_guard.has_successor_version().await.map_err(|e| { + /// Reload the dataset while already holding the write lock. + async fn reload_under_write_lock( + dataset: &mut tokio::sync::RwLockWriteGuard<'_, Dataset>, + ) -> Result<()> { + let dataset_uri = dataset.uri().to_string(); + let current_version = dataset.version().version; + log::debug!( + "Reload (under write lock) for uri={}, current_version={}", + dataset_uri, + current_version + ); + + let has_successor_version = dataset.has_successor_version().await.map_err(|e| { lance_core::Error::from(NamespaceError::Internal { message: format!("Failed to check dataset staleness: {:?}", e), }) })?; if has_successor_version { - write_guard.checkout_latest().await.map_err(|e| { + dataset.checkout_latest().await.map_err(|e| { lance_core::Error::from(NamespaceError::Internal { message: format!("Failed to checkout latest: {:?}", e), }) @@ -3868,6 +3886,9 @@ mod tests { use rstest::rstest; use std::collections::{HashMap, HashSet}; use std::sync::Arc; + use std::time::Duration; + use tokio::sync::oneshot; + use tokio::time::timeout; async fn create_manifest_namespace( root: &str, @@ -4466,6 +4487,57 @@ mod tests { assert_eq!(manifest_data_paths(&manifest_ns).await, data_paths_before); } + /// A reader that arrives after a waiting `get_mut()` must be served after + /// it. tokio's `RwLock` guarantees that only once the writer is queued on + /// the write lock: if `get_mut()` first queued as a reader to reload, the + /// later reader would be admitted alongside it, and a reader that holds + /// its permit until the writer is done would then block the write lock + /// forever. + #[tokio::test] + async fn test_get_mut_is_served_before_later_readers() { + let temp_dir = TempStdDir::default(); + let manifest_ns = create_manifest_namespace(temp_dir.to_str().unwrap(), false).await; + let wrapper = manifest_ns.manifest_dataset.clone(); + + // Hold the lock exclusively so both tasks below park on their first + // lock request, in spawn order. The current-thread test runtime polls + // each spawned task before this task resumes from `yield_now`. + let blocker = wrapper.0.write().await; + let writer = { + let wrapper = wrapper.clone(); + tokio::spawn(async move { + wrapper.get_mut().await.unwrap(); + }) + }; + tokio::task::yield_now().await; + + // The late reader takes a single read permit and keeps it until told + // to let go. A raw lock rather than `get()` so that it never releases + // and re-acquires the permit around a reload probe, which would leave + // a window for the writer's request to slip in ahead of it. + let (release_tx, release_rx) = oneshot::channel::<()>(); + let late_reader = { + let wrapper = wrapper.clone(); + tokio::spawn(async move { + let guard = wrapper.0.read().await; + release_rx.await.unwrap(); + drop(guard); + }) + }; + tokio::task::yield_now().await; + + drop(blocker); + timeout(Duration::from_secs(10), writer) + .await + .expect("get_mut() waited on a reader that arrived after it") + .unwrap(); + release_tx.send(()).unwrap(); + timeout(Duration::from_secs(10), late_reader) + .await + .expect("the late reader did not finish after the writer released the lock") + .unwrap(); + } + #[tokio::test] async fn test_manifest_noop_delete_uses_latest_snapshot() { let temp_dir = TempStdDir::default();