diff --git a/vortex-cuda/src/pooled_read_at/file.rs b/vortex-cuda/src/pooled_read_at/file.rs index afa8d736f2e..dea1e26f693 100644 --- a/vortex-cuda/src/pooled_read_at/file.rs +++ b/vortex-cuda/src/pooled_read_at/file.rs @@ -12,6 +12,7 @@ use std::os::unix::fs::MetadataExt; use std::path::Path; use std::sync::Arc; use std::sync::LazyLock; +use std::sync::mpsc; use std::time::SystemTime; use cudarc::driver::CudaFunction; @@ -65,6 +66,24 @@ const MAPPED_COPY_THREADS: u32 = 256; #[cfg(target_os = "linux")] static MAPPED_FILL: Mutex<()> = Mutex::new(()); +/// Releases a closed reader's mapping once its copy stream has drained. +type MappedRelease = Box; + +/// Releases of mapped files, each run once its copy stream has drained, off the scan's threads. +static MAPPED_RELEASE_QUEUE: LazyLock>> = LazyLock::new(|| { + let (sender, receiver) = mpsc::channel::(); + std::thread::Builder::new() + .name("vortex-mapped-release".to_string()) + .spawn(move || { + // Runs until the process exits: the sender is a process-wide static. + for job in receiver { + job(); + } + }) + .ok() + .map(|_| sender) +}); + /// Options controlling how [`PooledFileReadAt`] opens and reads a local file. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct PooledFileReadAtOptions { @@ -231,6 +250,19 @@ impl FileMapping { /// cold-cache measurement, or to return the address space. Mappings that open readers still use /// stay mapped until those readers close. pub fn release_file_mappings() { + // Closed readers hand their mappings to the release queue; wait until it has dropped them. + if let Some(queue) = MAPPED_RELEASE_QUEUE.as_ref() { + let (sender, receiver) = mpsc::sync_channel::<()>(1); + if queue + .send(Box::new(move || { + let _ = sender.send(()); + })) + .is_ok() + { + // An error means the release thread is gone, leaving nothing to wait for. + let _ = receiver.recv(); + } + } let released: Vec<_> = FILE_MAPPINGS .lock() .drain(..) @@ -389,10 +421,24 @@ impl Drop for MappedFile { let Some(file) = self.file.take() else { return; }; - // Queued copies may still read the mapping, which the last release unmaps. - if self.stream.synchronize().is_err() { - // Without proof the copies finished, never unmap memory the device may read. - let _leaked = ManuallyDrop::new(file); + // Queued copies may still read the mapping, which the last release unmaps. The reader + // drops when the scan has requested its last segment, typically on the thread that + // decodes the scan's batches; waiting there for the copy stream to drain stalled every + // decode and selection launch behind the scan's remaining copies. + let stream = self.stream.clone(); + let release: MappedRelease = Box::new(move || { + if stream.synchronize().is_err() { + // Without proof the copies finished, never unmap memory the device may read. + let _leaked = ManuallyDrop::new(file); + } + }); + match MAPPED_RELEASE_QUEUE.as_ref() { + Some(queue) => { + if let Err(unsent) = queue.send(release) { + (unsent.0)(); + } + } + None => release(), } } } diff --git a/vortex-cuda/src/pooled_read_at/file/tests.rs b/vortex-cuda/src/pooled_read_at/file/tests.rs index c2a5121c216..d75440e6c5c 100644 --- a/vortex-cuda/src/pooled_read_at/file/tests.rs +++ b/vortex-cuda/src/pooled_read_at/file/tests.rs @@ -491,6 +491,42 @@ async fn released_mappings_are_remapped() -> VortexResult<()> { Ok(()) } +/// Closing a reader releases its mapping in the background once its copies drain; releasing +/// cached mappings afterwards must still unmap it before returning, as cold-cache eviction needs. +#[crate::test] +async fn closed_reader_mappings_are_released_before_return() -> VortexResult<()> { + let (_runtime, reader, _file) = mapped_reader()?; + let Some(mapped) = reader.mapped.as_ref() else { + return Ok(()); + }; + let mapping = Arc::downgrade( + mapped + .file + .as_ref() + .vortex_expect("open file has a mapping"), + ); + #[cfg(target_os = "linux")] + mapped + .file + .as_ref() + .vortex_expect("open file has a mapping") + .mapping() + .advise(Advice::PopulateRead)?; + // Leave copies queued on the copy stream when the reader closes. + let buffer = reader + .read_at(0, MAPPED_FILE_BYTES, Alignment::of::()) + .await?; + // Only the release must not race other tests' use of the global cache; the guard is not + // held across an await. + let cache = MAPPING_CACHE.lock(); + drop(reader); + release_file_mappings(); + drop(cache); + assert!(mapping.upgrade().is_none()); + assert_bytes(buffer, 0, MAPPED_FILE_BYTES).await?; + Ok(()) +} + #[crate::test] async fn mapped_read_past_end_fails() -> VortexResult<()> { let (_runtime, reader, _file) = mapped_reader()?;