From ceea11433fc464a204440af88281eb9a04663cd5 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Fri, 23 Oct 2026 12:18:15 +0100 Subject: [PATCH] fix(cuda): bound concurrent stream-pool initialization Serialize first use within each stream-pool slot while retaining lock-free access to initialized slots. Include lazy-creation, round-robin, and concurrent-capacity regression coverage. Original-commit: ae60fd7950 Signed-off-by: Alexander Droste --- Cargo.lock | 2 +- vortex-cuda/Cargo.toml | 2 +- vortex-cuda/src/stream_pool.rs | 137 ++++++++++++++++++++++++++------- 3 files changed, 113 insertions(+), 28 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a471337bea4..5e7e3fa9b69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11731,7 +11731,6 @@ dependencies = [ name = "vortex-cuda" version = "0.1.0" dependencies = [ - "arc-swap", "arrow-schema 59.3.1", "async-trait", "bindgen", @@ -11744,6 +11743,7 @@ dependencies = [ "mimalloc", "num-traits", "object_store 0.13.2", + "once_cell", "parking_lot", "prost 0.14.4", "rstest", diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index eaffffd1251..0e82ddb1042 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -21,7 +21,6 @@ default = [] _test-harness = [] [dependencies] -arc-swap = { workspace = true } arrow-schema = { workspace = true, features = ["ffi"] } async-trait = { workspace = true } cudarc = { workspace = true, features = ["f16"] } @@ -30,6 +29,7 @@ itertools = { workspace = true } kanal = { workspace = true } num-traits = { workspace = true } object_store = { workspace = true, features = ["fs"] } +once_cell = { workspace = true } parking_lot = { workspace = true } prost = { workspace = true } tokio = { workspace = true, features = ["fs"] } diff --git a/vortex-cuda/src/stream_pool.rs b/vortex-cuda/src/stream_pool.rs index 6f8e66798c5..7f1548e311b 100644 --- a/vortex-cuda/src/stream_pool.rs +++ b/vortex-cuda/src/stream_pool.rs @@ -7,9 +7,9 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use arc_swap::ArcSwapOption; use cudarc::driver::CudaContext; use cudarc::driver::CudaStream; +use once_cell::sync::OnceCell; use vortex::error::VortexResult; use vortex::error::vortex_err; @@ -17,13 +17,12 @@ use crate::stream::VortexCudaStream; /// A pool of CUDA streams that hands out streams in a round-robin fashion. /// -/// Uses lock-free slot access via `ArcSwap` for high-concurrency scenarios. -/// Streams are lazily created on first access to each slot and remain alive -/// for the lifetime of the pool. +/// Streams are lazily initialized once per slot and remain alive for the lifetime of the pool. +/// Access to initialized slots is lock-free; failed initialization can be retried. pub struct VortexCudaStreamPool { context: Arc, /// Fixed-size array of slots, each holding an optional stream. - slots: Box<[ArcSwapOption]>, + slots: Box<[OnceCell>]>, /// Round-robin counter for slot selection. next_index: AtomicUsize, } @@ -46,7 +45,7 @@ impl VortexCudaStreamPool { /// * `max_capacity` - Maximum number of streams to maintain in the pool. pub fn new(context: Arc, max_capacity: usize) -> Self { let slots = (0..max_capacity) - .map(|_| ArcSwapOption::empty()) + .map(|_| OnceCell::new()) .collect::>() .into_boxed_slice(); @@ -61,35 +60,23 @@ impl VortexCudaStreamPool { /// /// Uses round-robin slot selection. If the selected slot has a stream, /// it is reused. Otherwise, a new stream is created for that slot. - /// All operations are lock-free. + /// Access to initialized slots is lock-free; initialization is serialized per slot. pub fn stream(&self) -> VortexResult { let slot_idx = self.next_index.fetch_add(1, Ordering::Relaxed) % self.slots.len(); - let slot = &self.slots[slot_idx]; + let stream = self.slots[slot_idx].get_or_try_init(|| { + self.context + .new_stream() + .map_err(|e| vortex_err!("Failed to create CUDA stream: {}", e)) + })?; - // Fast path: stream already exists in slot. - if let Some(stream) = slot.load_full() { - return Ok(VortexCudaStream(stream)); - } - - // Slow path: create a new stream. - // Note: CudaContext::new_stream() already returns Arc. - let new_stream = self - .context - .new_stream() - .map_err(|e| vortex_err!("Failed to create CUDA stream: {}", e))?; - - // Store it in the slot. If another thread raced us, that's fine - - // we'll just use our newly created stream this time. - slot.store(Some(Arc::clone(&new_stream))); - - Ok(VortexCudaStream(new_stream)) + Ok(VortexCudaStream(Arc::clone(stream))) } /// Returns the current number of initialized streams in the pool. pub fn live_stream_count(&self) -> usize { self.slots .iter() - .filter(|slot| slot.load().is_some()) + .filter(|slot| slot.get().is_some()) .count() } @@ -98,3 +85,101 @@ impl VortexCudaStreamPool { self.slots.len() } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::Barrier; + use std::thread; + + use cudarc::driver::CudaContext; + use vortex::error::VortexResult; + use vortex::error::vortex_err; + use vortex::utils::aliases::hash_set::HashSet; + + use super::VortexCudaStreamPool; + + #[crate::test] + fn lazy_creation() -> VortexResult<()> { + let context = + CudaContext::new(0).map_err(|e| vortex_err!("Failed to create CUDA context: {e}"))?; + let pool = VortexCudaStreamPool::new(context, 3); + + assert_eq!(pool.max_capacity(), 3); + assert_eq!(pool.live_stream_count(), 0); + for expected_live in 1..=3 { + let _stream = pool.stream()?; + assert_eq!(pool.live_stream_count(), expected_live); + } + let _stream = pool.stream()?; + assert_eq!(pool.live_stream_count(), 3); + Ok(()) + } + + #[crate::test] + fn round_robin_reuse() -> VortexResult<()> { + let context = + CudaContext::new(0).map_err(|e| vortex_err!("Failed to create CUDA context: {e}"))?; + let pool = VortexCudaStreamPool::new(context, 3); + let streams = (0..pool.max_capacity()) + .map(|_| pool.stream()) + .collect::>>()?; + let handles = streams + .iter() + .map(|stream| stream.cu_stream() as usize) + .collect::>(); + assert_eq!(handles.len(), pool.max_capacity()); + + for _ in 0..3 { + for expected in &streams { + let stream = pool.stream()?; + assert!(Arc::ptr_eq(&stream.0, &expected.0)); + assert_eq!(stream.cu_stream(), expected.cu_stream()); + } + } + assert_eq!(pool.live_stream_count(), pool.max_capacity()); + Ok(()) + } + + #[crate::test] + fn concurrent_initialization_respects_capacity() -> VortexResult<()> { + let context = + CudaContext::new(0).map_err(|e| vortex_err!("Failed to create CUDA context: {e}"))?; + for capacity in [1, 4] { + let pool = VortexCudaStreamPool::new(Arc::clone(&context), capacity); + let barrier = Barrier::new(32); + let streams = thread::scope(|scope| { + let workers = (0..32) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + (0..8) + .map(|_| pool.stream()) + .collect::>>() + }) + }) + .collect::>(); + workers + .into_iter() + .map(|worker| worker.join().expect("stream worker panicked")) + .collect::>>() + })?; + + // Keep every returned stream alive so transient handles cannot be recycled. + let handles = streams + .iter() + .flatten() + .map(|stream| stream.cu_stream() as usize) + .collect::>(); + assert_eq!(handles.len(), capacity); + assert_eq!(pool.live_stream_count(), capacity); + let stored = (0..capacity) + .map(|_| pool.stream()) + .collect::>>()?; + for stream in streams.iter().flatten() { + assert!(stored.iter().any(|slot| Arc::ptr_eq(&stream.0, &slot.0))); + } + } + Ok(()) + } +}