Skip to content

Commit a74fdc7

Browse files
authored
perf(cuda): release closed mappings off the decode thread (#10718)
## Problem Closing a mapped reader can stall decoding while waiting for GPU copies and releasing page tables. ## Solution Move synchronization and mapping release to a background worker; drain queued releases before explicit cache eviction. ## Assessment Falls back to synchronous release if queuing fails and retains mappings if completion is unknown. Depends on #10717. Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>
1 parent c7a6cc9 commit a74fdc7

2 files changed

Lines changed: 86 additions & 4 deletions

File tree

‎vortex-cuda/src/pooled_read_at/file.rs‎

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use std::os::unix::fs::MetadataExt;
1212
use std::path::Path;
1313
use std::sync::Arc;
1414
use std::sync::LazyLock;
15+
use std::sync::mpsc;
1516
use std::time::SystemTime;
1617

1718
use cudarc::driver::CudaFunction;
@@ -65,6 +66,24 @@ const MAPPED_COPY_THREADS: u32 = 256;
6566
#[cfg(target_os = "linux")]
6667
static MAPPED_FILL: Mutex<()> = Mutex::new(());
6768

69+
/// Releases a closed reader's mapping once its copy stream has drained.
70+
type MappedRelease = Box<dyn FnOnce() + Send>;
71+
72+
/// Releases of mapped files, each run once its copy stream has drained, off the scan's threads.
73+
static MAPPED_RELEASE_QUEUE: LazyLock<Option<mpsc::Sender<MappedRelease>>> = LazyLock::new(|| {
74+
let (sender, receiver) = mpsc::channel::<MappedRelease>();
75+
std::thread::Builder::new()
76+
.name("vortex-mapped-release".to_string())
77+
.spawn(move || {
78+
// Runs until the process exits: the sender is a process-wide static.
79+
for job in receiver {
80+
job();
81+
}
82+
})
83+
.ok()
84+
.map(|_| sender)
85+
});
86+
6887
/// Options controlling how [`PooledFileReadAt`] opens and reads a local file.
6988
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
7089
pub struct PooledFileReadAtOptions {
@@ -231,6 +250,19 @@ impl FileMapping {
231250
/// cold-cache measurement, or to return the address space. Mappings that open readers still use
232251
/// stay mapped until those readers close.
233252
pub fn release_file_mappings() {
253+
// Closed readers hand their mappings to the release queue; wait until it has dropped them.
254+
if let Some(queue) = MAPPED_RELEASE_QUEUE.as_ref() {
255+
let (sender, receiver) = mpsc::sync_channel::<()>(1);
256+
if queue
257+
.send(Box::new(move || {
258+
let _ = sender.send(());
259+
}))
260+
.is_ok()
261+
{
262+
// An error means the release thread is gone, leaving nothing to wait for.
263+
let _ = receiver.recv();
264+
}
265+
}
234266
let released: Vec<_> = FILE_MAPPINGS
235267
.lock()
236268
.drain(..)
@@ -389,10 +421,24 @@ impl Drop for MappedFile {
389421
let Some(file) = self.file.take() else {
390422
return;
391423
};
392-
// Queued copies may still read the mapping, which the last release unmaps.
393-
if self.stream.synchronize().is_err() {
394-
// Without proof the copies finished, never unmap memory the device may read.
395-
let _leaked = ManuallyDrop::new(file);
424+
// Queued copies may still read the mapping, which the last release unmaps. The reader
425+
// drops when the scan has requested its last segment, typically on the thread that
426+
// decodes the scan's batches; waiting there for the copy stream to drain stalled every
427+
// decode and selection launch behind the scan's remaining copies.
428+
let stream = self.stream.clone();
429+
let release: MappedRelease = Box::new(move || {
430+
if stream.synchronize().is_err() {
431+
// Without proof the copies finished, never unmap memory the device may read.
432+
let _leaked = ManuallyDrop::new(file);
433+
}
434+
});
435+
match MAPPED_RELEASE_QUEUE.as_ref() {
436+
Some(queue) => {
437+
if let Err(unsent) = queue.send(release) {
438+
(unsent.0)();
439+
}
440+
}
441+
None => release(),
396442
}
397443
}
398444
}

‎vortex-cuda/src/pooled_read_at/file/tests.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -491,6 +491,42 @@ async fn released_mappings_are_remapped() -> VortexResult<()> {
491491
Ok(())
492492
}
493493

494+
/// Closing a reader releases its mapping in the background once its copies drain; releasing
495+
/// cached mappings afterwards must still unmap it before returning, as cold-cache eviction needs.
496+
#[crate::test]
497+
async fn closed_reader_mappings_are_released_before_return() -> VortexResult<()> {
498+
let (_runtime, reader, _file) = mapped_reader()?;
499+
let Some(mapped) = reader.mapped.as_ref() else {
500+
return Ok(());
501+
};
502+
let mapping = Arc::downgrade(
503+
mapped
504+
.file
505+
.as_ref()
506+
.vortex_expect("open file has a mapping"),
507+
);
508+
#[cfg(target_os = "linux")]
509+
mapped
510+
.file
511+
.as_ref()
512+
.vortex_expect("open file has a mapping")
513+
.mapping()
514+
.advise(Advice::PopulateRead)?;
515+
// Leave copies queued on the copy stream when the reader closes.
516+
let buffer = reader
517+
.read_at(0, MAPPED_FILE_BYTES, Alignment::of::<u8>())
518+
.await?;
519+
// Only the release must not race other tests' use of the global cache; the guard is not
520+
// held across an await.
521+
let cache = MAPPING_CACHE.lock();
522+
drop(reader);
523+
release_file_mappings();
524+
drop(cache);
525+
assert!(mapping.upgrade().is_none());
526+
assert_bytes(buffer, 0, MAPPED_FILE_BYTES).await?;
527+
Ok(())
528+
}
529+
494530
#[crate::test]
495531
async fn mapped_read_past_end_fails() -> VortexResult<()> {
496532
let (_runtime, reader, _file) = mapped_reader()?;

0 commit comments

Comments
 (0)