@@ -7,23 +7,22 @@ use std::sync::Arc;
77use std:: sync:: atomic:: AtomicUsize ;
88use std:: sync:: atomic:: Ordering ;
99
10- use arc_swap:: ArcSwapOption ;
1110use cudarc:: driver:: CudaContext ;
1211use cudarc:: driver:: CudaStream ;
12+ use once_cell:: sync:: OnceCell ;
1313use vortex:: error:: VortexResult ;
1414use vortex:: error:: vortex_err;
1515
1616use crate :: stream:: VortexCudaStream ;
1717
1818/// A pool of CUDA streams that hands out streams in a round-robin fashion.
1919///
20- /// Uses lock-free slot access via `ArcSwap` for high-concurrency scenarios.
21- /// Streams are lazily created on first access to each slot and remain alive
22- /// for the lifetime of the pool.
20+ /// Streams are lazily initialized once per slot and remain alive for the lifetime of the pool.
21+ /// Access to initialized slots is lock-free; failed initialization can be retried.
2322pub struct VortexCudaStreamPool {
2423 context : Arc < CudaContext > ,
2524 /// Fixed-size array of slots, each holding an optional stream.
26- slots : Box < [ ArcSwapOption < CudaStream > ] > ,
25+ slots : Box < [ OnceCell < Arc < CudaStream > > ] > ,
2726 /// Round-robin counter for slot selection.
2827 next_index : AtomicUsize ,
2928}
@@ -46,7 +45,7 @@ impl VortexCudaStreamPool {
4645 /// * `max_capacity` - Maximum number of streams to maintain in the pool.
4746 pub fn new ( context : Arc < CudaContext > , max_capacity : usize ) -> Self {
4847 let slots = ( 0 ..max_capacity)
49- . map ( |_| ArcSwapOption :: empty ( ) )
48+ . map ( |_| OnceCell :: new ( ) )
5049 . collect :: < Vec < _ > > ( )
5150 . into_boxed_slice ( ) ;
5251
@@ -61,35 +60,23 @@ impl VortexCudaStreamPool {
6160 ///
6261 /// Uses round-robin slot selection. If the selected slot has a stream,
6362 /// it is reused. Otherwise, a new stream is created for that slot.
64- /// All operations are lock-free.
63+ /// Access to initialized slots is lock-free; initialization is serialized per slot .
6564 pub fn stream ( & self ) -> VortexResult < VortexCudaStream > {
6665 let slot_idx = self . next_index . fetch_add ( 1 , Ordering :: Relaxed ) % self . slots . len ( ) ;
67- let slot = & self . slots [ slot_idx] ;
66+ let stream = self . slots [ slot_idx] . get_or_try_init ( || {
67+ self . context
68+ . new_stream ( )
69+ . map_err ( |e| vortex_err ! ( "Failed to create CUDA stream: {}" , e) )
70+ } ) ?;
6871
69- // Fast path: stream already exists in slot.
70- if let Some ( stream) = slot. load_full ( ) {
71- return Ok ( VortexCudaStream ( stream) ) ;
72- }
73-
74- // Slow path: create a new stream.
75- // Note: CudaContext::new_stream() already returns Arc<CudaStream>.
76- let new_stream = self
77- . context
78- . new_stream ( )
79- . map_err ( |e| vortex_err ! ( "Failed to create CUDA stream: {}" , e) ) ?;
80-
81- // Store it in the slot. If another thread raced us, that's fine -
82- // we'll just use our newly created stream this time.
83- slot. store ( Some ( Arc :: clone ( & new_stream) ) ) ;
84-
85- Ok ( VortexCudaStream ( new_stream) )
72+ Ok ( VortexCudaStream ( Arc :: clone ( stream) ) )
8673 }
8774
8875 /// Returns the current number of initialized streams in the pool.
8976 pub fn live_stream_count ( & self ) -> usize {
9077 self . slots
9178 . iter ( )
92- . filter ( |slot| slot. load ( ) . is_some ( ) )
79+ . filter ( |slot| slot. get ( ) . is_some ( ) )
9380 . count ( )
9481 }
9582
@@ -98,3 +85,101 @@ impl VortexCudaStreamPool {
9885 self . slots . len ( )
9986 }
10087}
88+
89+ #[ cfg( test) ]
90+ mod tests {
91+ use std:: sync:: Arc ;
92+ use std:: sync:: Barrier ;
93+ use std:: thread;
94+
95+ use cudarc:: driver:: CudaContext ;
96+ use vortex:: error:: VortexResult ;
97+ use vortex:: error:: vortex_err;
98+ use vortex:: utils:: aliases:: hash_set:: HashSet ;
99+
100+ use super :: VortexCudaStreamPool ;
101+
102+ #[ crate :: test]
103+ fn lazy_creation ( ) -> VortexResult < ( ) > {
104+ let context =
105+ CudaContext :: new ( 0 ) . map_err ( |e| vortex_err ! ( "Failed to create CUDA context: {e}" ) ) ?;
106+ let pool = VortexCudaStreamPool :: new ( context, 3 ) ;
107+
108+ assert_eq ! ( pool. max_capacity( ) , 3 ) ;
109+ assert_eq ! ( pool. live_stream_count( ) , 0 ) ;
110+ for expected_live in 1 ..=3 {
111+ let _stream = pool. stream ( ) ?;
112+ assert_eq ! ( pool. live_stream_count( ) , expected_live) ;
113+ }
114+ let _stream = pool. stream ( ) ?;
115+ assert_eq ! ( pool. live_stream_count( ) , 3 ) ;
116+ Ok ( ( ) )
117+ }
118+
119+ #[ crate :: test]
120+ fn round_robin_reuse ( ) -> VortexResult < ( ) > {
121+ let context =
122+ CudaContext :: new ( 0 ) . map_err ( |e| vortex_err ! ( "Failed to create CUDA context: {e}" ) ) ?;
123+ let pool = VortexCudaStreamPool :: new ( context, 3 ) ;
124+ let streams = ( 0 ..pool. max_capacity ( ) )
125+ . map ( |_| pool. stream ( ) )
126+ . collect :: < VortexResult < Vec < _ > > > ( ) ?;
127+ let handles = streams
128+ . iter ( )
129+ . map ( |stream| stream. cu_stream ( ) as usize )
130+ . collect :: < HashSet < _ > > ( ) ;
131+ assert_eq ! ( handles. len( ) , pool. max_capacity( ) ) ;
132+
133+ for _ in 0 ..3 {
134+ for expected in & streams {
135+ let stream = pool. stream ( ) ?;
136+ assert ! ( Arc :: ptr_eq( & stream. 0 , & expected. 0 ) ) ;
137+ assert_eq ! ( stream. cu_stream( ) , expected. cu_stream( ) ) ;
138+ }
139+ }
140+ assert_eq ! ( pool. live_stream_count( ) , pool. max_capacity( ) ) ;
141+ Ok ( ( ) )
142+ }
143+
144+ #[ crate :: test]
145+ fn concurrent_initialization_respects_capacity ( ) -> VortexResult < ( ) > {
146+ let context =
147+ CudaContext :: new ( 0 ) . map_err ( |e| vortex_err ! ( "Failed to create CUDA context: {e}" ) ) ?;
148+ for capacity in [ 1 , 4 ] {
149+ let pool = VortexCudaStreamPool :: new ( Arc :: clone ( & context) , capacity) ;
150+ let barrier = Barrier :: new ( 32 ) ;
151+ let streams = thread:: scope ( |scope| {
152+ let workers = ( 0 ..32 )
153+ . map ( |_| {
154+ scope. spawn ( || {
155+ barrier. wait ( ) ;
156+ ( 0 ..8 )
157+ . map ( |_| pool. stream ( ) )
158+ . collect :: < VortexResult < Vec < _ > > > ( )
159+ } )
160+ } )
161+ . collect :: < Vec < _ > > ( ) ;
162+ workers
163+ . into_iter ( )
164+ . map ( |worker| worker. join ( ) . expect ( "stream worker panicked" ) )
165+ . collect :: < VortexResult < Vec < _ > > > ( )
166+ } ) ?;
167+
168+ // Keep every returned stream alive so transient handles cannot be recycled.
169+ let handles = streams
170+ . iter ( )
171+ . flatten ( )
172+ . map ( |stream| stream. cu_stream ( ) as usize )
173+ . collect :: < HashSet < _ > > ( ) ;
174+ assert_eq ! ( handles. len( ) , capacity) ;
175+ assert_eq ! ( pool. live_stream_count( ) , capacity) ;
176+ let stored = ( 0 ..capacity)
177+ . map ( |_| pool. stream ( ) )
178+ . collect :: < VortexResult < Vec < _ > > > ( ) ?;
179+ for stream in streams. iter ( ) . flatten ( ) {
180+ assert ! ( stored. iter( ) . any( |slot| Arc :: ptr_eq( & stream. 0 , & slot. 0 ) ) ) ;
181+ }
182+ }
183+ Ok ( ( ) )
184+ }
185+ }
0 commit comments