Skip to main content

cubecl_server/stream/
event.rs

1use crate::{
2    config::streaming::StreamingLogLevel,
3    logging::ServerLogger,
4    memory_management::{ErrorGraph, FailureId, ManagedMemoryId, SharedMemoryBindings},
5    server::{BufferBinding, ServerError},
6    stream::{FailureStore, Failures, StreamFactory, StreamMemory, StreamPool, base},
7};
8use core::any::Any;
9use cubecl_environment::collections::HashMap;
10use cubecl_environment::stream::StreamId;
11use std::{
12    boxed::Box,
13    format,
14    sync::{Arc, mpsc::SyncSender},
15    vec::Vec,
16};
17
18/// Trait defining the backend operations for managing streams and events.
19///
20/// This trait provides the necessary methods for initializing streams, flushing them to create events,
21/// and waiting on events for synchronization purposes.
22pub trait EventStreamBackend: 'static {
23    /// The type representing a stream in this backend, which exposes the
24    /// memory its kernels see through [`StreamMemory`].
25    type Stream: core::fmt::Debug + StreamMemory;
26    /// The type representing an event in this backend.
27    type Event: Send + 'static;
28
29    /// Initializes and returns a new stream associated with the given stream ID.
30    fn create_stream(&self) -> Self::Stream;
31    /// Returns the cursor of the given handle on the given stream.
32    fn handle_cursor(stream: &Self::Stream, handle: &BufferBinding) -> u64;
33    /// Flushes the given stream, ensuring all pending operations are submitted, and returns an event
34    /// that can be used for synchronization.
35    ///
36    /// `failures` is the device's failure store, handed down because a flush
37    /// may drive the stream's periodic memory cleanup, and cleaning up is
38    /// where a pool sheds the failures its slices carry.
39    fn flush(stream: &mut Self::Stream, failures: &mut ErrorGraph) -> Self::Event;
40    /// Makes the stream wait for the specified event to complete before proceeding with further operations.
41    fn wait_event(stream: &mut Self::Stream, event: Self::Event);
42    /// Wait for the given event synching the CPU.
43    fn wait_event_sync(event: Self::Event) -> Result<(), ServerError>;
44}
45
46/// Manages multiple streams with synchronization logic based on shared bindings.
47///
48/// This struct handles the creation and alignment of streams to ensure proper synchronization
49/// when bindings (e.g., buffers) are shared across different streams.
50#[derive(Debug)]
51pub struct MultiStream<B: EventStreamBackend> {
52    /// The map of stream IDs to their corresponding stream wrappers.
53    streams: StreamPool<EventStreamBackendWrapper<B>>,
54    /// Every failure the device is still holding, and the write scope's
55    /// scratch — see [`Failures`].
56    failures: Failures,
57    /// The logger used by the server.
58    pub logger: Arc<ServerLogger>,
59    max_streams: usize,
60    gc: GcThread<B>,
61    shared_bindings_pool: Vec<(ManagedMemoryId, StreamId, u64)>,
62}
63
64/// A wrapper around a backend stream that includes synchronization metadata.
65///
66/// This includes the stream itself, a map of last synchronized cursors from other streams,
67/// and the current cursor position for this stream.
68/// A backend stream plus the synchronization metadata the pool keeps beside
69/// it. Public only because it is the pool's stream type in
70/// [`FailureStore::Factory`]; backends reach the stream itself, not this.
71pub struct StreamWrapper<B: EventStreamBackend> {
72    /// The underlying backend stream.
73    stream: B::Stream,
74    /// The current cursor position, representing the logical progress or version of operations on this stream.
75    cursor: u64,
76    /// A map tracking the last synchronized cursor positions from other streams.
77    last_synced: HashMap<usize, u64>,
78}
79
80/// Streams that are synchronized correctly after a [`MultiStream::resolve`] is called.
81pub struct ResolvedStreams<'a, B: EventStreamBackend> {
82    /// The cursor on the current stream.
83    ///
84    /// This cursor should be use for new allocations happening on the current stream.
85    pub cursor: u64,
86    streams: &'a mut StreamPool<EventStreamBackendWrapper<B>>,
87    failures: &'a mut ErrorGraph,
88    analysis: SharedBindingAnalysis,
89    gc: &'a GcThread<B>,
90    /// The current stream where new tasks can be sent safely.
91    pub current: StreamId,
92}
93
94#[derive(Debug)]
95/// A task to be enqueue on the gc stream that will be clearned after an event is reached.
96pub struct GcTask<B: EventStreamBackend> {
97    to_drop: Box<dyn Any + Send + 'static>,
98    /// The event to sync making sure the bindings in the batch are ready to be reused by other streams.
99    event: B::Event,
100}
101
102impl<B: EventStreamBackend> GcTask<B> {
103    /// Creates a new task that will be clearned when the event is reached.
104    pub fn new<T: Send + 'static>(to_drop: T, event: B::Event) -> Self {
105        Self {
106            to_drop: Box::new(to_drop),
107            event,
108        }
109    }
110}
111
112#[derive(Debug)]
113/// The factory a [`MultiStream`]'s pool is built from. Public only because it
114/// names the pool in [`FailureStore::Factory`]; it has no surface of its own.
115pub struct EventStreamBackendWrapper<B: EventStreamBackend> {
116    backend: B,
117}
118
119impl<B: EventStreamBackend> StreamMemory for StreamWrapper<B> {
120    fn failure(&self, binding: &BufferBinding) -> Option<FailureId> {
121        self.stream.failure(binding)
122    }
123
124    fn taint(&mut self, binding: &BufferBinding, failure: FailureId, failures: &mut ErrorGraph) {
125        self.stream.taint(binding, failure, failures)
126    }
127
128    fn written(&mut self, binding: &BufferBinding, failures: &mut ErrorGraph) {
129        self.stream.written(binding, failures)
130    }
131}
132
133impl<B: EventStreamBackend> StreamFactory for EventStreamBackendWrapper<B> {
134    type Stream = StreamWrapper<B>;
135
136    fn create(&mut self) -> Self::Stream {
137        StreamWrapper {
138            stream: self.backend.create_stream(),
139            cursor: 0,
140            last_synced: Default::default(),
141        }
142    }
143}
144
145#[derive(Debug)]
146struct GcThread<B: EventStreamBackend> {
147    sender: SyncSender<GcTask<B>>,
148}
149
150impl<B: EventStreamBackend> GcThread<B> {
151    fn new() -> GcThread<B> {
152        let (sender, recv) = std::sync::mpsc::sync_channel::<GcTask<B>>(8);
153
154        cubecl_environment::thread::spawn(move || {
155            while let Ok(event) = recv.recv() {
156                B::wait_event_sync(event.event).unwrap();
157                core::mem::drop(event.to_drop);
158            }
159        });
160
161        GcThread { sender }
162    }
163    fn register(&self, task: GcTask<B>) {
164        self.sender.send(task).unwrap()
165    }
166}
167
168fn stream_index(stream_id: &StreamId, max_streams: usize) -> usize {
169    stream_id.value as usize % max_streams
170}
171
172impl<'a, B: EventStreamBackend> ResolvedStreams<'a, B> {
173    /// Get the stream associated to the given [`stream_id`](StreamId).
174    pub fn get(&mut self, stream_id: &StreamId) -> &mut B::Stream {
175        let stream = self.streams.get_mut(stream_id);
176        &mut stream.stream
177    }
178
179    /// Get the stream associated to the [current `stream_id`](StreamId).
180    pub fn current(&mut self) -> &mut B::Stream {
181        let stream = self.streams.get_mut(&self.current);
182        &mut stream.stream
183    }
184
185    /// The current stream and the device's failure store together, for the
186    /// paths that hand the store to the stream's memory manager — every
187    /// reserve, bind and cleanup, since those are where slices shed the
188    /// failures they carry.
189    pub fn current_and_failures(&mut self) -> (&mut B::Stream, &mut ErrorGraph) {
190        let stream = self.streams.get_mut(&self.current);
191        (&mut stream.stream, self.failures)
192    }
193
194    /// [`current_and_failures`](Self::current_and_failures) for the stream at
195    /// `stream_id`.
196    pub fn get_and_failures(&mut self, stream_id: &StreamId) -> (&mut B::Stream, &mut ErrorGraph) {
197        let stream = self.streams.get_mut(stream_id);
198        (&mut stream.stream, self.failures)
199    }
200
201    /// Taint every allocation in `written` with `error` — see
202    /// [`base::taint`].
203    pub fn taint<'b>(
204        &mut self,
205        error: ServerError,
206        written: impl Iterator<Item = &'b BufferBinding>,
207    ) {
208        base::taint(self.streams, error, written, self.failures);
209    }
210
211    /// Release the failure on every allocation in `written` — see
212    /// [`base::written`].
213    pub fn written<'b>(&mut self, written: impl Iterator<Item = &'b BufferBinding>) {
214        base::written(self.streams, written, self.failures);
215    }
216
217    /// Enqueue a task to be cleaned.
218    pub fn gc(&mut self, gc: GcTask<B>) {
219        self.gc.sender.send(gc).unwrap();
220    }
221}
222
223impl<'a, B: EventStreamBackend> Drop for ResolvedStreams<'a, B> {
224    fn drop(&mut self) {
225        if self.analysis.pinned.is_empty() {
226            return;
227        }
228
229        let stream = self.streams.get_mut(&self.current);
230        let event_origin = B::flush(&mut stream.stream, self.failures);
231
232        let stream_gc = &mut unsafe { self.streams.get_special(0) }.stream;
233        B::wait_event(stream_gc, event_origin);
234        let event = B::flush(stream_gc, self.failures);
235
236        let pinned = core::mem::take(&mut self.analysis.pinned);
237        self.gc.register(GcTask::new(pinned, event));
238    }
239}
240
241impl<B: EventStreamBackend> MultiStream<B> {
242    /// Mutable access to the stream-creation backend, e.g. to change the
243    /// configuration new streams are created with. Already-created streams are
244    /// unaffected.
245    pub fn backend_mut(&mut self) -> &mut B {
246        &mut self.streams.factory_mut().backend
247    }
248
249    /// Creates an empty multi-stream.
250    pub fn new(logger: Arc<ServerLogger>, backend: B, max_streams: u8) -> Self {
251        let wrapper = EventStreamBackendWrapper { backend };
252        Self {
253            streams: StreamPool::new(wrapper, max_streams, 1),
254            failures: Failures::new(logger.clone()),
255            logger,
256            max_streams: max_streams as usize,
257            gc: GcThread::new(),
258            shared_bindings_pool: Vec::new(),
259        }
260    }
261
262    /// Synthetic [`StreamId`]s, one per initialized stream (see [`StreamPool::stream_ids`]).
263    pub fn stream_ids(&self) -> impl Iterator<Item = StreamId> + '_ {
264        self.streams.stream_ids()
265    }
266
267    /// Enqueue a task to be cleaned.
268    pub fn gc(&mut self, gc: GcTask<B>) {
269        self.gc.sender.send(gc).unwrap();
270    }
271
272    /// The backend stream on `stream_id`'s slot when that slot was ever
273    /// initialized, mutably — a lookup that must not create a stream (see
274    /// [`StreamPool::try_get_mut`]).
275    pub fn try_stream_mut(&mut self, stream_id: &StreamId) -> Option<&mut B::Stream> {
276        self.streams
277            .try_get_mut(stream_id)
278            .map(|wrapper| &mut wrapper.stream)
279    }
280
281    /// Resolves and returns a mutable reference to the stream for the given ID, performing any necessary
282    /// alignment based on the provided bindings.
283    ///
284    /// This method ensures that the stream is synchronized with any shared bindings from other streams
285    /// before returning the stream reference.
286    pub fn resolve<'a>(
287        &mut self,
288        stream_id: StreamId,
289        handles: impl Iterator<Item = &'a BufferBinding>,
290    ) -> ResolvedStreams<'_, B> {
291        let analysis = self.align_streams(stream_id, handles);
292
293        let stream = self.streams.get_mut(&stream_id);
294        stream.cursor += 1;
295
296        ResolvedStreams {
297            cursor: stream.cursor,
298            streams: &mut self.streams,
299            failures: self.failures.graph_mut(),
300            current: stream_id,
301            analysis,
302            gc: &self.gc,
303        }
304    }
305
306    /// Aligns the target stream with other streams based on shared bindings.
307    ///
308    /// This initializes the stream if it doesn't exist, analyzes which originating streams need flushing
309    /// for synchronization, flushes them, and waits on the events in the target stream.
310    fn align_streams<'a>(
311        &mut self,
312        stream_id: StreamId,
313        handles: impl Iterator<Item = &'a BufferBinding>,
314    ) -> SharedBindingAnalysis {
315        let analysis = self.update_shared_bindings(stream_id, handles);
316
317        self.apply_analysis(stream_id, analysis)
318    }
319
320    /// Updates and analyzes the bindings to determine which streams need alignment (flushing and waiting).
321    ///
322    /// This checks for shared bindings from other streams and determines if synchronization is needed
323    /// based on cursor positions.
324    pub(crate) fn update_shared_bindings<'a>(
325        &mut self,
326        stream_id: StreamId,
327        handles: impl Iterator<Item = &'a BufferBinding>,
328    ) -> SharedBindingAnalysis {
329        // We reset the memory pool for the info.
330        self.shared_bindings_pool.clear();
331
332        let mut analysis = SharedBindingAnalysis::default();
333
334        // We only consider handles whose stream is different from the current stream.
335        for handle in handles.filter(|handle| handle.stream != stream_id) {
336            let index = stream_index(&handle.stream, self.max_streams);
337            let stream = unsafe { self.streams.get_mut_index(index) };
338            let cursor_handle = B::handle_cursor(&stream.stream, handle);
339
340            self.shared_bindings_pool.push((
341                handle.memory.descriptor().id,
342                handle.stream,
343                cursor_handle,
344            ));
345            // Pinned unconditionally, even when the cursor check below decides
346            // no new wait is needed: the reverse-direction hazard (the origin
347            // stream freeing/reusing the memory under the in-flight consumer)
348            // exists either way.
349            analysis.pinned.push(handle.memory.clone());
350        }
351
352        let current = self.streams.get_mut(&stream_id);
353
354        for (handle_id, stream, cursor) in self.shared_bindings_pool.iter() {
355            let index = stream_index(stream, self.max_streams);
356
357            if let Some(last_synced) = current.last_synced.get(&index) {
358                if last_synced < cursor {
359                    self.logger.log_streaming(
360                        |level| matches!(level, StreamingLogLevel::Full),
361                        || {
362                            format!(
363                                "Binding on {} is shared on {} since it's not sync {} < {}",
364                                stream, stream_id, last_synced, cursor
365                            )
366                        },
367                    );
368                    analysis.shared(*handle_id, index);
369                }
370            } else {
371                self.logger.log_streaming(
372                    |level| matches!(level, StreamingLogLevel::Full),
373                    || {
374                        format!(
375                            "Binding on {} is shared on {} since it was never synced.",
376                            stream, stream_id,
377                        )
378                    },
379                );
380                analysis.shared(*handle_id, index);
381            }
382        }
383
384        analysis
385    }
386
387    pub(crate) fn apply_analysis(
388        &mut self,
389        stream_id: StreamId,
390        analysis: SharedBindingAnalysis,
391    ) -> SharedBindingAnalysis {
392        if analysis.slices.is_empty() {
393            return analysis;
394        }
395
396        let mut events = Vec::with_capacity(analysis.slices.len());
397
398        unsafe {
399            for origin in analysis.slices.keys() {
400                let stream = self.streams.get_mut_index(*origin);
401                let event = B::flush(&mut stream.stream, self.failures.graph_mut());
402
403                events.push(((origin, stream.cursor), event));
404            }
405        }
406
407        let stream = self.streams.get_mut(&stream_id);
408
409        for ((stream_origin, cursor_origin), event) in events {
410            stream.last_synced.insert(*stream_origin, cursor_origin);
411
412            self.logger.log_streaming(
413                |level| !matches!(level, StreamingLogLevel::Disabled),
414                || format!("Waiting on {stream_origin} from {stream_id}",),
415            );
416
417            B::wait_event(&mut stream.stream, event);
418        }
419
420        analysis
421    }
422}
423
424impl<B: EventStreamBackend> FailureStore for MultiStream<B> {
425    type Factory = EventStreamBackendWrapper<B>;
426
427    fn split(&mut self) -> (&mut StreamPool<Self::Factory>, &mut Failures) {
428        (&mut self.streams, &mut self.failures)
429    }
430
431    fn parts(&self) -> (&StreamPool<Self::Factory>, &Failures) {
432        (&self.streams, &self.failures)
433    }
434}
435
436impl<B: EventStreamBackend> core::fmt::Debug for StreamWrapper<B> {
437    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
438        f.debug_struct("StreamWrapper")
439            .field("stream", &self.stream)
440            .field("cursor", &self.cursor)
441            .field("last_synced", &self.last_synced)
442            .finish()
443    }
444}
445
446#[derive(Default, Debug)]
447pub(crate) struct SharedBindingAnalysis {
448    slices: HashMap<usize, Vec<ManagedMemoryId>>,
449    /// Every cross-stream binding of the task, kept alive until the consumer
450    /// stream's work completes (released by the GC thread after its event).
451    ///
452    /// The origin stream's pools consider a slice free once no handle/binding
453    /// references its descriptor, but the consumer's kernel may still be
454    /// running on the GPU after the CPU-side bindings were dropped at enqueue
455    /// time. Pinning the bindings here is what keeps the slice non-free until
456    /// the GC event fires, so `cleanup`/`try_reserve` on the origin stream
457    /// cannot dealloc or reuse memory that is still read by another stream.
458    pinned: SharedMemoryBindings,
459}
460
461/// Equality covers the sync analysis only; `pinned` is a lifetime mechanism,
462/// not part of the analysis result.
463impl PartialEq for SharedBindingAnalysis {
464    fn eq(&self, other: &Self) -> bool {
465        self.slices == other.slices
466    }
467}
468
469impl Eq for SharedBindingAnalysis {}
470
471impl SharedBindingAnalysis {
472    fn shared(&mut self, id: ManagedMemoryId, index: usize) {
473        match self.slices.get_mut(&index) {
474            Some(bindings) => bindings.push(id),
475            None => {
476                self.slices.insert(index, alloc::vec![id]);
477            }
478        }
479    }
480}
481
482#[cfg(test)]
483mod tests {
484    use crate::server::Handle;
485    use core::sync::atomic::{AtomicBool, Ordering};
486
487    use super::*;
488
489    /// A service for handles that never reach a device.
490    fn service() -> cubecl_common::device::ServiceId {
491        cubecl_common::device::ServiceId::of::<()>(cubecl_common::device::DeviceId::new(0, 0))
492    }
493
494    const MAX_STREAMS: u8 = 4;
495
496    #[test_log::test]
497    fn test_analysis_shared_bindings_1() {
498        let logger = Arc::new(ServerLogger::default());
499        let stream_1 = StreamId { value: 1 };
500        let stream_2 = StreamId { value: 2 };
501
502        let binding_1 = handle(stream_1);
503        let binding_2 = handle(stream_2);
504
505        let mut ms = MultiStream::new(logger, TestBackend, MAX_STREAMS);
506        ms.resolve(stream_1, [].into_iter());
507        ms.resolve(stream_2, [].into_iter());
508
509        let analysis = ms.update_shared_bindings(stream_1, [&binding_1, &binding_2].into_iter());
510
511        let mut expected = SharedBindingAnalysis::default();
512        expected.shared(
513            binding_2.memory.descriptor().id,
514            ms.streams.stream_index(&binding_2.stream),
515        );
516
517        assert_eq!(analysis, expected);
518    }
519
520    #[test_log::test]
521    fn test_analysis_shared_bindings_2() {
522        let logger = Arc::new(ServerLogger::default());
523        let stream_1 = StreamId { value: 1 };
524        let stream_2 = StreamId { value: 2 };
525
526        let binding_1 = handle(stream_1);
527        let binding_2 = handle(stream_2);
528        let binding_3 = handle(stream_1);
529
530        let mut ms = MultiStream::new(logger, TestBackend, 4);
531        ms.resolve(stream_1, [].into_iter());
532        ms.resolve(stream_2, [].into_iter());
533
534        let analysis =
535            ms.update_shared_bindings(stream_1, [&binding_1, &binding_2, &binding_3].into_iter());
536
537        let mut expected = SharedBindingAnalysis::default();
538        expected.shared(
539            binding_2.memory.descriptor().id,
540            ms.streams.stream_index(&binding_2.stream),
541        );
542
543        assert_eq!(analysis, expected);
544    }
545
546    #[test_log::test]
547    fn test_analysis_no_shared() {
548        let logger = Arc::new(ServerLogger::default());
549        let stream_1 = StreamId { value: 1 };
550        let stream_2 = StreamId { value: 2 };
551
552        let binding_1 = handle(stream_1);
553        let binding_2 = handle(stream_1);
554        let binding_3 = handle(stream_1);
555
556        let mut ms = MultiStream::new(logger, TestBackend, MAX_STREAMS);
557        ms.resolve(stream_1, [].into_iter());
558        ms.resolve(stream_2, [].into_iter());
559
560        let analysis =
561            ms.update_shared_bindings(stream_1, [&binding_1, &binding_2, &binding_3].into_iter());
562
563        let expected = SharedBindingAnalysis::default();
564
565        assert_eq!(analysis, expected);
566    }
567
568    #[test_log::test]
569    fn test_state() {
570        let logger = Arc::new(ServerLogger::default());
571        let stream_1 = StreamId { value: 1 };
572        let stream_2 = StreamId { value: 2 };
573
574        let binding_1 = handle(stream_1);
575        let binding_2 = handle(stream_2);
576        let binding_3 = handle(stream_1);
577
578        let mut ms = MultiStream::new(logger, TestBackend, MAX_STREAMS);
579        ms.resolve(stream_1, [].into_iter());
580        ms.resolve(stream_2, [].into_iter());
581
582        ms.resolve(stream_1, [&binding_1, &binding_2, &binding_3].into_iter());
583
584        let stream1 = ms.streams.get_mut(&stream_1);
585        let index_2 = stream_index(&stream_2, MAX_STREAMS as usize);
586        assert_eq!(stream1.last_synced.get(&index_2), Some(&1));
587        assert_eq!(stream1.cursor, 2);
588
589        let stream2 = ms.streams.get_mut(&stream_2);
590        assert!(stream2.last_synced.is_empty());
591        assert_eq!(stream2.cursor, 1);
592    }
593
594    #[test_log::test]
595    fn test_cross_stream_binding_pinned_until_gc_event() {
596        let logger = Arc::new(ServerLogger::default());
597        let stream_1 = StreamId { value: 1 };
598        let stream_2 = StreamId { value: 2 };
599
600        let gate = Arc::new(AtomicBool::new(false));
601        let mut ms = MultiStream::new(logger, GatedBackend { gate: gate.clone() }, MAX_STREAMS);
602        ms.resolve(stream_1, [].into_iter());
603        ms.resolve(stream_2, [].into_iter());
604
605        let handle = Handle::new(service(), stream_1, 10);
606        let observer = handle.memory.clone();
607        let binding = handle.binding();
608
609        drop(ms.resolve(stream_2, [&binding].into_iter()));
610        drop(binding);
611
612        // The GC thread is blocked on the (gated) consumer event, so the pinned
613        // binding must keep the memory non-free even though every user-side
614        // handle/binding is gone.
615        assert!(
616            !observer.is_free(),
617            "cross-stream binding must stay pinned while the consumer event is pending"
618        );
619
620        gate.store(true, Ordering::Release);
621        wait_until_free(&observer);
622    }
623
624    #[test_log::test]
625    fn test_already_synced_cross_stream_binding_still_pinned() {
626        let logger = Arc::new(ServerLogger::default());
627        let stream_1 = StreamId { value: 1 };
628        let stream_2 = StreamId { value: 2 };
629
630        let gate = Arc::new(AtomicBool::new(true));
631        let mut ms = MultiStream::new(logger, GatedBackend { gate: gate.clone() }, MAX_STREAMS);
632        ms.resolve(stream_1, [].into_iter());
633        ms.resolve(stream_2, [].into_iter());
634
635        // First resolve records stream_1 as synced on stream_2.
636        let handle_1 = Handle::new(service(), stream_1, 10);
637        let binding_1 = handle_1.binding();
638        drop(ms.resolve(stream_2, [&binding_1].into_iter()));
639        drop(binding_1);
640
641        // Close the gate for the second round.
642        gate.store(false, Ordering::Release);
643
644        let handle_2 = Handle::new(service(), stream_1, 10);
645        let observer = handle_2.memory.clone();
646        let binding_2 = handle_2.binding();
647
648        // stream_2 already synced past this binding's cursor, so the sync
649        // analysis is empty — but the binding must still be pinned: the origin
650        // stream could otherwise free/reuse the memory under the in-flight
651        // consumer.
652        let resolved = ms.resolve(stream_2, [&binding_2].into_iter());
653        assert!(resolved.analysis.slices.is_empty());
654        drop(resolved);
655        drop(binding_2);
656
657        assert!(
658            !observer.is_free(),
659            "already-synced cross-stream binding must still be pinned"
660        );
661
662        gate.store(true, Ordering::Release);
663        wait_until_free(&observer);
664    }
665
666    fn wait_until_free(observer: &crate::memory_management::ManagedMemoryHandle) {
667        let start = std::time::Instant::now();
668        while !observer.is_free() {
669            assert!(
670                start.elapsed() < std::time::Duration::from_secs(10),
671                "pinned binding was never released"
672            );
673            std::thread::yield_now();
674        }
675    }
676
677    fn handle(stream: StreamId) -> BufferBinding {
678        Handle::new(service(), stream, 10).binding()
679    }
680
681    struct TestBackend;
682
683    #[derive(Debug, Default)]
684    struct TestStream;
685
686    #[derive(Debug)]
687    struct TestEvent {}
688
689    /// A backend whose events complete only once the shared `gate` opens,
690    /// emulating GPU work still in flight on the consumer stream.
691    struct GatedBackend {
692        gate: Arc<AtomicBool>,
693    }
694
695    #[derive(Debug, Default)]
696    struct GatedStream {
697        gate: Arc<AtomicBool>,
698    }
699
700    #[derive(Debug)]
701    struct GatedEvent {
702        gate: Arc<AtomicBool>,
703    }
704
705    /// The test streams manage no memory, so nothing carries a failure.
706    impl StreamMemory for GatedStream {
707        fn failure(&self, _binding: &BufferBinding) -> Option<FailureId> {
708            None
709        }
710
711        fn taint(
712            &mut self,
713            _binding: &BufferBinding,
714            _failure: FailureId,
715            _failures: &mut ErrorGraph,
716        ) {
717        }
718
719        fn written(&mut self, _binding: &BufferBinding, _failures: &mut ErrorGraph) {}
720    }
721
722    impl StreamMemory for TestStream {
723        fn failure(&self, _binding: &BufferBinding) -> Option<FailureId> {
724            None
725        }
726
727        fn taint(
728            &mut self,
729            _binding: &BufferBinding,
730            _failure: FailureId,
731            _failures: &mut ErrorGraph,
732        ) {
733        }
734
735        fn written(&mut self, _binding: &BufferBinding, _failures: &mut ErrorGraph) {}
736    }
737
738    impl EventStreamBackend for GatedBackend {
739        type Stream = GatedStream;
740        type Event = GatedEvent;
741
742        fn create_stream(&self) -> Self::Stream {
743            GatedStream {
744                gate: self.gate.clone(),
745            }
746        }
747
748        fn flush(stream: &mut Self::Stream, _failures: &mut ErrorGraph) -> Self::Event {
749            GatedEvent {
750                gate: stream.gate.clone(),
751            }
752        }
753
754        fn wait_event(_stream: &mut Self::Stream, _event: Self::Event) {}
755
756        fn wait_event_sync(event: Self::Event) -> Result<(), ServerError> {
757            while !event.gate.load(Ordering::Acquire) {
758                std::thread::yield_now();
759            }
760            Ok(())
761        }
762
763        fn handle_cursor(_stream: &Self::Stream, _handle: &BufferBinding) -> u64 {
764            0
765        }
766    }
767
768    impl EventStreamBackend for TestBackend {
769        type Stream = TestStream;
770        type Event = TestEvent;
771
772        fn create_stream(&self) -> Self::Stream {
773            TestStream
774        }
775
776        fn flush(_stream: &mut Self::Stream, _failures: &mut ErrorGraph) -> Self::Event {
777            TestEvent {}
778        }
779
780        fn wait_event(_stream: &mut Self::Stream, _event: Self::Event) {}
781
782        fn wait_event_sync(_event: Self::Event) -> Result<(), ServerError> {
783            Ok(())
784        }
785
786        fn handle_cursor(_stream: &Self::Stream, _handle: &BufferBinding) -> u64 {
787            0
788        }
789    }
790}