Skip to main content

commonware_runtime/
mocks.rs

1//! Mock implementations of runtime primitives for testing.
2
3use crate::{
4    Blob, BlobVersion, BufMut, Clock, Error, Handle, IoBufs, IoBufsMut, Metrics, Name, ReadOptions,
5    Spawner, Supervisor, WriteOptions,
6    signal::Signal,
7    telemetry::metrics::{Metric, Registered},
8};
9use bytes::{Bytes, BytesMut};
10use commonware_macros::boxed;
11use commonware_utils::{
12    channel::{fallible::OneshotExt, oneshot},
13    sync::Mutex,
14};
15use governor::clock::{Clock as GovernorClock, ReasonablyRealtime};
16use rand::{TryCryptoRng, TryRng};
17use std::{
18    future::{Future, poll_fn},
19    mem,
20    sync::Arc,
21    task::Poll,
22};
23
24cfg_if::cfg_if! {
25    if #[cfg(any(test, feature = "test-utils"))] {
26        use crate::{IoBufMut, utils::reschedule};
27        use futures::poll;
28        use std::sync::atomic::{AtomicUsize, Ordering};
29    }
30}
31
32/// In-memory storage with exclusive logical opens and durable snapshot inspection.
33#[cfg(any(test, feature = "test-utils"))]
34#[derive(Clone, Default)]
35pub struct Storage {
36    inner: crate::storage::memory::Storage,
37    opens: Arc<crate::storage::memory::open::Opens>,
38}
39
40#[cfg(any(test, feature = "test-utils"))]
41impl Storage {
42    /// Create an empty memory storage backend.
43    pub fn new() -> Self {
44        Self::default()
45    }
46
47    /// Compute a CRC32C digest of all durable blob contents.
48    pub fn audit(&self) -> [u8; 4] {
49        self.inner.audit()
50    }
51
52    /// Return a copy of a blob's durable raw contents without interpreting its container header.
53    pub fn raw_blob(&self, partition: &str, name: &[u8]) -> Option<Vec<u8>> {
54        self.inner.raw_blob(partition, name)
55    }
56
57    /// Return a copy of a blob's durable logical contents, or `None` when the blob is missing or
58    /// its container header does not resolve.
59    pub fn logical_blob(&self, partition: &str, name: &[u8]) -> Option<Vec<u8>> {
60        self.inner.logical_blob(partition, name)
61    }
62
63    /// Install durable raw contents without validating the blob's container header.
64    ///
65    /// This retires the prior incarnation and permits a new open while old handles remain alive.
66    pub fn set_raw_blob(&self, partition: &str, name: &[u8], content: Vec<u8>) {
67        self.opens
68            .replace(partition, Some(name), || {
69                self.inner.set_raw_blob(partition, name, content);
70                Ok(())
71            })
72            .expect("installing a raw memory image cannot fail");
73    }
74}
75
76#[cfg(any(test, feature = "test-utils"))]
77impl crate::Storage for Storage {
78    type Blob = crate::storage::memory::open::Blob<crate::storage::memory::Blob>;
79
80    async fn open_versioned(
81        &self,
82        partition: &str,
83        name: &[u8],
84        versions: std::ops::RangeInclusive<BlobVersion>,
85    ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
86        let opened = self.opens.open(
87            partition,
88            name,
89            self.inner.open_versioned(partition, name, versions),
90        )?;
91        Ok(opened.finish())
92    }
93
94    async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
95        self.opens
96            .remove(partition, name, self.inner.remove(partition, name))
97    }
98
99    async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
100        self.inner.scan(partition).await
101    }
102}
103
104/// Default buffer size (64 KB). Controls both how much data the stream
105/// pulls per recv and the backpressure threshold for send.
106const DEFAULT_BUFFER_SIZE: usize = 64 * 1024;
107
108/// A mock channel struct that is used internally by Sink and Stream.
109pub struct Channel {
110    /// Stores the bytes sent by the sink that are not yet read by the stream.
111    buffer: BytesMut,
112
113    /// If the stream is waiting to read bytes, the waiter stores the number of
114    /// bytes that the stream is waiting for, as well as the oneshot sender that
115    /// the sink uses to send the bytes to the stream directly.
116    waiter: Option<(usize, oneshot::Sender<Bytes>)>,
117
118    /// Target buffer size, used to bound both the stream's local buffer
119    /// and the shared buffer (backpressure threshold).
120    buffer_size: usize,
121
122    /// If the sink is blocked waiting for the buffer to drain, this holds
123    /// the oneshot sender that the stream uses to wake the sink.
124    drain_waiter: Option<oneshot::Sender<()>>,
125
126    /// Tracks whether the sink is still alive and able to send messages.
127    sink_alive: bool,
128
129    /// Tracks whether the stream is still alive and able to receive messages.
130    stream_alive: bool,
131}
132
133impl Channel {
134    /// Returns an async-safe Sink/Stream pair with default buffer size.
135    pub fn init() -> (Sink, Stream) {
136        Self::init_with_buffer_size(DEFAULT_BUFFER_SIZE)
137    }
138
139    /// Returns an async-safe Sink/Stream pair with the specified buffer size.
140    pub fn init_with_buffer_size(buffer_size: usize) -> (Sink, Stream) {
141        let channel = Arc::new(Mutex::new(Self {
142            buffer: BytesMut::new(),
143            waiter: None,
144            buffer_size,
145            drain_waiter: None,
146            sink_alive: true,
147            stream_alive: true,
148        }));
149        (
150            Sink {
151                channel: channel.clone(),
152                state: SinkState::Open,
153            },
154            Stream {
155                channel,
156                buffer: BytesMut::new(),
157                poisoned: false,
158            },
159        )
160    }
161
162    /// Restores bytes that were detached from the front of the shared buffer.
163    fn restore_front(&mut self, data: Bytes) {
164        if data.is_empty() {
165            return;
166        }
167
168        let mut restored = BytesMut::with_capacity(data.len() + self.buffer.len());
169        restored.extend_from_slice(&data);
170        restored.extend_from_slice(&self.buffer);
171        self.buffer = restored;
172    }
173
174    /// Marks the sink as closed and wakes any waiter.
175    fn close_sink(&mut self) {
176        self.sink_alive = false;
177
178        // If there is a waiter, resolve it by dropping the oneshot sender.
179        self.waiter.take();
180    }
181}
182
183struct RecvWaiterGuard {
184    channel: Arc<Mutex<Channel>>,
185    active: bool,
186}
187
188impl RecvWaiterGuard {
189    const fn new(channel: Arc<Mutex<Channel>>) -> Self {
190        Self {
191            channel,
192            active: true,
193        }
194    }
195
196    const fn disarm(&mut self) {
197        self.active = false;
198    }
199}
200
201impl Drop for RecvWaiterGuard {
202    fn drop(&mut self) {
203        if !self.active {
204            return;
205        }
206
207        self.channel.lock().waiter.take();
208    }
209}
210
211/// A mock sink that implements the Sink trait.
212pub struct Sink {
213    channel: Arc<Mutex<Channel>>,
214    state: SinkState,
215}
216
217/// Lifecycle state for the mock sink half.
218enum SinkState {
219    /// Sends may be attempted.
220    Open,
221    /// A send is currently in progress.
222    Sending,
223    /// The sink has been closed.
224    Closed,
225}
226
227impl Sink {
228    fn close(&mut self) {
229        if matches!(self.state, SinkState::Closed) {
230            return;
231        }
232        self.channel.lock().close_sink();
233        self.state = SinkState::Closed;
234    }
235}
236
237impl crate::Sink for Sink {
238    async fn send(&mut self, bufs: impl Into<IoBufs> + Send) -> Result<(), Error> {
239        match self.state {
240            SinkState::Open => {}
241            SinkState::Sending => {
242                self.close();
243                return Err(Error::Closed);
244            }
245            SinkState::Closed => return Err(Error::Closed),
246        }
247
248        let drain_recv = {
249            let mut channel = self.channel.lock();
250
251            // If the receiver is dead, we cannot send any more messages.
252            if !channel.stream_alive {
253                channel.close_sink();
254                self.state = SinkState::Closed;
255                return Err(Error::SendFailed);
256            }
257
258            channel.buffer.put(bufs.into());
259
260            // If there is a waiter and the buffer is large enough,
261            // resolve the waiter (while clearing the waiter field).
262            if channel
263                .waiter
264                .as_ref()
265                .is_some_and(|(requested, _)| *requested <= channel.buffer.len())
266            {
267                // Send up to buffer_size bytes (but at least requested amount)
268                let (requested, os_send) = channel.waiter.take().unwrap();
269                let send_amount = channel.buffer.len().min(requested.max(channel.buffer_size));
270                let data = channel.buffer.split_to(send_amount).freeze();
271
272                // A canceled recv should behave like a buffered transport:
273                // preserve the bytes and allow a subsequent recv to consume them.
274                if let Err(data) = os_send.send(data) {
275                    channel.restore_front(data);
276                    if !channel.stream_alive {
277                        channel.close_sink();
278                        self.state = SinkState::Closed;
279                        return Err(Error::SendFailed);
280                    }
281                }
282            }
283
284            // If the buffer exceeds the write limit, block until the
285            // receiver drains enough data.
286            if channel.buffer.len() > channel.buffer_size {
287                assert!(channel.drain_waiter.is_none());
288                let (os_send, os_recv) = oneshot::channel();
289                channel.drain_waiter = Some(os_send);
290                os_recv
291            } else {
292                return Ok(());
293            }
294        };
295
296        // Mark the sink as sending before awaiting so cancellation can be
297        // detected by the next send.
298        self.state = SinkState::Sending;
299
300        // Wait for the receiver to drain the buffer.
301        match drain_recv.await {
302            Ok(()) => {
303                self.state = SinkState::Open;
304                Ok(())
305            }
306            Err(_) => {
307                self.close();
308                Err(Error::SendFailed)
309            }
310        }
311    }
312}
313
314impl Drop for Sink {
315    fn drop(&mut self) {
316        self.close();
317    }
318}
319
320/// A mock stream that implements the Stream trait.
321pub struct Stream {
322    channel: Arc<Mutex<Channel>>,
323    /// Local buffer for data that has been received but not yet consumed.
324    buffer: BytesMut,
325    poisoned: bool,
326}
327
328impl crate::Stream for Stream {
329    async fn recv(&mut self, len: usize) -> Result<IoBufs, Error> {
330        if self.poisoned {
331            return Err(Error::Closed);
332        }
333
334        let os_recv = {
335            let mut channel = self.channel.lock();
336
337            // Pull data from channel buffer into local buffer.
338            let target = len.max(channel.buffer_size);
339            let pull_amount = channel
340                .buffer
341                .len()
342                .min(target.saturating_sub(self.buffer.len()));
343            if pull_amount > 0 {
344                let data = channel.buffer.split_to(pull_amount);
345                self.buffer.extend_from_slice(&data);
346
347                // Wake a blocked sender if the buffer drained below the limit.
348                if channel.buffer.len() <= channel.buffer_size
349                    && let Some(sender) = channel.drain_waiter.take()
350                {
351                    sender.send_lossy(());
352                }
353            }
354
355            // If we have enough, return immediately.
356            if self.buffer.len() >= len {
357                return Ok(IoBufs::from(self.buffer.split_to(len).freeze()));
358            }
359
360            // If the sink is dead, we cannot receive any more messages.
361            if !channel.sink_alive {
362                self.poisoned = true;
363                return Err(Error::RecvFailed);
364            }
365
366            // Set up waiter for remaining amount.
367            let remaining = len - self.buffer.len();
368            assert!(channel.waiter.is_none());
369            let (os_send, os_recv) = oneshot::channel();
370            channel.waiter = Some((remaining, os_send));
371            os_recv
372        };
373
374        let mut waiter_guard = RecvWaiterGuard::new(self.channel.clone());
375
376        // Pre-poison so that cancellation  leaves the stream permanently closed.
377        self.poisoned = true;
378
379        // Wait for the waiter to be resolved.
380        let data = match os_recv.await {
381            Ok(data) => {
382                waiter_guard.disarm();
383                self.poisoned = false;
384                data
385            }
386            Err(_) => {
387                waiter_guard.disarm();
388                return Err(Error::RecvFailed);
389            }
390        };
391        self.buffer.extend_from_slice(&data);
392
393        assert!(self.buffer.len() >= len);
394        Ok(IoBufs::from(self.buffer.split_to(len).freeze()))
395    }
396
397    fn peek(&self, max_len: usize) -> &[u8] {
398        let len = max_len.min(self.buffer.len());
399        &self.buffer[..len]
400    }
401}
402
403impl Drop for Stream {
404    fn drop(&mut self) {
405        let mut channel = self.channel.lock();
406        channel.stream_alive = false;
407
408        // Wake a blocked sender so it can observe the closed stream.
409        channel.drain_waiter.take();
410    }
411}
412
413/// A sync deferred by a [DelayedSyncBlob], held open until explicitly completed.
414pub struct DeferredSync {
415    /// Completes the sync with the provided result (success runs the inner blob's sync).
416    pub release: oneshot::Sender<Result<(), Error>>,
417
418    /// Resolves once the deferred sync's handle begins waiting on `release`.
419    pub blocked: oneshot::Receiver<()>,
420}
421
422/// Coordinates durability operations for a [DelayedSyncContext] or [DelayedSyncBlob].
423///
424/// Every started sync parks in a deferred queue (in start order) until a test
425/// releases it or [Self::unblock] runs. [Self::arm] additionally installs a one-shot gate that blocks
426/// the next durability operation and counts operations from that point on
427/// ([Self::calls]). The gate is pushed onto the deferred queue when [Self::arm]
428/// is called, before any operation reaches it.
429#[derive(Clone, Default)]
430pub struct PendingSyncs {
431    state: Arc<Mutex<State>>,
432}
433
434/// State shared by all clones of a [PendingSyncs].
435#[derive(Default)]
436struct State {
437    /// Deferred syncs in start order.
438    syncs: Vec<DeferredSync>,
439    /// One-shot gate blocking the next durability operation (see [PendingSyncs::arm]).
440    gate: SyncGateState,
441    /// Sticky: stop parking started syncs (see [PendingSyncs::unblock]).
442    unblocked: bool,
443    /// Sticky: syncs resolve to an injected error (see [PendingSyncs::arm_fail]).
444    fail: bool,
445    /// Started syncs issued.
446    starts: usize,
447    /// Started syncs whose completion futures have begun executing.
448    entered: usize,
449    /// Started syncs that completed durably.
450    completions: usize,
451}
452
453impl State {
454    /// Creates a waiter parked in the deferred queue.
455    fn defer(&mut self) -> SyncWaiter {
456        let (release, release_rx) = oneshot::channel();
457        let (entered, blocked) = oneshot::channel();
458        self.syncs.push(DeferredSync { release, blocked });
459        SyncWaiter {
460            entered,
461            release: release_rx,
462        }
463    }
464
465    /// Records a durability operation if the gate is armed, returning the
466    /// one-shot gate waiter if it has not been consumed yet.
467    const fn observe(&mut self) -> Option<SyncWaiter> {
468        if !self.gate.tracking {
469            return None;
470        }
471        self.gate.calls += 1;
472        self.gate.waiter.take()
473    }
474
475    /// Parks a new deferred sync, unless [PendingSyncs::unblock] already ran.
476    fn park(&mut self) -> Option<SyncWaiter> {
477        if self.unblocked {
478            return None;
479        }
480        Some(self.defer())
481    }
482}
483
484/// Forwards [Supervisor], [Clock], [GovernorClock], [ReasonablyRealtime],
485/// [Metrics], [TryRng], and [TryCryptoRng] to the wrapped context for test
486/// context wrappers with one extra field (named by the second argument).
487macro_rules! forward_context {
488    ($wrapper:ident, $field:ident) => {
489        impl<E: Supervisor> Supervisor for $wrapper<E> {
490            fn name(&self) -> Name {
491                self.inner.name()
492            }
493
494            fn child(&self, label: &'static str) -> Self {
495                Self {
496                    inner: self.inner.child(label),
497                    $field: self.$field.clone(),
498                }
499            }
500
501            fn with_attribute(self, key: &'static str, value: impl std::fmt::Display) -> Self {
502                Self {
503                    inner: self.inner.with_attribute(key, value),
504                    $field: self.$field,
505                }
506            }
507        }
508
509        impl<E: Clock> Clock for $wrapper<E> {
510            fn current(&self) -> std::time::SystemTime {
511                self.inner.current()
512            }
513
514            fn sleep(
515                &self,
516                duration: std::time::Duration,
517            ) -> impl Future<Output = ()> + Send + 'static + use<E> {
518                self.inner.sleep(duration)
519            }
520
521            fn sleep_until(
522                &self,
523                deadline: std::time::SystemTime,
524            ) -> impl Future<Output = ()> + Send + 'static + use<E> {
525                self.inner.sleep_until(deadline)
526            }
527        }
528
529        impl<E: Clock> GovernorClock for $wrapper<E> {
530            type Instant = std::time::SystemTime;
531
532            fn now(&self) -> Self::Instant {
533                self.current()
534            }
535        }
536
537        impl<E: Clock> ReasonablyRealtime for $wrapper<E> {}
538
539        impl<E: Metrics> Metrics for $wrapper<E> {
540            fn register<N: Into<String>, H: Into<String>, M: Metric>(
541                &self,
542                name: N,
543                help: H,
544                metric: M,
545            ) -> Registered<M> {
546                self.inner.register(name, help, metric)
547            }
548
549            fn encode(&self) -> String {
550                self.inner.encode()
551            }
552        }
553
554        impl<E: TryRng> TryRng for $wrapper<E> {
555            type Error = E::Error;
556
557            fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
558                self.inner.try_next_u32()
559            }
560
561            fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
562                self.inner.try_next_u64()
563            }
564
565            fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
566                self.inner.try_fill_bytes(dest)
567            }
568        }
569
570        impl<E: TryCryptoRng> TryCryptoRng for $wrapper<E> {}
571    };
572}
573
574/// Snapshot of the options observed by a [RecordingContext] or [RecordingBlob].
575#[cfg(any(test, feature = "test-utils"))]
576#[derive(Clone, Debug, Default, Eq, PartialEq)]
577pub struct RecordingSnapshot {
578    /// Options supplied to read operations, in call order, with one entry per range of a batched
579    /// read.
580    pub reads: Vec<ReadOptions>,
581    /// Range count of each batched read, in call order.
582    pub batches: Vec<usize>,
583    /// Options supplied to write operations, in call order.
584    pub writes: Vec<WriteOptions>,
585}
586
587/// Shared observations produced by recording storage wrappers.
588#[cfg(any(test, feature = "test-utils"))]
589#[derive(Clone, Default)]
590pub struct Recordings {
591    state: Arc<Mutex<RecordingSnapshot>>,
592}
593
594#[cfg(any(test, feature = "test-utils"))]
595impl Recordings {
596    /// Return a snapshot of all observations recorded so far.
597    pub fn snapshot(&self) -> RecordingSnapshot {
598        self.state.lock().clone()
599    }
600
601    /// Remove all recorded observations.
602    pub fn clear(&self) {
603        *self.state.lock() = RecordingSnapshot::default();
604    }
605
606    fn read(&self, options: ReadOptions) {
607        self.state.lock().reads.push(options);
608    }
609
610    fn batch(&self, ranges: usize, options: ReadOptions) {
611        let mut state = self.state.lock();
612        state.reads.extend(std::iter::repeat_n(options, ranges));
613        state.batches.push(ranges);
614    }
615
616    fn write(&self, options: WriteOptions) {
617        self.state.lock().writes.push(options);
618    }
619}
620
621/// Context wrapper that records options supplied to every opened blob.
622#[cfg(any(test, feature = "test-utils"))]
623#[derive(Clone)]
624pub struct RecordingContext<E> {
625    /// Wrapped context.
626    pub inner: E,
627    /// Observations shared by this context and all blobs opened through it.
628    pub recordings: Recordings,
629}
630
631#[cfg(any(test, feature = "test-utils"))]
632impl<E> RecordingContext<E> {
633    /// Wrap `inner` and return both the context and its shared observations.
634    pub fn new(inner: E) -> (Self, Recordings) {
635        let recordings = Recordings::default();
636        (
637            Self {
638                inner,
639                recordings: recordings.clone(),
640            },
641            recordings,
642        )
643    }
644}
645
646#[cfg(any(test, feature = "test-utils"))]
647forward_context!(RecordingContext, recordings);
648
649#[cfg(any(test, feature = "test-utils"))]
650impl<E: Spawner> Spawner for RecordingContext<E> {
651    fn shared(mut self, blocking: bool) -> Self {
652        self.inner = self.inner.shared(blocking);
653        self
654    }
655
656    fn dedicated(mut self) -> Self {
657        self.inner = self.inner.dedicated();
658        self
659    }
660
661    fn spawn<F, Fut, T>(self, f: F) -> Handle<T>
662    where
663        F: FnOnce(Self) -> Fut + Send + 'static,
664        Fut: Future<Output = T> + Send + 'static,
665        T: Send + 'static,
666    {
667        let recordings = self.recordings;
668        self.inner.spawn(move |inner| f(Self { inner, recordings }))
669    }
670
671    async fn stop(self, value: i32, timeout: Option<std::time::Duration>) -> Result<(), Error> {
672        self.inner.stop(value, timeout).await
673    }
674
675    fn stopped(&self) -> Signal {
676        self.inner.stopped()
677    }
678}
679
680#[cfg(any(test, feature = "test-utils"))]
681impl<E: crate::Storage> crate::Storage for RecordingContext<E> {
682    type Blob = RecordingBlob<E::Blob>;
683
684    async fn open_versioned(
685        &self,
686        partition: &str,
687        name: &[u8],
688        versions: std::ops::RangeInclusive<BlobVersion>,
689    ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
690        let (inner, len, version) = self.inner.open_versioned(partition, name, versions).await?;
691        Ok((
692            RecordingBlob {
693                inner,
694                recordings: self.recordings.clone(),
695            },
696            len,
697            version,
698        ))
699    }
700
701    async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
702        self.inner.remove(partition, name).await
703    }
704
705    async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
706        self.inner.scan(partition).await
707    }
708}
709
710/// Blob wrapper that records read and write options and the range count of each batched read,
711/// then delegates each operation.
712#[cfg(any(test, feature = "test-utils"))]
713pub struct RecordingBlob<B> {
714    inner: B,
715    recordings: Recordings,
716}
717
718#[cfg(any(test, feature = "test-utils"))]
719impl<B: Blob> Blob for RecordingBlob<B> {
720    async fn read_at_buf(
721        &self,
722        offset: u64,
723        len: usize,
724        bufs: impl Into<IoBufsMut> + Send,
725        options: ReadOptions,
726    ) -> Result<IoBufsMut, Error> {
727        self.recordings.read(options);
728        self.inner.read_at_buf(offset, len, bufs, options).await
729    }
730
731    async fn read_at(
732        &self,
733        offset: u64,
734        len: usize,
735        options: ReadOptions,
736    ) -> Result<IoBufsMut, Error> {
737        self.recordings.read(options);
738        self.inner.read_at(offset, len, options).await
739    }
740
741    fn read_many(
742        &self,
743        ranges: &[(u64, usize)],
744        options: ReadOptions,
745    ) -> impl futures::Stream<Item = Result<(usize, IoBufsMut), Error>> + Send {
746        self.recordings.batch(ranges.len(), options);
747        self.inner.read_many(ranges, options)
748    }
749
750    async fn write_at(
751        &self,
752        offset: u64,
753        bufs: impl Into<IoBufs> + Send,
754        options: WriteOptions,
755    ) -> Result<(), Error> {
756        self.recordings.write(options);
757        self.inner.write_at(offset, bufs, options).await
758    }
759
760    async fn resize(&self, len: u64) -> Result<(), Error> {
761        self.inner.resize(len).await
762    }
763
764    async fn sync(&self) -> Result<(), Error> {
765        self.inner.sync().await
766    }
767
768    async fn start_sync(&self) -> Handle<()> {
769        self.inner.start_sync().await
770    }
771}
772
773/// Blob wrapper that yields after the first backend poll of each read, allowing a test to
774/// hand the unresolved read future to another task even when the backend completes immediately.
775#[cfg(any(test, feature = "test-utils"))]
776pub struct MigratingReadBlob<B> {
777    /// Wrapped blob.
778    inner: B,
779    /// Whether each read must remain pending after its first backend poll.
780    require_pending: bool,
781    /// Reads started through this blob.
782    reads: AtomicUsize,
783}
784
785#[cfg(any(test, feature = "test-utils"))]
786impl<B> MigratingReadBlob<B> {
787    /// Wrap `inner`, optionally requiring the first backend poll of each read to return pending.
788    /// Set `require_pending` when a test must exercise registered backend I/O.
789    pub const fn new(inner: B, require_pending: bool) -> Self {
790        Self {
791            inner,
792            require_pending,
793            reads: AtomicUsize::new(0),
794        }
795    }
796
797    /// Number of reads started through this blob.
798    pub fn reads(&self) -> usize {
799        self.reads.load(Ordering::Relaxed)
800    }
801}
802
803#[cfg(any(test, feature = "test-utils"))]
804impl<B: Blob> Blob for MigratingReadBlob<B> {
805    async fn read_at(
806        &self,
807        offset: u64,
808        len: usize,
809        options: ReadOptions,
810    ) -> Result<IoBufsMut, Error> {
811        self.read_at_buf(offset, len, IoBufMut::with_capacity(len), options)
812            .await
813    }
814
815    async fn read_at_buf(
816        &self,
817        offset: u64,
818        len: usize,
819        bufs: impl Into<IoBufsMut> + Send,
820        options: ReadOptions,
821    ) -> Result<IoBufsMut, Error> {
822        self.reads.fetch_add(1, Ordering::Relaxed);
823        let mut read = Box::pin(self.inner.read_at_buf(offset, len, bufs, options));
824
825        // Capture exactly one backend poll, then yield before exposing its result. Requiring a
826        // pending result lets a test establish that a task handoff carries unresolved backend I/O.
827        let first_poll = poll!(&mut read);
828        if self.require_pending {
829            assert!(
830                first_poll.is_pending(),
831                "blob read completed before registering pending I/O"
832            );
833        }
834        reschedule().await;
835
836        match first_poll {
837            Poll::Ready(result) => result,
838            Poll::Pending => read.await,
839        }
840    }
841
842    async fn write_at(
843        &self,
844        offset: u64,
845        bufs: impl Into<IoBufs> + Send,
846        options: WriteOptions,
847    ) -> Result<(), Error> {
848        self.inner.write_at(offset, bufs, options).await
849    }
850
851    async fn resize(&self, len: u64) -> Result<(), Error> {
852        self.inner.resize(len).await
853    }
854
855    async fn sync(&self) -> Result<(), Error> {
856        self.inner.sync().await
857    }
858
859    async fn start_sync(&self) -> Handle<()> {
860        self.inner.start_sync().await
861    }
862}
863
864/// Context wrapper whose blobs defer [Blob::start_sync] and can gate blocking syncs in tests.
865#[derive(Clone)]
866pub struct DelayedSyncContext<E> {
867    pub inner: E,
868    pub pending: PendingSyncs,
869}
870
871forward_context!(DelayedSyncContext, pending);
872
873impl<E: Spawner> Spawner for DelayedSyncContext<E> {
874    fn shared(mut self, blocking: bool) -> Self {
875        self.inner = self.inner.shared(blocking);
876        self
877    }
878
879    fn dedicated(mut self) -> Self {
880        self.inner = self.inner.dedicated();
881        self
882    }
883
884    fn spawn<F, Fut, T>(self, f: F) -> Handle<T>
885    where
886        F: FnOnce(Self) -> Fut + Send + 'static,
887        Fut: Future<Output = T> + Send + 'static,
888        T: Send + 'static,
889    {
890        let pending = self.pending;
891        self.inner.spawn(move |inner| f(Self { inner, pending }))
892    }
893
894    async fn stop(self, value: i32, timeout: Option<std::time::Duration>) -> Result<(), Error> {
895        self.inner.stop(value, timeout).await
896    }
897
898    fn stopped(&self) -> Signal {
899        self.inner.stopped()
900    }
901}
902
903impl<E: crate::Storage> crate::Storage for DelayedSyncContext<E> {
904    type Blob = DelayedSyncBlob<E::Blob>;
905
906    async fn open_versioned(
907        &self,
908        partition: &str,
909        name: &[u8],
910        versions: std::ops::RangeInclusive<BlobVersion>,
911    ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
912        let (inner, len, version) = self.inner.open_versioned(partition, name, versions).await?;
913        Ok((
914            DelayedSyncBlob {
915                inner: Arc::new(inner),
916                pending: self.pending.clone(),
917            },
918            len,
919            version,
920        ))
921    }
922
923    async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
924        self.inner.remove(partition, name).await
925    }
926
927    async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
928        self.inner.scan(partition).await
929    }
930}
931
932/// Blob wrapper that parks each started sync and supports one-shot blocking sync tracking.
933///
934/// A started sync keeps the wrapped blob, and with it the logical open, alive until its handle
935/// completes or drops.
936pub struct DelayedSyncBlob<B> {
937    inner: Arc<B>,
938    pending: PendingSyncs,
939}
940
941impl<B> DelayedSyncBlob<B> {
942    /// Wrap `inner`, returning the blob and the list its deferred syncs are pushed onto.
943    pub fn new(inner: B) -> (Self, PendingSyncs) {
944        let pending = PendingSyncs::default();
945        (
946            Self {
947                inner: Arc::new(inner),
948                pending: pending.clone(),
949            },
950            pending,
951        )
952    }
953}
954
955impl<B: Blob> Blob for DelayedSyncBlob<B> {
956    async fn read_at_buf(
957        &self,
958        offset: u64,
959        len: usize,
960        bufs: impl Into<IoBufsMut> + Send,
961        options: ReadOptions,
962    ) -> Result<IoBufsMut, Error> {
963        self.inner.read_at_buf(offset, len, bufs, options).await
964    }
965
966    async fn read_at(
967        &self,
968        offset: u64,
969        len: usize,
970        options: ReadOptions,
971    ) -> Result<IoBufsMut, Error> {
972        self.inner.read_at(offset, len, options).await
973    }
974
975    async fn write_at(
976        &self,
977        offset: u64,
978        bufs: impl Into<IoBufs> + Send,
979        options: WriteOptions,
980    ) -> Result<(), Error> {
981        if !options.contains(WriteOptions::SYNC) || !self.pending.tracking() {
982            return self.inner.write_at(offset, bufs, options).await;
983        }
984        self.inner
985            .write_at(offset, bufs, options.without(WriteOptions::SYNC))
986            .await?;
987        self.sync().await
988    }
989
990    async fn resize(&self, len: u64) -> Result<(), Error> {
991        self.inner.resize(len).await
992    }
993
994    async fn sync(&self) -> Result<(), Error> {
995        self.pending.wait().await?;
996        self.inner.sync().await
997    }
998
999    async fn start_sync(&self) -> Handle<()> {
1000        let pending = self.pending.clone();
1001        let inner = self.inner.clone();
1002        let waiter = {
1003            let mut state = pending.state.lock();
1004            state.starts += 1;
1005            // An armed gate takes precedence over parking.
1006            state.observe().or_else(|| state.park())
1007        };
1008        Handle::from_future(async move {
1009            let fail = {
1010                let mut state = pending.state.lock();
1011                state.entered += 1;
1012                state.fail
1013            };
1014            match waiter {
1015                Some(waiter) => waiter.wait().await?,
1016                None if fail => return Err(injected_sync_failure()),
1017                None => {}
1018            }
1019            inner.sync().await?;
1020            pending.state.lock().completions += 1;
1021            Ok(())
1022        })
1023    }
1024}
1025
1026/// Take the oldest pending sync, panicking if none was started.
1027pub fn next_pending_sync(pending: &PendingSyncs) -> DeferredSync {
1028    let mut pending = pending.lock();
1029    assert!(!pending.is_empty(), "no pending sync was started");
1030    pending.remove(0)
1031}
1032
1033/// Complete the oldest `count` pending syncs successfully.
1034pub fn release_next_pending_syncs(pending: &PendingSyncs, count: usize) {
1035    let syncs = {
1036        let mut pending = pending.lock();
1037        assert!(
1038            pending.len() >= count,
1039            "not enough pending syncs: have {}, need {count}",
1040            pending.len()
1041        );
1042        pending.drain(..count).collect::<Vec<_>>()
1043    };
1044    for sync in syncs {
1045        let _ = sync.release.send(Ok(()));
1046    }
1047}
1048
1049/// Complete all pending syncs successfully.
1050pub fn release_pending_syncs(pending: &PendingSyncs) {
1051    for sync in mem::take(&mut *pending.lock()) {
1052        let _ = sync.release.send(Ok(()));
1053    }
1054}
1055
1056/// Drive `fut` to completion, releasing any parked syncs each time it stalls.
1057#[boxed]
1058pub async fn drive_pending_syncs<T>(pending: &PendingSyncs, fut: impl Future<Output = T>) -> T {
1059    let mut fut = std::pin::pin!(fut);
1060    poll_fn(|cx| match fut.as_mut().poll(cx) {
1061        Poll::Ready(out) => Poll::Ready(out),
1062        Poll::Pending => {
1063            // A concurrent task may park a new sync after this release, so
1064            // self-wake to check again on the next scheduler tick.
1065            release_pending_syncs(pending);
1066            cx.waker().wake_by_ref();
1067            Poll::Pending
1068        }
1069    })
1070    .await
1071}
1072
1073/// Fail all pending syncs with an injected I/O error.
1074pub fn fail_pending_syncs(pending: &PendingSyncs) {
1075    for sync in mem::take(&mut *pending.lock()) {
1076        let _ = sync.release.send(Err(injected_sync_failure()));
1077    }
1078}
1079
1080/// The error injected by [fail_pending_syncs] and [PendingSyncs::arm_fail].
1081fn injected_sync_failure() -> Error {
1082    Error::Io(std::io::Error::other("injected sync failure").into())
1083}
1084
1085struct SyncWaiter {
1086    entered: oneshot::Sender<()>,
1087    release: oneshot::Receiver<Result<(), Error>>,
1088}
1089
1090impl SyncWaiter {
1091    async fn wait(self) -> Result<(), Error> {
1092        self.entered.send_lossy(());
1093        self.release.await.map_err(|_| Error::Closed)??;
1094        Ok(())
1095    }
1096}
1097
1098#[derive(Default)]
1099struct SyncGateState {
1100    tracking: bool,
1101    calls: usize,
1102    waiter: Option<SyncWaiter>,
1103}
1104
1105impl PendingSyncs {
1106    /// Locks the deferred sync queue.
1107    pub fn lock(&self) -> commonware_utils::sync::MappedMutexGuard<'_, Vec<DeferredSync>> {
1108        commonware_utils::sync::MutexGuard::map(self.state.lock(), |state| &mut state.syncs)
1109    }
1110
1111    /// Begins counting durability operations and blocks the next one behind a
1112    /// one-shot gate (pushed onto the deferred queue so tests can release it).
1113    ///
1114    /// Once the gate is consumed, started syncs park in the deferred queue as
1115    /// usual while [Self::calls] keeps counting.
1116    pub fn arm(&self) {
1117        let mut state = self.state.lock();
1118        assert!(!state.gate.tracking, "sync gate already armed");
1119        assert!(
1120            state.gate.waiter.is_none(),
1121            "sync gate already has a waiter"
1122        );
1123        state.gate.tracking = true;
1124        state.gate.calls = 0;
1125        let waiter = state.defer();
1126        state.gate.waiter = Some(waiter);
1127    }
1128
1129    /// Returns the number of durability operations observed since [Self::arm].
1130    pub fn calls(&self) -> usize {
1131        self.state.lock().gate.calls
1132    }
1133
1134    fn tracking(&self) -> bool {
1135        self.state.lock().gate.tracking
1136    }
1137
1138    /// Releases every parked sync and permanently stops parking new ones: future started syncs
1139    /// proceed immediately (or fail, after [Self::arm_fail]). Unlike [release_pending_syncs],
1140    /// which drains the queue once, this is sticky. A gate armed via [Self::arm] parks in the
1141    /// same queue, so this releases it too.
1142    pub fn unblock(&self) {
1143        let (drained, fail) = {
1144            let mut state = self.state.lock();
1145            state.unblocked = true;
1146            (mem::take(&mut state.syncs), state.fail)
1147        };
1148        for sync in drained {
1149            let result = if fail {
1150                Err(injected_sync_failure())
1151            } else {
1152                Ok(())
1153            };
1154            let _ = sync.release.send(result);
1155        }
1156    }
1157
1158    /// Arms every parked and future started sync to resolve to an injected error once released
1159    /// via [Self::unblock]. The release helpers send explicit results and ignore this.
1160    pub fn arm_fail(&self) {
1161        self.state.lock().fail = true;
1162    }
1163
1164    /// Number of started syncs issued.
1165    pub fn starts(&self) -> usize {
1166        self.state.lock().starts
1167    }
1168
1169    /// Number of started syncs whose completion futures have begun executing, parked or not.
1170    pub fn entered(&self) -> usize {
1171        self.state.lock().entered
1172    }
1173
1174    /// Number of started syncs that completed durably.
1175    pub fn completions(&self) -> usize {
1176        self.state.lock().completions
1177    }
1178
1179    async fn wait(&self) -> Result<(), Error> {
1180        let waiter = self.state.lock().observe();
1181        match waiter {
1182            Some(waiter) => waiter.wait().await,
1183            None => Ok(()),
1184        }
1185    }
1186}
1187
1188/// Controls a [WriteFaultContext]: while armed, every `write_at` fails with an
1189/// injected error. Successful writes are counted.
1190#[derive(Clone, Default)]
1191pub struct WriteFaults {
1192    state: Arc<Mutex<WriteFaultState>>,
1193}
1194
1195#[derive(Default)]
1196struct WriteFaultState {
1197    fail: bool,
1198    writes: u64,
1199}
1200
1201impl WriteFaults {
1202    /// Start failing writes.
1203    pub fn arm(&self) {
1204        self.state.lock().fail = true;
1205    }
1206
1207    /// Stop failing writes.
1208    pub fn disarm(&self) {
1209        self.state.lock().fail = false;
1210    }
1211
1212    /// The number of successful writes so far.
1213    pub fn writes(&self) -> u64 {
1214        self.state.lock().writes
1215    }
1216
1217    fn check(&self) -> Result<(), Error> {
1218        if self.state.lock().fail {
1219            return Err(Error::Io(
1220                std::io::Error::other("injected write failure").into(),
1221            ));
1222        }
1223        Ok(())
1224    }
1225
1226    fn note(&self) {
1227        self.state.lock().writes += 1;
1228    }
1229}
1230
1231/// Context wrapper whose blobs fail `write_at` while the shared [WriteFaults]
1232/// is armed, counting successful writes. Unlike [DelayedSyncContext], this injects failures
1233/// into inline writes issued before any blob sync starts.
1234#[derive(Clone)]
1235pub struct WriteFaultContext<E> {
1236    pub inner: E,
1237    pub faults: WriteFaults,
1238}
1239
1240forward_context!(WriteFaultContext, faults);
1241
1242impl<E: crate::Storage> crate::Storage for WriteFaultContext<E> {
1243    type Blob = WriteFaultBlob<E::Blob>;
1244
1245    async fn open_versioned(
1246        &self,
1247        partition: &str,
1248        name: &[u8],
1249        versions: std::ops::RangeInclusive<BlobVersion>,
1250    ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
1251        let (inner, len, version) = self.inner.open_versioned(partition, name, versions).await?;
1252        Ok((
1253            WriteFaultBlob {
1254                inner,
1255                faults: self.faults.clone(),
1256            },
1257            len,
1258            version,
1259        ))
1260    }
1261
1262    async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
1263        self.inner.remove(partition, name).await
1264    }
1265
1266    async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
1267        self.inner.scan(partition).await
1268    }
1269}
1270
1271/// Blob wrapper that fails `write_at` while its [WriteFaults] is armed.
1272pub struct WriteFaultBlob<B> {
1273    inner: B,
1274    faults: WriteFaults,
1275}
1276
1277impl<B: Blob> Blob for WriteFaultBlob<B> {
1278    async fn read_at_buf(
1279        &self,
1280        offset: u64,
1281        len: usize,
1282        bufs: impl Into<IoBufsMut> + Send,
1283        options: ReadOptions,
1284    ) -> Result<IoBufsMut, Error> {
1285        self.inner.read_at_buf(offset, len, bufs, options).await
1286    }
1287
1288    async fn read_at(
1289        &self,
1290        offset: u64,
1291        len: usize,
1292        options: ReadOptions,
1293    ) -> Result<IoBufsMut, Error> {
1294        self.inner.read_at(offset, len, options).await
1295    }
1296
1297    async fn write_at(
1298        &self,
1299        offset: u64,
1300        bufs: impl Into<IoBufs> + Send,
1301        options: WriteOptions,
1302    ) -> Result<(), Error> {
1303        self.faults.check()?;
1304        self.inner.write_at(offset, bufs, options).await?;
1305        self.faults.note();
1306        Ok(())
1307    }
1308
1309    async fn resize(&self, len: u64) -> Result<(), Error> {
1310        self.inner.resize(len).await
1311    }
1312
1313    async fn sync(&self) -> Result<(), Error> {
1314        self.inner.sync().await
1315    }
1316
1317    async fn start_sync(&self) -> Handle<()> {
1318        self.inner.start_sync().await
1319    }
1320}
1321
1322/// Context wrapper whose blobs fail `sync` and `start_sync` for a single partition.
1323#[derive(Clone)]
1324pub struct SyncFaultContext<E> {
1325    pub inner: E,
1326    pub fail_partition: String,
1327}
1328
1329forward_context!(SyncFaultContext, fail_partition);
1330
1331impl<E: crate::Storage> crate::Storage for SyncFaultContext<E> {
1332    type Blob = SyncFaultBlob<E::Blob>;
1333
1334    async fn open_versioned(
1335        &self,
1336        partition: &str,
1337        name: &[u8],
1338        versions: std::ops::RangeInclusive<BlobVersion>,
1339    ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
1340        let (inner, len, version) = self.inner.open_versioned(partition, name, versions).await?;
1341        Ok((
1342            SyncFaultBlob {
1343                inner,
1344                faulty: partition == self.fail_partition,
1345            },
1346            len,
1347            version,
1348        ))
1349    }
1350
1351    async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
1352        self.inner.remove(partition, name).await
1353    }
1354
1355    async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
1356        self.inner.scan(partition).await
1357    }
1358}
1359
1360/// Blob wrapper that fails `sync` and `start_sync` when marked faulty.
1361pub struct SyncFaultBlob<B> {
1362    inner: B,
1363    faulty: bool,
1364}
1365
1366impl<B: Blob> Blob for SyncFaultBlob<B> {
1367    async fn read_at_buf(
1368        &self,
1369        offset: u64,
1370        len: usize,
1371        bufs: impl Into<IoBufsMut> + Send,
1372        options: ReadOptions,
1373    ) -> Result<IoBufsMut, Error> {
1374        self.inner.read_at_buf(offset, len, bufs, options).await
1375    }
1376
1377    async fn read_at(
1378        &self,
1379        offset: u64,
1380        len: usize,
1381        options: ReadOptions,
1382    ) -> Result<IoBufsMut, Error> {
1383        self.inner.read_at(offset, len, options).await
1384    }
1385
1386    async fn write_at(
1387        &self,
1388        offset: u64,
1389        bufs: impl Into<IoBufs> + Send,
1390        options: WriteOptions,
1391    ) -> Result<(), Error> {
1392        self.inner.write_at(offset, bufs, options).await
1393    }
1394
1395    async fn resize(&self, len: u64) -> Result<(), Error> {
1396        self.inner.resize(len).await
1397    }
1398
1399    async fn sync(&self) -> Result<(), Error> {
1400        if self.faulty {
1401            let err = std::io::Error::other("injected partition sync fault");
1402            return Err(Error::Io(err.into()));
1403        }
1404        self.inner.sync().await
1405    }
1406
1407    async fn start_sync(&self) -> Handle<()> {
1408        if self.faulty {
1409            return Handle::ready(self.sync().await);
1410        }
1411        self.inner.start_sync().await
1412    }
1413}
1414
1415#[cfg(test)]
1416mod tests {
1417    use super::*;
1418    use crate::{Clock, IoBufMut, Runner, Sink, Spawner, Storage as _, Stream, deterministic};
1419    use commonware_macros::select;
1420    use std::{thread::sleep, time::Duration};
1421
1422    #[test]
1423    fn recording_context_preserves_data_and_records_options() {
1424        deterministic::Runner::default().start(|context| async move {
1425            let (context, recordings) = RecordingContext::new(context);
1426            let (blob, _) = context.open("recording", b"blob").await.unwrap();
1427
1428            blob.write_at(0, b"data", WriteOptions::DONT_CACHE)
1429                .await
1430                .unwrap();
1431            let read = blob.read_at(0, 4, ReadOptions::DONT_CACHE).await.unwrap();
1432            assert_eq!(read.coalesce(), b"data");
1433
1434            let read = blob
1435                .read_at_buf(0, 4, IoBufMut::with_capacity(4), ReadOptions::default())
1436                .await
1437                .unwrap();
1438            assert_eq!(read.coalesce(), b"data");
1439
1440            assert_eq!(
1441                recordings.snapshot(),
1442                RecordingSnapshot {
1443                    reads: vec![ReadOptions::DONT_CACHE, ReadOptions::default()],
1444                    batches: Vec::new(),
1445                    writes: vec![WriteOptions::DONT_CACHE],
1446                }
1447            );
1448            recordings.clear();
1449            assert_eq!(recordings.snapshot(), RecordingSnapshot::default());
1450        });
1451    }
1452
1453    async fn assert_read_options_forwarded<E: crate::Storage>(
1454        context: &E,
1455        recordings: &Recordings,
1456        partition: &str,
1457    ) {
1458        let (blob, _) = context.open(partition, b"blob").await.unwrap();
1459        blob.write_at(0, b"data", WriteOptions::default())
1460            .await
1461            .unwrap();
1462        recordings.clear();
1463
1464        let read = blob.read_at(0, 4, ReadOptions::DONT_CACHE).await.unwrap();
1465        assert_eq!(read.coalesce(), b"data");
1466
1467        let read = blob
1468            .read_at_buf(0, 4, IoBufMut::with_capacity(4), ReadOptions::DONT_CACHE)
1469            .await
1470            .unwrap();
1471        assert_eq!(read.coalesce(), b"data");
1472
1473        assert_eq!(
1474            recordings.snapshot(),
1475            RecordingSnapshot {
1476                reads: vec![ReadOptions::DONT_CACHE, ReadOptions::DONT_CACHE],
1477                batches: Vec::new(),
1478                writes: Vec::new(),
1479            }
1480        );
1481    }
1482
1483    #[test]
1484    fn delayed_sync_blob_forwards_read_options() {
1485        deterministic::Runner::default().start(|context| async move {
1486            let (inner, recordings) = RecordingContext::new(context);
1487            let context = DelayedSyncContext {
1488                inner,
1489                pending: PendingSyncs::default(),
1490            };
1491
1492            assert_read_options_forwarded(&context, &recordings, "delayed_sync").await;
1493        });
1494    }
1495
1496    #[test]
1497    fn write_fault_blob_forwards_read_options() {
1498        deterministic::Runner::default().start(|context| async move {
1499            let (inner, recordings) = RecordingContext::new(context);
1500            let context = WriteFaultContext {
1501                inner,
1502                faults: WriteFaults::default(),
1503            };
1504
1505            assert_read_options_forwarded(&context, &recordings, "write_fault").await;
1506        });
1507    }
1508
1509    #[test]
1510    fn sync_fault_blob_forwards_read_options() {
1511        deterministic::Runner::default().start(|context| async move {
1512            let (inner, recordings) = RecordingContext::new(context);
1513            let context = SyncFaultContext {
1514                inner,
1515                fail_partition: "sync_fault".to_string(),
1516            };
1517
1518            assert_read_options_forwarded(&context, &recordings, "sync_fault").await;
1519        });
1520    }
1521
1522    #[test]
1523    fn test_send_recv() {
1524        let (mut sink, mut stream) = Channel::init();
1525        let data = b"hello world";
1526
1527        let executor = deterministic::Runner::default();
1528        executor.start(|_| async move {
1529            sink.send(data.as_slice()).await.unwrap();
1530            let received = stream.recv(data.len()).await.unwrap();
1531            assert_eq!(received.coalesce(), data);
1532        });
1533    }
1534
1535    #[test]
1536    fn test_send_recv_partial_multiple() {
1537        let (mut sink, mut stream) = Channel::init();
1538        let data = b"hello";
1539        let data2 = b" world";
1540
1541        let executor = deterministic::Runner::default();
1542        executor.start(|_| async move {
1543            sink.send(data.as_slice()).await.unwrap();
1544            sink.send(data2.as_slice()).await.unwrap();
1545            let received = stream.recv(5).await.unwrap();
1546            assert_eq!(received.coalesce(), b"hello");
1547            let received = stream.recv(5).await.unwrap();
1548            assert_eq!(received.coalesce(), b" worl");
1549            let received = stream.recv(1).await.unwrap();
1550            assert_eq!(received.coalesce(), b"d");
1551        });
1552    }
1553
1554    #[test]
1555    fn test_send_recv_async() {
1556        let (mut sink, mut stream) = Channel::init();
1557        let data = b"hello world";
1558
1559        let executor = deterministic::Runner::default();
1560        executor.start(|_| async move {
1561            let (received, _) = futures::try_join!(stream.recv(data.len()), async {
1562                sleep(Duration::from_millis(50));
1563                sink.send(data.as_slice()).await
1564            })
1565            .unwrap();
1566            assert_eq!(received.coalesce(), data);
1567        });
1568    }
1569
1570    #[test]
1571    fn test_recv_error_sink_dropped_while_waiting() {
1572        let (sink, mut stream) = Channel::init();
1573
1574        let executor = deterministic::Runner::default();
1575        executor.start(|context| async move {
1576            futures::join!(
1577                async {
1578                    let result = stream.recv(5).await;
1579                    assert!(matches!(result, Err(Error::RecvFailed)));
1580                    let result = stream.recv(5).await;
1581                    assert!(matches!(result, Err(Error::Closed)));
1582                },
1583                async {
1584                    // Wait for the stream to start waiting
1585                    context.sleep(Duration::from_millis(50)).await;
1586                    drop(sink);
1587                }
1588            );
1589        });
1590    }
1591
1592    #[test]
1593    fn test_recv_error_sink_dropped_before_recv() {
1594        let (sink, mut stream) = Channel::init();
1595        drop(sink); // Drop sink immediately
1596
1597        let executor = deterministic::Runner::default();
1598        executor.start(|_| async move {
1599            let result = stream.recv(5).await;
1600            assert!(matches!(result, Err(Error::RecvFailed)));
1601            let result = stream.recv(5).await;
1602            assert!(matches!(result, Err(Error::Closed)));
1603        });
1604    }
1605
1606    #[test]
1607    fn test_send_error_stream_dropped() {
1608        let (mut sink, mut stream) = Channel::init();
1609
1610        let executor = deterministic::Runner::default();
1611        executor.start(|context| async move {
1612            // Send some bytes
1613            assert!(sink.send(b"7 bytes".as_slice()).await.is_ok());
1614
1615            // Spawn a task to initiate recv's where the first one will succeed and then will drop.
1616            let handle = context.child("recv").spawn(|_| async move {
1617                let _ = stream.recv(5).await;
1618                let _ = stream.recv(5).await;
1619            });
1620
1621            // Give the async task a moment to start
1622            context.sleep(Duration::from_millis(50)).await;
1623
1624            // Drop the stream by aborting the handle
1625            handle.abort();
1626            assert!(matches!(handle.await, Err(Error::Closed)));
1627
1628            // Try to send a message. The stream is dropped, so this should fail.
1629            let result = sink.send(b"hello world".as_slice()).await;
1630            assert!(matches!(result, Err(Error::SendFailed)));
1631            let result = sink.send(b"hello world".as_slice()).await;
1632            assert!(matches!(result, Err(Error::Closed)));
1633        });
1634    }
1635
1636    #[test]
1637    fn test_send_error_stream_dropped_before_send() {
1638        let (mut sink, stream) = Channel::init();
1639        drop(stream); // Drop stream immediately
1640
1641        let executor = deterministic::Runner::default();
1642        executor.start(|_| async move {
1643            let result = sink.send(b"hello world".as_slice()).await;
1644            assert!(matches!(result, Err(Error::SendFailed)));
1645            let result = sink.send(b"hello world".as_slice()).await;
1646            assert!(matches!(result, Err(Error::Closed)));
1647        });
1648    }
1649
1650    #[test]
1651    fn test_recv_timeout() {
1652        let (_sink, mut stream) = Channel::init();
1653
1654        // If there is no data to read, test that the recv function just blocks.
1655        // The timeout should return first.
1656        let executor = deterministic::Runner::default();
1657        executor.start(|context| async move {
1658            select! {
1659                v = stream.recv(5) => {
1660                    panic!("unexpected value: {v:?}");
1661                },
1662                _ = context.sleep(Duration::from_millis(100)) => "timeout",
1663            };
1664        });
1665    }
1666
1667    #[test]
1668    fn test_peek_empty() {
1669        let (_sink, stream) = Channel::init();
1670
1671        // Peek on a fresh stream should return empty slice
1672        assert!(stream.peek(10).is_empty());
1673    }
1674
1675    #[test]
1676    fn test_peek_after_partial_recv() {
1677        let (mut sink, mut stream) = Channel::init();
1678
1679        let executor = deterministic::Runner::default();
1680        executor.start(|_| async move {
1681            // Send more data than we'll consume
1682            sink.send(b"hello world".as_slice()).await.unwrap();
1683
1684            // Recv only part of it
1685            let received = stream.recv(5).await.unwrap();
1686            assert_eq!(received.coalesce(), b"hello");
1687
1688            // Peek should show the remaining data
1689            assert_eq!(stream.peek(100), b" world");
1690
1691            // Peek with smaller max_len
1692            assert_eq!(stream.peek(3), b" wo");
1693
1694            // Peek doesn't consume - can peek again
1695            assert_eq!(stream.peek(100), b" world");
1696
1697            // Recv consumes the peeked data
1698            let received = stream.recv(6).await.unwrap();
1699            assert_eq!(received.coalesce(), b" world");
1700
1701            // Peek is now empty
1702            assert!(stream.peek(100).is_empty());
1703        });
1704    }
1705
1706    #[test]
1707    fn test_peek_after_recv_wakeup() {
1708        let (mut sink, mut stream) = Channel::init_with_buffer_size(64);
1709
1710        let executor = deterministic::Runner::default();
1711        executor.start(|context| async move {
1712            // Spawn recv that will block waiting
1713            let (tx, rx) = oneshot::channel();
1714            let recv_handle = context.child("recv").spawn(|_| async move {
1715                let data = stream.recv(3).await.unwrap();
1716                tx.send(stream).ok();
1717                data
1718            });
1719
1720            // Let recv set up waiter
1721            context.sleep(Duration::from_millis(10)).await;
1722
1723            // Send more than requested
1724            sink.send(b"ABCDEFGHIJ".as_slice()).await.unwrap();
1725
1726            // Recv gets its 3 bytes
1727            let received = recv_handle.await.unwrap();
1728            assert_eq!(received.coalesce(), b"ABC");
1729
1730            // Get stream back and verify peek sees remaining data
1731            let stream = rx.await.unwrap();
1732            assert_eq!(stream.peek(100), b"DEFGHIJ");
1733        });
1734    }
1735
1736    #[test]
1737    fn test_peek_multiple_sends() {
1738        let (mut sink, mut stream) = Channel::init();
1739
1740        let executor = deterministic::Runner::default();
1741        executor.start(|_| async move {
1742            // Send multiple chunks
1743            sink.send(b"aaa".as_slice()).await.unwrap();
1744            sink.send(b"bbb".as_slice()).await.unwrap();
1745            sink.send(b"ccc".as_slice()).await.unwrap();
1746
1747            // Recv less than total
1748            let received = stream.recv(4).await.unwrap();
1749            assert_eq!(received.coalesce(), b"aaab");
1750
1751            // Peek should show remaining
1752            assert_eq!(stream.peek(100), b"bbccc");
1753        });
1754    }
1755
1756    #[test]
1757    fn test_buffer_size_limit() {
1758        // Use a small buffer capacity for testing
1759        let (mut sink, mut stream) = Channel::init_with_buffer_size(10);
1760
1761        let executor = deterministic::Runner::default();
1762        executor.start(|context| async move {
1763            // Send more than buffer capacity concurrently with recv
1764            // so the sender can drain via backpressure.
1765            let send_handle = context.child("sender").spawn(|_| async move {
1766                sink.send(b"0123456789ABCDEF".as_slice()).await.unwrap();
1767                sink
1768            });
1769
1770            // Recv a small amount - should only pull up to capacity (10 bytes)
1771            let received = stream.recv(2).await.unwrap();
1772            assert_eq!(received.coalesce(), b"01");
1773
1774            // Peek should show remaining buffered data (8 bytes, not 14)
1775            assert_eq!(stream.peek(100), b"23456789");
1776
1777            // The rest should still be in the channel buffer
1778            // Recv more to pull the remaining data
1779            let received = stream.recv(8).await.unwrap();
1780            assert_eq!(received.coalesce(), b"23456789");
1781
1782            // Now peek should show next chunk from channel (up to capacity)
1783            let received = stream.recv(2).await.unwrap();
1784            assert_eq!(received.coalesce(), b"AB");
1785
1786            assert_eq!(stream.peek(100), b"CDEF");
1787
1788            // Ensure the sender completes
1789            send_handle.await.unwrap();
1790        });
1791    }
1792
1793    #[test]
1794    fn test_recv_before_send() {
1795        // Use a small buffer capacity for testing
1796        let (mut sink, mut stream) = Channel::init_with_buffer_size(10);
1797
1798        let executor = deterministic::Runner::default();
1799        executor.start(|context| async move {
1800            // Start recv before send (will wait)
1801            let recv_handle = context
1802                .child("recv")
1803                .spawn(|_| async move { stream.recv(3).await.unwrap() });
1804
1805            // Give recv time to set up waiter
1806            context.sleep(Duration::from_millis(10)).await;
1807
1808            // Send more than capacity
1809            sink.send(b"ABCDEFGHIJKLMNOP".as_slice()).await.unwrap();
1810
1811            // Recv should get its 3 bytes
1812            let received = recv_handle.await.unwrap();
1813            assert_eq!(received.coalesce(), b"ABC");
1814        });
1815    }
1816}