This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3435-6f6d68bb29d00cd34a93a4d598f882b97020f676 in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit 59946648ef56dcd5b70d46cc7ef90b25a00ed9fb Author: JeonDaehong <[email protected]> AuthorDate: Fri Nov 6 07:45:40 2026 +0000 fix(arrow): release a positional delete file whose load fails (#3435) * fix(arrow): release a positional delete file whose load fails try_start_pos_del_load marks a positional delete file Loading, and only the success path marked it Loaded. If opening or parsing the file failed, or the load was dropped because another delete file failed, the entry stayed Loading: a retry on the same loader, a clone of the same ArrowReader, or a concurrent load of the same file waited for a notification that never came. The Load action now carries a PosDelLoadGuard. finish() marks the file loaded; dropping the guard without finishing removes the Loading marker and wakes the waiters. A woken waiter checks the state again and loads the file itself if it isn't loaded, instead of assuming it is. Closes #3434. * refactor(arrow): assert the Loading invariant when abandoning a load Only try_start_pos_del_load inserts into positional_deletes, and only the guard that it returns changes the entry afterwards, so the entry is always the guard's Loading marker when it is abandoned. * test(arrow): cover waiters that load after an abandoned positional load Hold the Load guard, poll load_file_for_task for the same file until it waits, then drop the guard: the waiter loads the file itself, or returns the open error when the file is gone. Also load a positional delete file that opens but fails to parse, then rewrite it and load it again on the same loader, which hangs without the guard. * fix(arrow): return a failed positional delete load's error to its waiters Waiters on a failed load used to retry it one after another, so a missing file shared by many tasks was opened once per task. The load now records its error and the tasks waiting on that attempt return it. A load dropped without an error, when a sibling's error stops load_deletes, still lets its waiters retry, and later calls start a new attempt either way. Abandoning a load also tolerates a poisoned lock, since it runs from Drop. * test(arrow): keep the sibling's load in flight when another delete file fails On the current-thread test runtime the local filesystem never yields, so the valid file was only loaded on the second call and its guard was never dropped mid-load. A test storage now yields before reads and holds one file until the test opens a gate. * fix(arrow): pass the whole load error to waiters --- .../src/arrow/caching_delete_file_loader.rs | 446 +++++++++++++++++++-- crates/iceberg/src/arrow/delete_filter.rs | 208 ++++++++-- 2 files changed, 598 insertions(+), 56 deletions(-) diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 5582e4397..0711a58b3 100644 --- a/crates/iceberg/src/arrow/caching_delete_file_loader.rs +++ b/crates/iceberg/src/arrow/caching_delete_file_loader.rs @@ -23,7 +23,7 @@ use bytes::Bytes; use futures::{StreamExt, TryStreamExt}; use tokio::sync::oneshot::{Receiver, channel}; -use super::delete_filter::{DeleteFilter, PosDelLoadAction}; +use super::delete_filter::{DeleteFilter, PosDelLoadAction, PosDelLoadGuard}; use crate::arrow::delete_file_loader::BasicDeleteFileLoader; use crate::arrow::scan_metrics::ScanMetrics; use crate::arrow::{arrow_primitive_to_literal, arrow_schema_to_schema}; @@ -57,7 +57,7 @@ enum DeleteFileContext { ExistingEqDel, ExistingPosDel, PosDels { - file_path: String, + guard: PosDelLoadGuard, stream: ArrowRecordBatchStream, }, FreshEqDel { @@ -80,7 +80,7 @@ enum DeleteFileContext { // results are fully merged into the DeleteFileManager's state enum ParsedDeleteFileContext { DelVecs { - file_path: String, + guard: PosDelLoadGuard, results: HashMap<String, DeleteVector>, }, // A single deletion vector decoded from a Puffin blob, keyed by the data file it applies to. @@ -229,12 +229,12 @@ impl CachingDeleteFileLoader { while let Some(item) = results_stream.next().await { match item? { - ParsedDeleteFileContext::DelVecs { file_path, results } => { + ParsedDeleteFileContext::DelVecs { guard, results } => { for (data_file_path, delete_vector) in results.into_iter() { del_filter.upsert_delete_vector(data_file_path, delete_vector); } // Mark the positional delete file as fully loaded so waiters can proceed - del_filter.finish_pos_del_load(&file_path); + guard.finish(); } ParsedDeleteFileContext::DelVec { data_file_path, @@ -271,25 +271,37 @@ impl CachingDeleteFileLoader { return Self::load_deletion_vector(task, basic_delete_file_loader).await; } - match del_filter.try_start_pos_del_load(task.file_path()) { - PosDelLoadAction::AlreadyLoaded => Ok(DeleteFileContext::ExistingPosDel), - PosDelLoadAction::WaitFor(notified) => { - // Positional deletes are accessed synchronously by ArrowReader. - // We must wait here to ensure the data is ready before returning, - // otherwise ArrowReader might get an empty/partial result. - notified.await; - Ok(DeleteFileContext::ExistingPosDel) + loop { + match del_filter.try_start_pos_del_load(task.file_path()) { + PosDelLoadAction::AlreadyLoaded => { + return Ok(DeleteFileContext::ExistingPosDel); + } + PosDelLoadAction::WaitFor(wait) => { + // Positional deletes are accessed synchronously by ArrowReader. + // We must wait here to ensure the data is ready before returning, + // otherwise ArrowReader might get an empty/partial result. + // A failed load returns its error here. A load dropped without + // one (cancelled) ends the wait too, so check again. + wait.wait().await?; + } + PosDelLoadAction::Load(guard) => { + let stream = match basic_delete_file_loader + .parquet_to_batch_stream( + task.file_path(), + task.file_size_in_bytes(), + task.key_metadata(), + ) + .await + { + Ok(stream) => stream, + Err(err) => { + guard.fail(&err); + return Err(err); + } + }; + return Ok(DeleteFileContext::PosDels { guard, stream }); + } } - PosDelLoadAction::Load => Ok(DeleteFileContext::PosDels { - file_path: task.file_path().to_string(), - stream: basic_delete_file_loader - .parquet_to_batch_stream( - task.file_path(), - task.file_size_in_bytes(), - task.key_metadata(), - ) - .await?, - }), } } @@ -435,12 +447,17 @@ impl CachingDeleteFileLoader { match ctx { DeleteFileContext::ExistingEqDel => Ok(ParsedDeleteFileContext::EqDel), DeleteFileContext::ExistingPosDel => Ok(ParsedDeleteFileContext::ExistingPosDel), - DeleteFileContext::PosDels { file_path, stream } => { - let del_vecs = Self::parse_positional_deletes_record_batch_stream(stream).await?; - Ok(ParsedDeleteFileContext::DelVecs { - file_path, - results: del_vecs, - }) + DeleteFileContext::PosDels { guard, stream } => { + match Self::parse_positional_deletes_record_batch_stream(stream).await { + Ok(del_vecs) => Ok(ParsedDeleteFileContext::DelVecs { + guard, + results: del_vecs, + }), + Err(err) => { + guard.fail(&err); + Err(err) + } + } } DeleteFileContext::DelVec { data_file_path, @@ -1791,4 +1808,375 @@ mod tests { assert_eq!(err.kind(), ErrorKind::DataInvalid); assert!(err.message().contains("expected 2 from record_count")); } + + #[tokio::test] + async fn test_pos_del_load_can_be_retried_after_failure() { + let tmp_dir = TempDir::new().unwrap(); + let file_scan_tasks = setup(tmp_dir.path()); + let task = &file_scan_tasks[0]; + let pos_del = task + .deletes() + .iter() + .find(|delete| delete.file_type() == DataContentType::PositionDeletes) + .unwrap() + .clone(); + let path = pos_del.file_path().to_string(); + let hidden = format!("{path}.hidden"); + let load = |loader: &CachingDeleteFileLoader| { + loader.load_deletes(std::slice::from_ref(&pos_del), task.schema_ref()) + }; + + // The first load fails because the file is not there yet. + std::fs::rename(&path, &hidden).unwrap(); + let loader = CachingDeleteFileLoader::new(FileIO::new_with_fs(), 10, Runtime::current()); + assert!(load(&loader).await.unwrap().is_err()); + + // Once it is, the same loader loads it instead of waiting on the failed load. + std::fs::rename(&hidden, &path).unwrap(); + let retried = tokio::time::timeout(std::time::Duration::from_secs(5), load(&loader)) + .await + .expect("retry waited on the failed load") + .unwrap() + .unwrap(); + let fresh = load(&CachingDeleteFileLoader::new( + FileIO::new_with_fs(), + 10, + Runtime::current(), + )) + .await + .unwrap() + .unwrap(); + + let positions = |filter: &DeleteFilter| -> Vec<u64> { + filter + .get_delete_vector(task) + .unwrap() + .lock() + .unwrap() + .iter() + .collect() + }; + assert!(!positions(&fresh).is_empty()); + assert_eq!(positions(&retried), positions(&fresh)); + } + + /// The load of a positional delete file taken by the test, and tasks polled until + /// they wait on it. + struct PendingPosDelLoad { + guard: PosDelLoadGuard, + waiters: Vec<futures::future::BoxFuture<'static, Result<DeleteFileContext>>>, + path: String, + } + + /// Takes the load of the first test task's positional delete file, then polls + /// `load_file_for_task` for the same file from `waiters` tasks until each waits on it. + async fn wait_on_pos_del_load(tmp_dir: &TempDir, waiters: usize) -> PendingPosDelLoad { + let task = setup(tmp_dir.path()).remove(0); + let pos_del = task + .deletes() + .iter() + .find(|delete| delete.file_type() == DataContentType::PositionDeletes) + .unwrap() + .clone(); + let loader = CachingDeleteFileLoader::new(FileIO::new_with_fs(), 10, Runtime::current()); + + let PosDelLoadAction::Load(guard) = loader + .delete_filter + .try_start_pos_del_load(pos_del.file_path()) + else { + panic!("expected Load for a new file"); + }; + let mut pending = Vec::with_capacity(waiters); + for _ in 0..waiters { + let (pos_del, schema) = (pos_del.clone(), task.schema_ref()); + let (basic, filter) = ( + loader.basic_delete_file_loader.clone(), + loader.delete_filter.clone(), + ); + let mut waiter: futures::future::BoxFuture<'static, _> = Box::pin(async move { + CachingDeleteFileLoader::load_file_for_task(&pos_del, basic, filter, schema).await + }); + assert!(futures::poll!(&mut waiter).is_pending()); + pending.push(waiter); + } + PendingPosDelLoad { + guard, + waiters: pending, + path: pos_del.file_path().to_string(), + } + } + + #[tokio::test] + async fn test_pos_del_waiter_loads_file_after_cancelled_load() { + let tmp_dir = TempDir::new().unwrap(); + let PendingPosDelLoad { + guard, mut waiters, .. + } = wait_on_pos_del_load(&tmp_dir, 1).await; + + // Dropped without an error, as when a sibling's error stops `load_deletes`. + drop(guard); + + let ctx = tokio::time::timeout(std::time::Duration::from_secs(5), waiters.remove(0)) + .await + .expect("waiter did not wake after the load was cancelled") + .unwrap(); + assert!(matches!(ctx, DeleteFileContext::PosDels { .. })); + } + + #[tokio::test] + async fn test_pos_del_waiters_return_the_error_of_a_failed_load() { + let tmp_dir = TempDir::new().unwrap(); + let PendingPosDelLoad { guard, waiters, .. } = wait_on_pos_del_load(&tmp_dir, 2).await; + + // The file is still there, so a waiter that opened it again would succeed. + guard.fail(&Error::new(ErrorKind::Unexpected, "load failed")); + + for waiter in waiters { + let err = tokio::time::timeout(std::time::Duration::from_secs(5), waiter) + .await + .expect("waiter did not wake after the load failed") + .err() + .unwrap(); + assert_eq!(err.kind(), ErrorKind::Unexpected); + assert!(err.message().contains("load failed"), "{err}"); + } + } + + #[tokio::test] + async fn test_pos_del_waiters_retry_after_cancelled_load() { + let tmp_dir = TempDir::new().unwrap(); + let PendingPosDelLoad { + guard, + waiters, + path, + } = wait_on_pos_del_load(&tmp_dir, 2).await; + + // Each waiter loads the file itself after a cancelled load, and gets its own error. + std::fs::remove_file(&path).unwrap(); + drop(guard); + + for waiter in waiters { + let err = tokio::time::timeout(std::time::Duration::from_secs(5), waiter) + .await + .expect("waiter did not wake after the load was cancelled") + .err() + .unwrap(); + assert!(err.to_string().contains("Failed to open file"), "{err}"); + } + } + + /// Local filesystem storage that yields before every read, so the loads of several + /// delete files in one `load_deletes` call start together, and holds reads of + /// `gated.parquet` until the gate is opened, so one of them is still in flight when + /// another fails. `LocalFsStorage` never yields, so on the current-thread test + /// runtime each load would otherwise finish before the next one starts. + mod gated_storage { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + use bytes::Bytes; + use futures::stream::BoxStream; + + use crate::Result; + use crate::io::{ + FileMetadata, FileRead, FileWrite, InputFile, LocalFsStorage, OutputFile, Storage, + StorageConfig, StorageFactory, + }; + + #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] + pub(super) struct GatedStorage { + #[serde(skip)] + inner: LocalFsStorage, + #[serde(skip)] + open: Arc<AtomicBool>, + } + + impl GatedStorage { + async fn wait_for_gate(&self, path: &str) { + tokio::task::yield_now().await; + while path.ends_with("gated.parquet") && !self.open.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + } + } + + #[async_trait::async_trait] + #[typetag::serde(name = "test_gated_storage")] + impl Storage for GatedStorage { + async fn metadata(&self, path: &str) -> Result<FileMetadata> { + self.inner.metadata(path).await + } + + fn new_input(&self, path: &str) -> Result<InputFile> { + Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) + } + + fn new_output(&self, path: &str) -> Result<OutputFile> { + Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) + } + + async fn exists(&self, path: &str) -> Result<bool> { + self.inner.exists(path).await + } + + async fn read(&self, path: &str) -> Result<Bytes> { + self.wait_for_gate(path).await; + self.inner.read(path).await + } + + async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> { + self.wait_for_gate(path).await; + self.inner.reader(path).await + } + + async fn write(&self, path: &str, bs: Bytes) -> Result<()> { + self.inner.write(path, bs).await + } + + async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> { + self.inner.writer(path).await + } + + async fn delete(&self, path: &str) -> Result<()> { + self.inner.delete(path).await + } + + async fn delete_prefix(&self, path: &str) -> Result<()> { + self.inner.delete_prefix(path).await + } + + async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> { + self.inner.delete_stream(paths).await + } + } + + #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] + pub(super) struct GatedStorageFactory { + #[serde(skip)] + pub(super) open: Arc<AtomicBool>, + } + + #[typetag::serde(name = "test_gated_storage_factory")] + impl StorageFactory for GatedStorageFactory { + fn build(&self, _config: &StorageConfig) -> Result<Arc<dyn Storage>> { + Ok(Arc::new(GatedStorage { + inner: LocalFsStorage, + open: Arc::clone(&self.open), + })) + } + } + } + + /// Writes a position delete file at `path` for `data.parquet` with `positions`. + fn write_pos_del_file(path: &str, positions: Vec<i64>) -> FileScanTaskDeleteFile { + let batch = RecordBatch::try_new( + crate::arrow::delete_filter::tests::create_pos_del_schema(), + vec![ + Arc::new(StringArray::from_iter_values(vec![ + "data.parquet"; + positions.len() + ])), + Arc::new(Int64Array::from_iter_values(positions)), + ], + ) + .unwrap(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + FileScanTaskDeleteFile::builder() + .with_file_path(path.to_string()) + .with_file_size_in_bytes(std::fs::metadata(path).unwrap().len()) + .with_file_type(DataContentType::PositionDeletes) + .with_file_format(DataFileFormat::Parquet) + .with_partition_spec_id(0) + .build() + .unwrap() + } + + /// A failing positional delete file ends `load_deletes` early while the load of the + /// other file in the same call is still in flight, dropping its guard. That file + /// still loads later. + #[tokio::test] + async fn test_pos_del_load_dropped_by_a_failing_sibling_can_be_loaded_later() { + let tmp_dir = TempDir::new().unwrap(); + let schema = Arc::new(Schema::builder().build().unwrap()); + let invalid = write_pos_del_file( + tmp_dir.path().join("invalid.parquet").to_str().unwrap(), + vec![-1], + ); + // Its read waits for the gate, so its load is in flight when `invalid` fails. + let valid = write_pos_del_file( + tmp_dir.path().join("gated.parquet").to_str().unwrap(), + vec![1, 3], + ); + let factory = Arc::new(gated_storage::GatedStorageFactory::default()); + let gate = Arc::clone(&factory.open); + let file_io = crate::io::FileIOBuilder::new(factory).build(); + let loader = CachingDeleteFileLoader::new(file_io, 10, Runtime::current()); + + let err = loader + .load_deletes(&[invalid, valid.clone()], Arc::clone(&schema)) + .await + .unwrap() + .err() + .unwrap(); + assert!(err.message().contains("negative position"), "{err}"); + + gate.store(true, std::sync::atomic::Ordering::SeqCst); + let filter = tokio::time::timeout( + std::time::Duration::from_secs(5), + loader.load_deletes(&[valid], schema), + ) + .await + .expect("load waited on the dropped load") + .unwrap() + .unwrap(); + let positions: Vec<u64> = filter + .get_delete_vector_for_path("data.parquet") + .unwrap() + .lock() + .unwrap() + .iter() + .collect(); + assert_eq!(positions, vec![1, 3]); + } + + #[tokio::test] + async fn test_pos_del_load_can_be_retried_after_parse_failure() { + let tmp_dir = TempDir::new().unwrap(); + let path = tmp_dir.path().join("pos-del.parquet"); + let path = path.to_str().unwrap(); + let schema = Arc::new(Schema::builder().build().unwrap()); + let loader = CachingDeleteFileLoader::new(FileIO::new_with_fs(), 10, Runtime::current()); + + // The file opens but fails to parse, so the load is abandoned after the open. + let invalid = write_pos_del_file(path, vec![-1]); + let err = loader + .load_deletes(&[invalid], Arc::clone(&schema)) + .await + .unwrap() + .err() + .unwrap(); + assert!(err.message().contains("negative position"), "{err}"); + + // Once the file is valid, the same loader loads it instead of waiting. + let valid = write_pos_del_file(path, vec![1, 3]); + let filter = tokio::time::timeout( + std::time::Duration::from_secs(5), + loader.load_deletes(&[valid], schema), + ) + .await + .expect("retry waited on the failed load") + .unwrap() + .unwrap(); + let positions: Vec<u64> = filter + .get_delete_vector_for_path("data.parquet") + .unwrap() + .lock() + .unwrap() + .iter() + .collect(); + assert_eq!(positions, vec![1, 3]); + } } diff --git a/crates/iceberg/src/arrow/delete_filter.rs b/crates/iceberg/src/arrow/delete_filter.rs index 29843295d..977917030 100644 --- a/crates/iceberg/src/arrow/delete_filter.rs +++ b/crates/iceberg/src/arrow/delete_filter.rs @@ -16,7 +16,7 @@ // under the License. use std::collections::HashMap; -use std::sync::{Arc, Mutex, RwLock}; +use std::sync::{Arc, Mutex, OnceLock, PoisonError, RwLock}; use tokio::sync::Notify; use tokio::sync::futures::OwnedNotified; @@ -42,8 +42,8 @@ enum EqDelState { #[derive(Debug)] enum PosDelState { /// The file is currently being loaded by a task. - /// The notifier allows other tasks to wait for completion. - Loading(Arc<Notify>), + /// Other tasks wait on the attempt for it to finish or fail. + Loading(Arc<PosDelAttempt>), /// The file has been fully loaded and merged into the delete vector map. Loaded, } @@ -61,18 +61,83 @@ pub(crate) struct DeleteFilter { runtime: Runtime, } +/// One attempt to load a positional delete file, shared with the tasks waiting on it. +#[derive(Debug, Default)] +pub(crate) struct PosDelAttempt { + notify: Arc<Notify>, + /// Set when the load fails, so the tasks waiting on it get the error instead of + /// each opening the file again. + failure: OnceLock<(ErrorKind, String)>, +} + +/// A wait on another task's positional delete load, from [`PosDelLoadAction::WaitFor`]. +pub(crate) struct PosDelWait { + notified: OwnedNotified, + attempt: Arc<PosDelAttempt>, +} + +impl PosDelWait { + /// Waits for the load to end. Returns its error if it failed. `Ok` means it finished + /// or was dropped without an error, so the caller must call + /// [`DeleteFilter::try_start_pos_del_load`] again rather than assume the file is loaded. + pub(crate) async fn wait(self) -> Result<()> { + self.notified.await; + match self.attempt.failure.get() { + Some((kind, message)) => Err(Error::new(*kind, message.clone())), + None => Ok(()), + } + } +} + /// Action to take when trying to start loading a positional delete file pub(crate) enum PosDelLoadAction { - /// The file is not loaded, the caller should load it. - Load, + /// The file is not loaded, the caller should load it and call + /// [`PosDelLoadGuard::finish`] once its delete vectors are merged. + Load(PosDelLoadGuard), /// The file is already loaded, nothing to do. AlreadyLoaded, /// The file is currently being loaded by another task. - /// The caller *must* await this future to ensure data availability before + /// The caller *must* await [`PosDelWait::wait`] to ensure data availability before /// returning, as subsequent access (get_delete_vector) is synchronous. The - /// future is created under the state lock so it cannot miss the loader's + /// wait is created under the state lock so it cannot miss the loader's /// `notify_waiters()` (which stores no permit). - WaitFor(OwnedNotified), + WaitFor(PosDelWait), +} + +/// Ownership of an in-progress positional delete load, returned by +/// [`DeleteFilter::try_start_pos_del_load`]. +/// +/// [`finish`](Self::finish) marks the file loaded. [`fail`](Self::fail) hands the error to +/// the tasks waiting on this attempt. Dropping the guard without finishing removes the +/// `Loading` marker and wakes the waiters, so a later attempt can load the file instead of +/// waiting forever. +pub(crate) struct PosDelLoadGuard { + filter: DeleteFilter, + file_path: String, + attempt: Arc<PosDelAttempt>, + finished: bool, +} + +impl PosDelLoadGuard { + /// Marks the file as loaded and wakes any waiters. + pub(crate) fn finish(mut self) { + self.finished = true; + self.filter.finish_pos_del_load(&self.file_path); + } + + /// Ends the load with `error`, which the tasks waiting on it return. Later calls to + /// [`DeleteFilter::try_start_pos_del_load`] start a new attempt. + pub(crate) fn fail(self, error: &Error) { + let _ = self.attempt.failure.set((error.kind(), error.to_string())); + } +} + +impl Drop for PosDelLoadGuard { + fn drop(&mut self) { + if !self.finished { + self.filter.abandon_pos_del_load(&self.file_path); + } + } } impl DeleteFilter { @@ -130,36 +195,61 @@ impl DeleteFilter { if let Some(state) = state.positional_deletes.get(file_path) { match state { PosDelState::Loaded => return PosDelLoadAction::AlreadyLoaded, - PosDelState::Loading(notify) => { - return PosDelLoadAction::WaitFor(notify.clone().notified_owned()); + PosDelState::Loading(attempt) => { + return PosDelLoadAction::WaitFor(PosDelWait { + notified: attempt.notify.clone().notified_owned(), + attempt: attempt.clone(), + }); } } } - let notifier = Arc::new(Notify::new()); + let attempt = Arc::new(PosDelAttempt::default()); state .positional_deletes - .insert(file_path.to_string(), PosDelState::Loading(notifier)); - - PosDelLoadAction::Load + .insert(file_path.to_string(), PosDelState::Loading(attempt.clone())); + + PosDelLoadAction::Load(PosDelLoadGuard { + filter: self.clone(), + file_path: file_path.to_string(), + attempt, + finished: false, + }) } /// Marks a positional delete file as successfully loaded and notifies any waiting tasks. - pub(crate) fn finish_pos_del_load(&self, file_path: &str) { + fn finish_pos_del_load(&self, file_path: &str) { let notify = { let mut state = self.state.write().unwrap(); - if let Some(PosDelState::Loading(notify)) = state + if let Some(PosDelState::Loading(attempt)) = state .positional_deletes .insert(file_path.to_string(), PosDelState::Loaded) { - Some(notify) + Some(attempt) } else { None } }; - if let Some(notify) = notify { - notify.notify_waiters(); + if let Some(attempt) = notify { + attempt.notify.notify_waiters(); + } + } + + /// Removes the `Loading` marker of a positional delete file whose load failed or was + /// dropped, and wakes the tasks waiting on it. + fn abandon_pos_del_load(&self, file_path: &str) { + // This runs from `Drop`, possibly while unwinding from a panic that poisoned the + // lock; a second panic here would abort. + let removed = self + .state + .write() + .unwrap_or_else(PoisonError::into_inner) + .positional_deletes + .remove(file_path); + debug_assert!(matches!(removed, Some(PosDelState::Loading(_)))); + if let Some(PosDelState::Loading(attempt)) = removed { + attempt.notify.notify_waiters(); } } @@ -319,25 +409,89 @@ pub(crate) mod tests { let filter = DeleteFilter::new(Runtime::current()); let path = "s3://bucket/pos-delete.parquet"; - assert!(matches!( - filter.try_start_pos_del_load(path), - PosDelLoadAction::Load - )); + let PosDelLoadAction::Load(guard) = filter.try_start_pos_del_load(path) else { + panic!("expected Load for a new file"); + }; - let PosDelLoadAction::WaitFor(notified) = filter.try_start_pos_del_load(path) else { + let PosDelLoadAction::WaitFor(wait) = filter.try_start_pos_del_load(path) else { panic!("expected WaitFor for an in-progress load"); }; // Loader completes and signals before the waiter awaits. - filter.finish_pos_del_load(path); + guard.finish(); - let waited = tokio::time::timeout(std::time::Duration::from_secs(5), notified).await; + let waited = tokio::time::timeout(std::time::Duration::from_secs(5), wait.wait()).await; assert!( - waited.is_ok(), + matches!(waited, Ok(Ok(()))), "WaitFor future must resolve after finish_pos_del_load" ); } + #[tokio::test] + async fn test_cancelled_pos_del_load_wakes_waiters_and_can_be_retried() { + let filter = DeleteFilter::new(Runtime::current()); + let path = "s3://bucket/pos-delete.parquet"; + + let PosDelLoadAction::Load(guard) = filter.try_start_pos_del_load(path) else { + panic!("expected Load for a new file"); + }; + let PosDelLoadAction::WaitFor(wait) = filter.try_start_pos_del_load(path) else { + panic!("expected WaitFor for an in-progress load"); + }; + + // The load is cancelled: its guard is dropped without `finish` or `fail`. + drop(guard); + + let waited = tokio::time::timeout(std::time::Duration::from_secs(5), wait.wait()).await; + assert!( + matches!(waited, Ok(Ok(()))), + "WaitFor future must resolve when the load is dropped" + ); + + // The file is not loaded, so the next caller is told to load it. + assert!(matches!( + filter.try_start_pos_del_load(path), + PosDelLoadAction::Load(_) + )); + } + + #[tokio::test] + async fn test_failed_pos_del_load_returns_its_error_to_waiters() { + let filter = DeleteFilter::new(Runtime::current()); + let path = "s3://bucket/pos-delete.parquet"; + + let PosDelLoadAction::Load(guard) = filter.try_start_pos_del_load(path) else { + panic!("expected Load for a new file"); + }; + let waits: Vec<_> = (0..2) + .map(|_| match filter.try_start_pos_del_load(path) { + PosDelLoadAction::WaitFor(wait) => wait, + _ => panic!("expected WaitFor for an in-progress load"), + }) + .collect(); + + guard.fail( + &Error::new(ErrorKind::DataInvalid, "bad file") + .with_source(std::io::Error::other("403 Forbidden")), + ); + + for wait in waits { + let err = tokio::time::timeout(std::time::Duration::from_secs(5), wait.wait()) + .await + .expect("WaitFor future must resolve when the load fails") + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains("bad file"), "{err}"); + assert!(err.message().contains("403 Forbidden"), "{err}"); + } + + // Only the waiters of that attempt get the error; a later call loads again. + assert!(matches!( + filter.try_start_pos_del_load(path), + PosDelLoadAction::Load(_) + )); + } + #[tokio::test] async fn test_delete_file_filter_load_deletes() { let tmp_dir = TempDir::new().unwrap();
