fix(cuda): bound concurrent stream-pool initialization

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"] }

vortex-cuda/src/stream_pool.rs

@@ -7,23 +7,22 @@ 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;
 
 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<CudaContext>,
     /// Fixed-size array of slots, each holding an optional stream.
-    slots: Box<[ArcSwapOption<CudaStream>]>,
+    slots: Box<[OnceCell<Arc<CudaStream>>]>,
     /// 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<CudaContext>, max_capacity: usize) -> Self {
         let slots = (0..max_capacity)
-            .map(|_| ArcSwapOption::empty())
+            .map(|_| OnceCell::new())
             .collect::<Vec<_>>()
             .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<VortexCudaStream> {
         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<CudaStream>.
-        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::<VortexResult<Vec<_>>>()?;
+        let handles = streams
+            .iter()
+            .map(|stream| stream.cu_stream() as usize)
+            .collect::<HashSet<_>>();
+        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::<VortexResult<Vec<_>>>()
+                        })
+                    })
+                    .collect::<Vec<_>>();
+                workers
+                    .into_iter()
+                    .map(|worker| worker.join().expect("stream worker panicked"))
+                    .collect::<VortexResult<Vec<_>>>()
+            })?;
+
+            // Keep every returned stream alive so transient handles cannot be recycled.
+            let handles = streams
+                .iter()
+                .flatten()
+                .map(|stream| stream.cu_stream() as usize)
+                .collect::<HashSet<_>>();
+            assert_eq!(handles.len(), capacity);
+            assert_eq!(pool.live_stream_count(), capacity);
+            let stored = (0..capacity)
+                .map(|_| pool.stream())
+                .collect::<VortexResult<Vec<_>>>()?;
+            for stream in streams.iter().flatten() {
+                assert!(stored.iter().any(|slot| Arc::ptr_eq(&stream.0, &slot.0)));
+            }
+        }
+        Ok(())
+    }
+}