arcium-primitives 0.8.5

pub mod config;
pub mod messages;

use std::{
    cell::Cell,
    collections::VecDeque,
    marker::PhantomData,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
};

pub use config::{Buffer, BufferConfig, SharedBufferConfig};
use futures::{stream::FuturesOrdered, StreamExt};
use log::{debug, error, info};
pub use messages::PrefetchHandle;
use parking_lot::RwLock;
use tokio::{
    sync::{
        mpsc::{self, UnboundedReceiver, UnboundedSender},
        oneshot,
    },
    task::JoinHandle,
};

use crate::correlated_randomness::{
    generator::{CorrelationGenerator, PipelinedCorrelationGenerator},
    stream::{
        buffered::{config::try_read_config, messages::Command},
        errors::CorrelatedStreamError,
        futures::Next,
        CorrelatedStream,
        NextVec,
        ResyncHandle,
    },
    CorrelatedBatch,
};

/// A unit of background work sent from the dispatcher to the generator task.
enum Work {
    /// Generate at least `n` elements and return them.
    Generate(usize),
    /// Advance the generator's logical position by `n` without materializing the elements.
    Skip(usize),
}

/// Buffering over a correlation generator, providing both demand-driven streaming
/// and proactive prefetching interfaces. The buffer is refilled in the background by a
/// dispatcher/generator task pair; `capacity` bounds aggregate outstanding demand and
/// `refill_threshold` drives proactive top-ups.
pub struct BufferedStream<PB: CorrelatedBatch, E> {
    command_sender: UnboundedSender<Command<PB, E>>,
    _unsync_marker: PhantomData<Cell<()>>,
    config: Arc<RwLock<BufferConfig>>,
    /// Logical position: cumulative items delivered through `next_n`. Written by the dispatcher
    /// (the single source of truth) and read synchronously here. See
    /// [`CorrelatedStream::position`].
    position: Arc<AtomicU64>,
    /// Elements currently held in the buffer, ready for immediate delivery. Written by the
    /// dispatcher whenever the buffer changes and read synchronously here. See
    /// [`CorrelatedStream::buffered`].
    buffered: Arc<AtomicU64>,
    dispatcher_handle: JoinHandle<()>,
    generator_handle: JoinHandle<()>,
}

impl<PB: CorrelatedBatch, E> Buffer for BufferedStream<PB, E> {
    fn config(&self) -> &Arc<RwLock<BufferConfig>> {
        &self.config
    }
}

impl<
        PB: CorrelatedBatch,
        E: From<CorrelatedStreamError> + Clone + Send + std::fmt::Debug + 'static,
    > BufferedStream<PB, E>
{
    /// Creates a new stream.
    pub fn new<G: CorrelationGenerator<PB> + Send + 'static>(
        generator: G,
        net: G::Net,
        config: BufferConfig,
    ) -> Self
    where
        E: From<G::Error>,
    {
        Self::new_with_shared_config(generator, net, Arc::new(RwLock::new(config)))
    }

    /// Creates a purely on-demand stream: no proactive refill and effectively unbounded admission,
    /// so every `next_n` triggers generation of exactly the shortfall (any surplus is buffered).
    /// Suited to streams whose generator is itself the buffer/limiter (e.g. a dealer client).
    pub fn new_on_demand<G: CorrelationGenerator<PB> + Send + 'static>(
        generator: G,
        net: G::Net,
    ) -> Self
    where
        E: From<G::Error>,
    {
        Self::new(
            generator,
            net,
            BufferConfig::lazy(BufferConfig::UNBOUNDED, 0),
        )
    }

    /// Like [`new`](Self::new), but for a [`PipelinedCorrelationGenerator`]: up to
    /// `max_orders_in_flight` generation orders may be outstanding at once, so consecutive
    /// request/generate/transfer cycles overlap instead of running stop-and-wait. Orders are
    /// issued and resolved in FIFO order; demand accounting nets out in-flight quantities, so
    /// pipelining never requests more than the unmet demand.
    ///
    /// # Depth bound
    ///
    /// Any depth is accepted here — this crate cannot see what the generator can replay. Callers
    /// pairing this with a dealer client **must** keep the depth within the dealer's reply-replay
    /// window (`REPLY_CACHE_CAPACITY` in `arcium-dealer`): a reconnect resends every outstanding
    /// request, and one whose cached reply has been evicted is served fresh, silently desyncing
    /// this party's stream from the others'. The dealer bundler builder enforces that bound.
    ///
    /// # Panics
    ///
    /// Panics if `max_orders_in_flight` is 0.
    pub fn new_pipelined<G: PipelinedCorrelationGenerator<PB> + Send + 'static>(
        generator: G,
        net: G::Net,
        config: BufferConfig,
        max_orders_in_flight: usize,
    ) -> Self
    where
        E: From<G::Error>,
    {
        assert!(
            max_orders_in_flight > 0,
            "pipeline depth must be at least 1"
        );
        let (parts, work_rx, items_tx, skip_tx) = Self::spawn_dispatcher(
            Arc::new(RwLock::new(config)),
            max_orders_in_flight,
            G::SUPPORTS_UNILATERAL_SKIP,
        );
        let generator_handle = tokio::spawn(pipelined_generator_loop(
            generator, net, work_rx, items_tx, skip_tx,
        ));
        parts.into_stream(generator_handle)
    }

    /// [`new_on_demand`](Self::new_on_demand) over a pipelined generator; see
    /// [`new_pipelined`](Self::new_pipelined).
    pub fn new_on_demand_pipelined<G: PipelinedCorrelationGenerator<PB> + Send + 'static>(
        generator: G,
        net: G::Net,
        max_orders_in_flight: usize,
    ) -> Self
    where
        E: From<G::Error>,
    {
        Self::new_pipelined(
            generator,
            net,
            BufferConfig::lazy(BufferConfig::UNBOUNDED, 0),
            max_orders_in_flight,
        )
    }

    /// Creates a new stream.
    pub fn new_with_shared_config<G: CorrelationGenerator<PB> + Send + 'static>(
        generator: G,
        net: G::Net,
        config: Arc<RwLock<BufferConfig>>,
    ) -> Self
    where
        E: From<G::Error>,
    {
        let (parts, work_rx, items_tx, skip_tx) =
            Self::spawn_dispatcher(config, 1, G::SUPPORTS_UNILATERAL_SKIP);
        let generator_handle =
            tokio::spawn(generator_loop(generator, net, work_rx, items_tx, skip_tx));
        parts.into_stream(generator_handle)
    }

    /// Spawns the dispatcher task and returns the stream's front-end parts plus the generator-side
    /// channel ends, so each constructor can spawn its own flavor of generator task.
    #[allow(clippy::type_complexity)]
    fn spawn_dispatcher(
        config: Arc<RwLock<BufferConfig>>,
        max_orders_in_flight: usize,
        supports_skip: bool,
    ) -> (
        StreamParts<PB, E>,
        UnboundedReceiver<Work>,
        UnboundedSender<Result<Vec<PB::Item>, E>>,
        UnboundedSender<Result<(), E>>,
    ) {
        let (cmd_tx, cmd_rx) = mpsc::unbounded_channel::<Command<PB, E>>();
        let (work_tx, work_rx) = mpsc::unbounded_channel::<Work>();
        let (items_tx, items_rx) = mpsc::unbounded_channel::<Result<Vec<PB::Item>, E>>();
        let (skip_tx, skip_rx) = mpsc::unbounded_channel::<Result<(), E>>();
        let position = Arc::new(AtomicU64::new(0));
        let buffered = Arc::new(AtomicU64::new(0));
        let dispatcher_handle = tokio::spawn(dispatcher_loop(
            cmd_rx,
            work_tx,
            items_rx,
            skip_rx,
            config.clone(),
            position.clone(),
            buffered.clone(),
            supports_skip,
            max_orders_in_flight,
        ));
        let parts = StreamParts {
            command_sender: cmd_tx,
            config,
            position,
            buffered,
            dispatcher_handle,
        };
        (parts, work_rx, items_tx, skip_tx)
    }

    /// Shuts down the buffer gracefully and waits for all background tasks to exit.
    ///
    /// Closing the command channel causes the dispatcher to exit, which in turn drops the work
    /// channel, causing the generator to exit.  Both tasks are awaited before returning.
    pub async fn stop(self) {
        let Self {
            command_sender,
            dispatcher_handle,
            generator_handle,
            ..
        } = self;
        drop(command_sender);
        let _ = dispatcher_handle.await;
        let _ = generator_handle.await;
    }
}

/// Front-end pieces of a stream whose dispatcher is already running; the constructor completes it
/// by attaching the generator task's handle.
struct StreamParts<PB: CorrelatedBatch, E> {
    command_sender: UnboundedSender<Command<PB, E>>,
    config: Arc<RwLock<BufferConfig>>,
    position: Arc<AtomicU64>,
    buffered: Arc<AtomicU64>,
    dispatcher_handle: JoinHandle<()>,
}

impl<PB: CorrelatedBatch, E> StreamParts<PB, E> {
    fn into_stream(self, generator_handle: JoinHandle<()>) -> BufferedStream<PB, E> {
        let Self {
            command_sender,
            config,
            position,
            buffered,
            dispatcher_handle,
        } = self;
        BufferedStream {
            command_sender,
            _unsync_marker: PhantomData,
            config,
            position,
            buffered,
            dispatcher_handle,
            generator_handle,
        }
    }
}

// ============================
// ===== Generator Task =======
// ============================

async fn generator_loop<
    PB: CorrelatedBatch,
    G: CorrelationGenerator<PB> + Send,
    E: From<G::Error> + Send,
>(
    mut generator: G,
    mut net: G::Net,
    mut work_rx: UnboundedReceiver<Work>,
    items_tx: UnboundedSender<Result<Vec<PB::Item>, E>>,
    skip_tx: UnboundedSender<Result<(), E>>,
) {
    let log_prefix = format!("<Generator<{}>>", std::any::type_name::<PB>());
    while let Some(work) = work_rx.recv().await {
        match work {
            Work::Generate(n) => {
                debug!("{log_prefix} generating {n} elements");
                let result = generator.run_for(n, &mut net).await.map_err(E::from);
                let stop = result.is_err();
                // Stop on either a generator error or a dropped dispatcher.
                if items_tx.send(result).is_err() || stop {
                    break;
                }
            }
            Work::Skip(n) => {
                debug!("{log_prefix} skipping {n} elements");
                let result = generator.skip(n, &mut net).await.map_err(E::from);
                let stop = result.is_err();
                if skip_tx.send(result).is_err() || stop {
                    break;
                }
            }
        }
    }
    info!("{log_prefix} exiting");
}

/// Like [`generator_loop`], but issues orders without awaiting them: new work is dispatched
/// through [`PipelinedCorrelationGenerator::issue_for`] while earlier orders' futures are still
/// pending, and results are forwarded in issue (FIFO) order. The dispatcher only sends a skip
/// once no generation orders are outstanding, so skips never overtake generates.
async fn pipelined_generator_loop<
    PB: CorrelatedBatch,
    G: PipelinedCorrelationGenerator<PB> + Send,
    E: From<G::Error> + Send,
>(
    mut generator: G,
    mut net: G::Net,
    mut work_rx: UnboundedReceiver<Work>,
    items_tx: UnboundedSender<Result<Vec<PB::Item>, E>>,
    skip_tx: UnboundedSender<Result<(), E>>,
) {
    let log_prefix = format!("<PipelinedGenerator<{}>>", std::any::type_name::<PB>());
    let mut in_flight = FuturesOrdered::new();
    loop {
        tokio::select! {
            work = work_rx.recv() => {
                match work {
                    Some(Work::Generate(n)) => {
                        debug!("{log_prefix} issuing order for {n} elements ({} already in flight)", in_flight.len());
                        in_flight.push_back(generator.issue_for(n, &mut net));
                    }
                    Some(Work::Skip(n)) => {
                        if !in_flight.is_empty() {
                            // The dispatcher only sends a skip once every order has resolved. A
                            // violation means the skip would overtake a generate, and since both
                            // advance the stream position, this party would silently desync from
                            // its peers. Checked in release too: shutting the stream down (which
                            // fails every consumer) beats producing wrong correlations.
                            error!(
                                "{log_prefix} skip of {n} dispatched with {} generate order(s) in \
                                 flight, shutting down",
                                in_flight.len()
                            );
                            break;
                        }
                        debug!("{log_prefix} skipping {n} elements");
                        let result = generator.skip(n, &mut net).await.map_err(E::from);
                        let stop = result.is_err();
                        if skip_tx.send(result).is_err() || stop {
                            break;
                        }
                    }
                    None => break,
                }
            }
            Some(result) = in_flight.next() => {
                let result = result.map_err(E::from);
                let stop = result.is_err();
                if items_tx.send(result).is_err() || stop {
                    break;
                }
            }
        }
    }
    info!("{log_prefix} exiting");
}

// ============================
// ===== Dispatcher Task ======
// ============================

type ItemsCollected<PB> = Vec<<PB as IntoIterator>::Item>;
type TotalNeeded = usize;
type BatchSender<PB, E> = oneshot::Sender<Result<Vec<<PB as IntoIterator>::Item>, E>>;

/// Dispatches a pending resync skip to the generator once no other work is in flight. Skips must
/// not overtake generation orders (both advance the stream position), so the skip waits until
/// every outstanding order has resolved.
fn maybe_skip(
    pending_skip: &mut usize,
    skip_in_flight: &mut bool,
    orders_in_flight: usize,
    work_tx: &UnboundedSender<Work>,
    log_prefix: &str,
) {
    if *pending_skip > 0 && !*skip_in_flight && orders_in_flight == 0 {
        debug!("{log_prefix} requesting skip of {} elements", *pending_skip);
        *skip_in_flight = true;
        let _ = work_tx.send(Work::Skip(*pending_skip));
        *pending_skip = 0;
    }
}

/// Issues a generation order if there is unmet demand and the pipeline window has room (a pending
/// resync skip takes priority and drains the pipeline first). Demand = `pending_batches` shortfall
/// plus `max(refill top-up, prefetch deficit)`, net of quantities already ordered — so pipelining
/// never requests more than the unmet demand. `Err` means the config lock timed out and the
/// dispatcher should stop.
#[allow(clippy::too_many_arguments)]
fn maybe_generate<PB: CorrelatedBatch, E: From<CorrelatedStreamError>>(
    orders_in_flight: &mut VecDeque<usize>,
    max_orders_in_flight: usize,
    skip_in_flight: bool,
    pending_resync: &Option<(usize, oneshot::Sender<Result<(), E>>)>,
    pending_skip: usize,
    config: &Arc<RwLock<BufferConfig>>,
    pending_batches: &VecDeque<(ItemsCollected<PB>, TotalNeeded, BatchSender<PB, E>)>,
    buffer_len: usize,
    prefetch_demand: usize,
    work_tx: &UnboundedSender<Work>,
    log_prefix: &str,
) -> Result<(), E> {
    if skip_in_flight
        || orders_in_flight.len() >= max_orders_in_flight
        || pending_resync.is_some()
        || pending_skip != 0
    {
        return Ok(());
    }
    let batch_shortfall: usize = pending_batches
        .iter()
        .map(|(collected, needed, _)| needed.saturating_sub(collected.len() + buffer_len))
        .sum();
    let buf_after = buffer_len.saturating_sub(batch_shortfall);
    let requested_in_flight: usize = orders_in_flight.iter().sum();
    let need = (batch_shortfall
        + try_read_config(config)?
            .refill_threshold()
            .saturating_sub(buf_after)
            .max(prefetch_demand))
    .saturating_sub(requested_in_flight);
    if need > 0 {
        debug!(
            "{log_prefix} requesting generation of {need} items ({} orders in flight)",
            orders_in_flight.len()
        );
        orders_in_flight.push_back(need);
        let _ = work_tx.send(Work::Generate(need));
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
async fn dispatcher_loop<
    PB: CorrelatedBatch,
    E: From<CorrelatedStreamError> + Clone + Send + std::fmt::Debug,
>(
    mut cmd_rx: UnboundedReceiver<Command<PB, E>>,
    work_tx: UnboundedSender<Work>,
    mut items_rx: UnboundedReceiver<Result<Vec<PB::Item>, E>>,
    mut skip_rx: UnboundedReceiver<Result<(), E>>,
    config: Arc<RwLock<BufferConfig>>,
    shared_position: Arc<AtomicU64>,
    shared_buffered: Arc<AtomicU64>,
    supports_skip: bool,
    max_orders_in_flight: usize,
) {
    let log_prefix = format!("<Dispatcher<{}>>", std::any::type_name::<PB>());

    // Initial capacity is just a hint (the buffer auto-grows on push), clamped so we never
    // pre-allocate a `BufferConfig::UNBOUNDED` capacity (e.g. from `new_on_demand`).
    let initial_cap = try_read_config(&config)
        .map(|c| c.capacity())
        .unwrap_or(0)
        .min(1 << 12);
    let mut buffer: VecDeque<PB::Item> = VecDeque::with_capacity(initial_cap);
    // Newly generated items still owed to outstanding prefetches.
    let mut prefetch_demand: usize = 0;
    // (remaining deficit of newly generated items, completion sender) per outstanding prefetch.
    let mut prefetch_completions: VecDeque<(usize, oneshot::Sender<Result<(), E>>)> =
        VecDeque::new();
    // Sizes of generation orders outstanding at the generator, in issue (FIFO) order. The count
    // is bounded by `max_orders_in_flight` (1 = classic stop-and-wait; pipelined constructors
    // raise it so consecutive orders overlap).
    let mut orders_in_flight: VecDeque<usize> = VecDeque::new();
    // A skip is exclusive: it is only dispatched with no orders outstanding, and blocks new ones.
    let mut skip_in_flight = false;
    // Outstanding `next_n` batch requests in FIFO order: (items collected, total needed, sender).
    let mut pending_batches: VecDeque<(ItemsCollected<PB>, TotalNeeded, BatchSender<PB, E>)> =
        VecDeque::new();
    // The logical position lives solely in `shared_position` (the dispatcher is its only writer):
    // advanced with `fetch_add`, read back with `load` where needed.
    // An in-flight resync awaiting generator skip completion: (deficit being skipped, completion).
    // The deficit (not an absolute target) is added to `position` on completion so concurrent
    // `next_n` increments are preserved.
    let mut pending_resync: Option<(usize, oneshot::Sender<Result<(), E>>)> = None;
    // Remaining elements to skip at the generator for the in-flight resync (not yet dispatched).
    let mut pending_skip: usize = 0;
    // If a config lock acquisition times out, we propagate this error to all consumers.
    let mut shutdown_err: Option<E> = None;
    // Set once the generator has exited, to stop polling a closed skip channel; shutdown is then
    // driven by `items_rx`, which still has the generator's final results (and error) to deliver.
    let mut skip_channel_closed = false;

    // Lock the config with timeout; on failure, record the error and `break` the outer loop.
    macro_rules! lock_cfg {
        () => {
            match try_read_config(&config) {
                Ok(g) => g,
                Err(e) => {
                    error!("{log_prefix} config lock timeout, shutting down");
                    shutdown_err = Some(e.into());
                    break;
                }
            }
        };
    }

    // A generator that fails a skip sends the error on `skip_tx` and then exits, dropping
    // `items_tx` too — so the error and the close of the item channel become ready together and
    // `select!` picks between them at random. Rescue the error before shutting down, so consumers
    // see the real cause rather than the generic `StreamClosed` the shutdown path substitutes.
    macro_rules! absorb_skip_error {
        () => {{
            while let Ok(result) = skip_rx.try_recv() {
                if let Err(e) = result {
                    shutdown_err.get_or_insert(e);
                }
            }
        }};
    }

    // Canonical unmet demand: batch shortfalls + prefetch deficits, net of buffered items.
    macro_rules! outstanding_demand {
        () => {{
            let batch_shortfall: usize = pending_batches
                .iter()
                .map(|(collected, needed, _)| needed.saturating_sub(collected.len()))
                .sum();
            (batch_shortfall + prefetch_demand).saturating_sub(buffer.len())
        }};
    }

    // Publishes the current buffer occupancy for synchronous front-end reads.
    let set_buffered =
        |buffer_len: usize| shared_buffered.store(buffer_len as u64, Ordering::Release);

    loop {
        // At start or after handling an event, (re)issue generation if there is unmet demand.
        if let Err(e) = maybe_generate::<PB, E>(
            &mut orders_in_flight,
            max_orders_in_flight,
            skip_in_flight,
            &pending_resync,
            pending_skip,
            &config,
            &pending_batches,
            buffer.len(),
            prefetch_demand,
            &work_tx,
            &log_prefix,
        ) {
            error!("{log_prefix} config lock timeout, shutting down");
            shutdown_err = Some(e);
            break;
        }

        tokio::select! {
            cmd = cmd_rx.recv() => {
                let Some(cmd) = cmd else {
                    info!("{log_prefix} command channel closed, shutting down");
                    break;
                };
                match cmd {
                    Command::RequestN { n_elements, completion } => {
                        debug!("{log_prefix} batch request for {n_elements} items");
                        if pending_resync.is_some() {
                            // Reaching here means this stream's owner issued a `next_n`
                            // concurrently with its own `resync`, so which side of the barrier the
                            // request lands on is nondeterministic — and `resync` targets an
                            // absolute position, so it cannot be guaranteed to match the side the
                            // other parties put it on. Serving it either way risks a silent
                            // cross-party desync, and parking it would only hide the race, so it
                            // is refused: callers must serialize `next_n` against `resync`.
                            let _ = completion.send(Err(CorrelatedStreamError::ResyncInProgress.into()));
                        } else if outstanding_demand!() + n_elements > lock_cfg!().capacity() {
                            // The config lock is taken only here, on the path that needs the
                            // admission bound.
                            // Rejected: nothing is delivered, so the position does not advance.
                            let _ = completion.send(Err(CorrelatedStreamError::RateLimitExceeded.into()));
                        } else if buffer.len() >= n_elements {
                            // Fully served from buffer — complete immediately.
                            let items: Vec<_> = buffer.drain(..n_elements).collect();
                            shared_position.fetch_add(n_elements as u64, Ordering::Release);
                            set_buffered(buffer.len());
                            let _ = completion.send(Ok(items));
                        } else {
                            // Partially served; need generation for the remainder. The request is
                            // admitted (committed to deliver in FIFO order), so advance now.
                            let collected: Vec<_> = buffer.drain(..).collect();
                            shared_position.fetch_add(n_elements as u64, Ordering::Release);
                            set_buffered(buffer.len());
                            pending_batches.push_back((collected, n_elements, completion));
                        }
                    }
                    Command::Resync { target, completion } => {
                        // Snapshot the single source of truth for this arm's checks.
                        let position = shared_position.load(Ordering::Relaxed);
                        debug!("{log_prefix} resync to {target} (position {position})");
                        if pending_resync.is_some() {
                            // Reject overlapping resyncs: queueing them would only widen the
                            // window for another desync. The caller must serialize resyncs.
                            let _ = completion.send(Err(CorrelatedStreamError::ResyncInProgress.into()));
                        } else if target < position {
                            let _ = completion.send(Err(CorrelatedStreamError::ResyncRewind {
                                current: position,
                                target,
                            }
                            .into()));
                        } else {
                            // Discard already-buffered elements first (free, local).
                            let skip = (target - position) as usize;
                            let drained = skip.min(buffer.len());
                            buffer.drain(..drained);
                            shared_position.fetch_add(drained as u64, Ordering::Release);
                            set_buffered(buffer.len());
                            let deficit = skip - drained;
                            if deficit == 0 {
                                let _ = completion.send(Ok(()));
                            } else if !supports_skip {
                                // By the lockstep argument, an interactive generator should always
                                // have the target buffered; a deficit signals a generation desync.
                                let _ = completion.send(Err(CorrelatedStreamError::ResyncUnsupported {
                                    generated: position + drained as u64,
                                    target,
                                }
                                .into()));
                            } else {
                                // Skip the remainder at the generator, then finalize on completion.
                                // We stash the deficit (not `target`) so completion advances the
                                // position relatively — defense in depth, even though `next_n` is
                                // rejected while this resync is pending so no concurrent
                                // increments can occur.
                                pending_resync = Some((deficit, completion));
                                pending_skip = deficit;
                                maybe_skip(&mut pending_skip, &mut skip_in_flight, orders_in_flight.len(), &work_tx, &log_prefix);
                            }
                        }
                    }
                    Command::Prefetch { n_elements, completion } => {
                        debug!("{log_prefix} prefetch {n_elements} items");
                        if buffer.len() >= n_elements {
                            // Buffer already covers the demand (includes n_elements == 0).
                            let _ = completion.send(Ok(()));
                        } else {
                            let cap = lock_cfg!().capacity();
                            if outstanding_demand!() + n_elements > cap {
                                let _ = completion.send(Err(CorrelatedStreamError::RateLimitExceeded.into()));
                            } else {
                                let deficit = n_elements - buffer.len();
                                prefetch_demand += deficit;
                                prefetch_completions.push_back((deficit, completion));
                            }
                        }
                    }
                }
            }

            result = items_rx.recv() => {
                let Some(result) = result else {
                    info!("{log_prefix} generator channel closed, shutting down");
                    absorb_skip_error!();
                    break;
                };
                match result {
                    Ok(items) => {
                        // Every batch settles exactly one outstanding order, for at least the
                        // quantity ordered. Violating either means the stream's idea of the
                        // generator's position no longer matches the generator's own — for a
                        // dealer stream, a silent desync from the other parties. Checked in
                        // release too: shutting down (which fails every consumer) beats that.
                        let Some(ordered) = orders_in_flight.pop_front() else {
                            error!("{log_prefix} received {} items with no order in flight, shutting down", items.len());
                            break;
                        };
                        if items.len() < ordered {
                            error!("{log_prefix} order for {ordered} items resolved with only {}, shutting down", items.len());
                            break;
                        }
                        let generated = items.len();
                        debug!("{log_prefix} received {generated} items (pending_batches: {}, buffer: {})",
                               pending_batches.iter().map(|(c, n, _)| n - c.len()).sum::<usize>(),
                               buffer.len());
                        let mut iter = items.into_iter();
                        // 1. Fill outstanding batch requests in FIFO order (admitted before the
                        //    resync, if any, so their positions precede its target).
                        while let Some((ref mut collected, needed, _)) = pending_batches.front_mut() {
                            let shortfall = *needed - collected.len();
                            collected.extend(iter.by_ref().take(shortfall));
                            if collected.len() < *needed { break; } // still waiting
                            let (collected, _, tx) = pending_batches.pop_front().unwrap();
                            let _ = tx.send(Ok(collected));
                        }
                        // 2. While a resync awaits its generator skip, leftover items below the
                        //    target are burned against the remaining deficit — buffering them
                        //    would deliver pre-target elements after the skip. Items beyond the
                        //    deficit are past the target and fall through to the buffer. Burning
                        //    can cover the whole deficit, completing the resync with no skip.
                        //    `pending_skip > 0` is implied (a zero deficit completes inline in the
                        //    `Resync` arm) but asserted here so completing the resync below always
                        //    follows an actual burn, never an empty batch.
                        if pending_resync.is_some() && !skip_in_flight && pending_skip > 0 {
                            let burn = pending_skip.min(iter.len());
                            iter.by_ref().take(burn).for_each(drop);
                            pending_skip -= burn;
                            if pending_skip == 0 {
                                if let Some((deficit, tx)) = pending_resync.take() {
                                    shared_position.fetch_add(deficit as u64, Ordering::Release);
                                    let _ = tx.send(Ok(()));
                                }
                            }
                        }
                        // 3. Buffer the rest; capacity is an admission bound, not a storage bound.
                        buffer.extend(iter);
                        set_buffered(buffer.len());
                        // 4. Credit all newly generated items against prefetch deficits (FIFO).
                        //    `generated` is the whole order, including anything step 2 burned
                        //    against a resync deficit, so prefetch completions are best-effort
                        //    across a resync: a prefetch can resolve `Ok` with fewer (even zero)
                        //    items actually buffered. That only costs the next `next_n` a wait —
                        //    crediting the burns to the *next* prefetch instead would leave this
                        //    one outstanding forever, since a resync also cancels the demand.
                        prefetch_demand = prefetch_demand.saturating_sub(generated);
                        let mut credit = generated;
                        while credit > 0 {
                            let Some((deficit, _)) = prefetch_completions.front_mut() else { break; };
                            let used = (*deficit).min(credit);
                            *deficit -= used;
                            credit -= used;
                            if *deficit > 0 { break; }
                            let (_, tx) = prefetch_completions.pop_front().unwrap();
                            let _ = tx.send(Ok(()));
                        }
                        // A resync may have been waiting for the pipeline to drain.
                        maybe_skip(&mut pending_skip, &mut skip_in_flight, orders_in_flight.len(), &work_tx, &log_prefix);
                    }
                    Err(e) => {
                        error!("{log_prefix} generation error, shutting down: {e:?}");
                        for (_, _, tx) in pending_batches.drain(..) { let _ = tx.send(Err(e.clone())); }
                        for (_, tx) in prefetch_completions.drain(..) { let _ = tx.send(Err(e.clone())); }
                        if let Some((_, tx)) = pending_resync.take() { let _ = tx.send(Err(e.clone())); }
                        return;
                    }
                }
            }

            result = skip_rx.recv(), if !skip_channel_closed => {
                let Some(result) = result else {
                    // The generator exited. Don't shut down here: it may have sent a final error
                    // (and earlier, still-unprocessed results) on `items_rx`, and both channels
                    // become ready together, so shutting down on whichever `select!` happens to
                    // pick would discard them. Disable this branch and let `items_rx` drain —
                    // an mpsc yields every buffered message before reporting closure, so the
                    // shutdown happens there with the real cause and nothing is lost.
                    info!("{log_prefix} skip channel closed, draining pending results");
                    skip_channel_closed = true;
                    continue;
                };
                skip_in_flight = false;
                match result {
                    Ok(()) => {
                        // The generator advanced its position; finalize the in-flight resync.
                        // Advancing by the deficit (rather than assigning the target) is defense
                        // in depth — `next_n` is rejected while a resync is pending, so no
                        // concurrent increments can have landed.
                        if let Some((deficit, tx)) = pending_resync.take() {
                            shared_position.fetch_add(deficit as u64, Ordering::Release);
                            let _ = tx.send(Ok(()));
                        }
                    }
                    Err(e) => {
                        error!("{log_prefix} skip error, shutting down: {e:?}");
                        for (_, _, tx) in pending_batches.drain(..) { let _ = tx.send(Err(e.clone())); }
                        for (_, tx) in prefetch_completions.drain(..) { let _ = tx.send(Err(e.clone())); }
                        if let Some((_, tx)) = pending_resync.take() { let _ = tx.send(Err(e.clone())); }
                        return;
                    }
                }
            }
        }
    }

    // Graceful shutdown: resolve all outstanding consumers/handles with the recorded error
    // (if shutdown was triggered by a lock timeout) or `StreamClosed` otherwise.
    let final_err: E = shutdown_err.unwrap_or_else(|| CorrelatedStreamError::StreamClosed.into());
    for (_, _, tx) in pending_batches.drain(..) {
        let _ = tx.send(Err(final_err.clone()));
    }
    for (_, tx) in prefetch_completions.drain(..) {
        let _ = tx.send(Err(final_err.clone()));
    }
    if let Some((_, tx)) = pending_resync.take() {
        let _ = tx.send(Err(final_err.clone()));
    }
}

impl<
        PB: CorrelatedBatch,
        E: From<CorrelatedStreamError> + Clone + Send + std::fmt::Debug + 'static,
    > CorrelatedStream<PB::Item> for BufferedStream<PB, E>
{
    type Error = E;

    fn next_n(&self, n_elements: usize) -> Result<NextVec<PB::Item, E>, CorrelatedStreamError> {
        if n_elements == 0 {
            return Ok(NextVec::default());
        }
        let max_allowed = self.max_request_size()?;
        if n_elements > max_allowed {
            return Err(CorrelatedStreamError::RequestTooLarge {
                requested: n_elements,
                max_allowed,
            });
        }
        let (tx, rx) = oneshot::channel();
        self.command_sender
            .send(Command::RequestN {
                n_elements,
                completion: tx,
            })
            .map_err(|e| CorrelatedStreamError::SendError(e.to_string()))?;
        Ok(NextVec {
            future: Next(rx),
            size: n_elements,
        })
    }

    fn prefetch_n(&self, n_elements: usize) -> PrefetchHandle<E> {
        let (tx, rx) = oneshot::channel();
        let max = match self.max_request_size() {
            Ok(m) => m,
            Err(e) => {
                let _ = tx.send(Err(e.into()));
                return PrefetchHandle::from(rx);
            }
        };
        if n_elements > max {
            let _ = tx.send(Err(CorrelatedStreamError::RequestTooLarge {
                requested: n_elements,
                max_allowed: max,
            }
            .into()));
            return PrefetchHandle::from(rx);
        }
        // If the dispatcher is gone, resolve the handle immediately.
        let cmd = Command::Prefetch {
            n_elements,
            completion: tx,
        };
        if let Err(e) = self.command_sender.send(cmd) {
            if let Command::Prefetch { completion, .. } = e.0 {
                let _ = completion.send(Err(CorrelatedStreamError::StreamClosed.into()));
            }
        }
        PrefetchHandle::from(rx)
    }

    fn position(&self) -> u64 {
        self.position.load(Ordering::Acquire)
    }

    fn buffered(&self) -> u64 {
        self.buffered.load(Ordering::Acquire)
    }

    fn resync(&self, target: u64) -> ResyncHandle<E> {
        let (tx, rx) = oneshot::channel();
        let cmd = Command::Resync {
            target,
            completion: tx,
        };
        if let Err(e) = self.command_sender.send(cmd) {
            if let Command::Resync { completion, .. } = e.0 {
                let _ = completion.send(Err(CorrelatedStreamError::StreamClosed.into()));
            }
        }
        ResyncHandle::from(rx)
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{
            atomic::{AtomicUsize, Ordering},
            Arc,
        },
        time::Duration,
    };

    use rand::{rngs::StdRng, SeedableRng};
    use typenum::U2;

    use crate::{
        algebra::elliptic_curve::{Curve25519Ristretto, ScalarField},
        correlated_randomness::{
            generator::CorrelationGenerator,
            singlets::{Singlet, Singlets},
            stream::{
                buffered::{Buffer, BufferConfig, BufferedStream},
                errors::CorrelatedStreamError,
                CorrelatedStream,
            },
        },
        random::Random,
        utils::TryFuture,
    };

    type Fq = ScalarField<Curve25519Ristretto>;
    type TestPB = Singlets<Fq, U2>;
    type TestItem = Singlet<Fq>;
    type TestErr = CorrelatedStreamError;

    // -----------------------
    // ===== Mock generator
    // -----------------------

    #[derive(Clone)]
    struct MockGenConfig {
        /// Sleep for this duration on every `run_for` call to simulate generation latency.
        delay: Duration,
        /// If `Some(threshold)`, fail (return `StreamClosed`) once this many items have been
        /// generated cumulatively.
        fail_after: Option<usize>,
    }

    impl Default for MockGenConfig {
        fn default() -> Self {
            Self {
                delay: Duration::from_millis(0),
                fail_after: None,
            }
        }
    }

    struct MockGen {
        rng: StdRng,
        cfg: MockGenConfig,
        /// Total items produced (visible to tests via `Arc`).
        items_produced: Arc<AtomicUsize>,
        /// Number of `run_for` calls made (i.e. number of generation batches).
        batches: Arc<AtomicUsize>,
    }

    impl MockGen {
        fn new(cfg: MockGenConfig) -> (Self, Arc<AtomicUsize>, Arc<AtomicUsize>) {
            let items_produced = Arc::new(AtomicUsize::new(0));
            let batches = Arc::new(AtomicUsize::new(0));
            let gen = Self {
                rng: StdRng::from_seed([0u8; 32]),
                cfg,
                items_produced: items_produced.clone(),
                batches: batches.clone(),
            };
            (gen, items_produced, batches)
        }
    }

    impl CorrelationGenerator<TestPB> for MockGen {
        type Net = ();
        type Error = TestErr;

        fn run(&mut self, _net: &mut ()) -> impl TryFuture<Ok = TestPB, Error = Self::Error> {
            async move { Err(CorrelatedStreamError::StreamClosed) }
        }

        fn run_for(
            &mut self,
            n: usize,
            _net: &mut (),
        ) -> impl TryFuture<Ok = Vec<TestItem>, Error = Self::Error> {
            async move {
                self.batches.fetch_add(1, Ordering::SeqCst);
                if self.cfg.delay > Duration::ZERO {
                    tokio::time::sleep(self.cfg.delay).await;
                }
                if let Some(threshold) = self.cfg.fail_after {
                    if self.items_produced.load(Ordering::SeqCst) + n > threshold {
                        return Err(CorrelatedStreamError::StreamClosed);
                    }
                }
                let items: Vec<TestItem> = (0..n)
                    .map(|_| {
                        Singlet::<Fq>::random_n::<Vec<_>>(&mut self.rng, 1)
                            .into_iter()
                            .next()
                            .unwrap()
                    })
                    .collect();
                self.items_produced.fetch_add(n, Ordering::SeqCst);
                Ok(items)
            }
        }
    }

    fn make_stream(
        cfg: MockGenConfig,
        buf_cfg: BufferConfig,
    ) -> (
        BufferedStream<TestPB, TestErr>,
        Arc<AtomicUsize>,
        Arc<AtomicUsize>,
    ) {
        let (gen, produced, batches) = MockGen::new(cfg);
        let stream = BufferedStream::<TestPB, TestErr>::new(gen, (), buf_cfg);
        (stream, produced, batches)
    }

    // ----------------------------------------------------------------------
    // ===== Skip-capable mock generator (dealer-like, deterministic-per-index)
    // ----------------------------------------------------------------------

    /// A generator that advertises [`CorrelationGenerator::SUPPORTS_UNILATERAL_SKIP`] and
    /// implements `skip` by advancing its produced counter without emitting items.
    struct SkipMockGen {
        rng: StdRng,
        produced: Arc<AtomicUsize>,
        skipped: Arc<AtomicUsize>,
    }

    impl SkipMockGen {
        fn new() -> (Self, Arc<AtomicUsize>, Arc<AtomicUsize>) {
            let produced = Arc::new(AtomicUsize::new(0));
            let skipped = Arc::new(AtomicUsize::new(0));
            let gen = Self {
                rng: StdRng::from_seed([7u8; 32]),
                produced: produced.clone(),
                skipped: skipped.clone(),
            };
            (gen, produced, skipped)
        }
    }

    impl CorrelationGenerator<TestPB> for SkipMockGen {
        type Net = ();
        type Error = TestErr;

        const SUPPORTS_UNILATERAL_SKIP: bool = true;

        fn run(&mut self, _net: &mut ()) -> impl TryFuture<Ok = TestPB, Error = Self::Error> {
            async move { Err(CorrelatedStreamError::StreamClosed) }
        }

        fn run_for(
            &mut self,
            n: usize,
            _net: &mut (),
        ) -> impl TryFuture<Ok = Vec<TestItem>, Error = Self::Error> {
            async move {
                let items: Vec<TestItem> = Singlet::<Fq>::random_n::<Vec<_>>(&mut self.rng, n);
                self.produced.fetch_add(n, Ordering::SeqCst);
                Ok(items)
            }
        }

        fn skip(
            &mut self,
            n: usize,
            _net: &mut (),
        ) -> impl TryFuture<Ok = (), Error = Self::Error> {
            async move {
                // Advance the logical position by `n` without materializing the elements.
                self.skipped.fetch_add(n, Ordering::SeqCst);
                self.produced.fetch_add(n, Ordering::SeqCst);
                Ok(())
            }
        }
    }

    /// A minimal generator with the default (false) skip capability, for testing the
    /// `ResyncUnsupported` path when a deficit cannot be covered from the buffer.
    struct NoSkipMockGen {
        rng: StdRng,
    }

    impl CorrelationGenerator<TestPB> for NoSkipMockGen {
        type Net = ();
        type Error = TestErr;

        fn run(&mut self, _net: &mut ()) -> impl TryFuture<Ok = TestPB, Error = Self::Error> {
            async move { Err(CorrelatedStreamError::StreamClosed) }
        }

        fn run_for(
            &mut self,
            n: usize,
            _net: &mut (),
        ) -> impl TryFuture<Ok = Vec<TestItem>, Error = Self::Error> {
            async move { Ok(Singlet::<Fq>::random_n::<Vec<_>>(&mut self.rng, n)) }
        }
    }

    // -----------------------
    // ===== Tests
    // -----------------------

    #[tokio::test]
    async fn next_n_resolves_batch_future() {
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::eager(16));
        let fut = stream.next_n(7).expect("request accepted");
        let items = fut.await.expect("batch resolves");
        assert_eq!(items.len(), 7);
    }

    #[tokio::test]
    async fn request_too_large_rejected() {
        let (stream, _, _) = make_stream(
            MockGenConfig::default(),
            BufferConfig::eager_with(8, 4), // capacity=8, max_request_size=4
        );
        match stream.next_n(5) {
            Err(CorrelatedStreamError::RequestTooLarge {
                requested: 5,
                max_allowed: 4,
            }) => {}
            Ok(_) => panic!("must reject n>max"),
            Err(e) => panic!("unexpected error: {e:?}"),
        }
    }

    #[tokio::test]
    async fn rate_limit_when_exceeding_capacity() {
        // Small capacity + slow generator → the first batch's shortfall (4) plus the second
        // request (4) exceeds capacity (4), tripping the admission bound.
        let cfg = MockGenConfig {
            delay: Duration::from_millis(200),
            ..Default::default()
        };
        let (stream, _, _) = make_stream(cfg, BufferConfig::eager(4));
        let _f1 = stream.next_n(4).unwrap();
        // Give the dispatcher a moment to register the first request before issuing the second.
        tokio::time::sleep(Duration::from_millis(20)).await;
        let f2 = stream.next_n(4).unwrap();
        let results = futures::future::join_all(f2).await;
        assert!(
            results
                .iter()
                .all(|r| matches!(r, Err(CorrelatedStreamError::RateLimitExceeded))),
            "expected all four to be rate-limited, got {results:?}"
        );
    }

    #[tokio::test]
    async fn prefetch_completes_and_serves_subsequent_requests_quickly() {
        let cfg = MockGenConfig {
            delay: Duration::from_millis(100),
            ..Default::default()
        };
        let (stream, produced, _) = make_stream(cfg, BufferConfig::eager(32));
        let handle = stream.prefetch_n(10);
        handle.await.expect("prefetch completes");
        assert!(produced.load(Ordering::SeqCst) >= 10);
        // Subsequent request should be served from buffer instantaneously
        // (or at least without waiting for another slow generation cycle).
        let start = std::time::Instant::now();
        let items = stream.next_n(10).unwrap().await.expect("served");
        assert_eq!(items.len(), 10);
        assert!(
            start.elapsed() < Duration::from_millis(80),
            "request should be served from prefetched buffer (took {:?})",
            start.elapsed()
        );
    }

    #[tokio::test]
    async fn sequential_prefetches_all_resolve() {
        // Regression: with refill_threshold=0 the first prefetch fills the buffer; the second
        // must resolve immediately from the buffer instead of waiting for a generation target
        // that never triggers.
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::lazy(16, 0));
        for i in 0..2 {
            tokio::time::timeout(Duration::from_secs(1), stream.prefetch_n(4))
                .await
                .unwrap_or_else(|_| panic!("prefetch {i} timed out"))
                .expect("prefetch completes");
        }
        // A larger prefetch only partially covered by the buffer must also resolve.
        tokio::time::timeout(Duration::from_secs(1), stream.prefetch_n(8))
            .await
            .expect("partially covered prefetch timed out")
            .expect("prefetch completes");
    }

    #[tokio::test]
    async fn generator_error_propagates_to_pending_consumers() {
        let cfg = MockGenConfig {
            delay: Duration::from_millis(20),
            fail_after: Some(0), // fail on first batch
        };
        let (stream, _, _) = make_stream(cfg, BufferConfig::eager(16));
        let futs = stream.next_n(4).unwrap();
        let results = futures::future::join_all(futs).await;
        assert!(
            results
                .iter()
                .all(|r| matches!(r, Err(CorrelatedStreamError::StreamClosed))),
            "all consumers should receive the generator error"
        );
    }

    #[tokio::test]
    async fn fifo_order_across_two_batches() {
        // Two requests issued back-to-back must be served in submission order.
        let cfg = MockGenConfig {
            delay: Duration::from_millis(40),
            ..Default::default()
        };
        let (stream, _, batches) = make_stream(cfg, BufferConfig::eager(32));
        let f1 = stream.next_n(3).unwrap();
        let f2 = stream.next_n(3).unwrap();
        let (a, b) = tokio::join!(f1, f2);
        let a = a.expect("first batch resolves");
        let b = b.expect("second batch resolves");
        assert_eq!(a.len(), 3);
        assert_eq!(b.len(), 3);
        // At least one batch should have been generated (count is implementation-defined).
        assert!(batches.load(Ordering::SeqCst) >= 1);
    }

    #[tokio::test]
    async fn buffer_config_setters_are_visible() {
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::eager(16));
        assert_eq!(stream.capacity().unwrap(), 16);
        stream.set_capacity(32).unwrap();
        assert_eq!(stream.capacity().unwrap(), 32);
    }

    // ── next_n validation
    // ──────────────────────────────────────────────────────────────────────

    #[tokio::test]
    async fn next_n_rejects_too_large() {
        let (stream, _, _) = make_stream(
            MockGenConfig::default(),
            BufferConfig::eager_with(8, 4), // capacity=8, max_request_size=4
        );
        match stream.next_n(5) {
            Err(CorrelatedStreamError::RequestTooLarge {
                requested: 5,
                max_allowed: 4,
            }) => {}
            Ok(_) => panic!("must reject n > max"),
            Err(e) => panic!("unexpected error: {e:?}"),
        }
    }

    #[tokio::test]
    async fn next_n_rate_limit_when_exceeding_capacity() {
        // Slow generator → first batch's shortfall is outstanding; a second batch whose demand
        // pushes the aggregate shortfall over capacity must be rejected with `RateLimitExceeded`.
        let cfg = MockGenConfig {
            delay: Duration::from_millis(200),
            ..Default::default()
        };
        let (stream, _, _) = make_stream(cfg, BufferConfig::eager(4));
        let _f1 = stream.next_n(4).unwrap();
        // Let the dispatcher register the first RequestBatch before issuing the second.
        tokio::time::sleep(Duration::from_millis(20)).await;
        let f2 = stream.next_n(4).unwrap();
        assert!(
            matches!(f2.await, Err(CorrelatedStreamError::RateLimitExceeded)),
            "second next_n should be rate-limited"
        );
    }

    #[tokio::test]
    async fn next_n_admitted_when_shortfall_fits_capacity() {
        // capacity=8, refill_threshold=0: a pending next_n(4) leaves a shortfall of 4, so a
        // concurrent next_n(4) (aggregate 8) fits exactly and must be admitted, not rejected.
        let cfg = MockGenConfig {
            delay: Duration::from_millis(50),
            ..Default::default()
        };
        let (stream, _, _) = make_stream(cfg, BufferConfig::lazy(8, 0));
        let f1 = stream.next_n(4).unwrap();
        tokio::time::sleep(Duration::from_millis(20)).await;
        let f2 = stream.next_n(4).unwrap();
        let (a, b) = tokio::join!(f1, f2);
        assert_eq!(a.expect("first batch resolves").len(), 4);
        assert_eq!(b.expect("second batch resolves").len(), 4);
    }

    #[tokio::test]
    async fn next_n_error_propagates() {
        // Generator always fails; the BatchFuture must resolve with the generator error.
        let cfg = MockGenConfig {
            delay: Duration::from_millis(20), // ensure batch is registered before failure
            fail_after: Some(0),
        };
        let (stream, _, _) = make_stream(cfg, BufferConfig::eager(16));
        let result = stream.next_n(4).unwrap().await;
        assert!(
            matches!(result, Err(CorrelatedStreamError::StreamClosed)),
            "expected generator error to propagate through BatchFuture, got {result:?}"
        );
    }

    // ── refill_threshold ───────────────────────────────────────────────────────────────────────

    #[tokio::test]
    async fn refill_threshold_drives_proactive_generation() {
        // lazy buffer: capacity=16, refill_threshold=8 → after a 4-item request, the dispatcher
        // should top up to 8 items in the buffer.
        let (stream, produced, _) =
            make_stream(MockGenConfig::default(), BufferConfig::lazy(16, 8));
        let _ = stream.next_n(4).unwrap().await.expect("served");
        // Allow the dispatcher to settle (post-serve refill cycle).
        tokio::time::sleep(Duration::from_millis(50)).await;
        // We requested 4 items; refill_threshold=8 means the buffer should hold at least 8
        // additional items beyond what was served, i.e. >= 12 total generated.
        let total = produced.load(Ordering::SeqCst);
        assert!(total >= 12, "expected >= 12 items generated, got {total}");
    }

    #[tokio::test]
    async fn refill_threshold_drives_proactive_generation_before_requests() {
        let (_stream, produced, _) =
            make_stream(MockGenConfig::default(), BufferConfig::lazy(16, 8));
        // Allow the dispatcher to settle
        tokio::time::sleep(Duration::from_millis(50)).await;
        // refill_threshold=8 means the buffer should hold at least 8
        let total = produced.load(Ordering::SeqCst);
        assert!(total >= 8, "expected >= 8 items generated, got {total}");
    }

    // ── position counter & resync ────────────────────────────────────────────────────────────

    #[tokio::test]
    async fn position_tracks_delivered_not_prefetched() {
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::eager(32));
        // Prefetching fills the buffer but must NOT advance the logical position.
        stream.prefetch_n(10).await.expect("prefetch completes");
        assert_eq!(stream.position(), 0, "prefetch must not advance position");
        // Delivering through next_n advances it by exactly the requested amount.
        let _ = stream.next_n(7).unwrap().await.expect("served");
        assert_eq!(stream.position(), 7);
        let _ = stream.next_n(3).unwrap().await.expect("served");
        assert_eq!(stream.position(), 10);
    }

    #[tokio::test]
    async fn buffered_tracks_ready_elements() {
        // Lazy buffer with no proactive refill, so occupancy only changes on prefetch/next_n.
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::lazy(32, 0));
        assert_eq!(stream.buffered(), 0);
        stream.prefetch_n(10).await.expect("prefetch completes");
        assert_eq!(stream.buffered(), 10, "prefetch fills the buffer");
        let _ = stream.next_n(4).unwrap().await.expect("served");
        assert_eq!(stream.buffered(), 6, "delivery drains the buffer");
    }

    #[tokio::test]
    async fn position_does_not_advance_on_rejected_request() {
        // Slow generator + tiny capacity → the second request is rate-limited and delivers
        // nothing, so the position must stay at the first request's amount.
        let cfg = MockGenConfig {
            delay: Duration::from_millis(200),
            ..Default::default()
        };
        let (stream, _, _) = make_stream(cfg, BufferConfig::eager(4));
        let _f1 = stream.next_n(4).unwrap();
        tokio::time::sleep(Duration::from_millis(20)).await;
        let f2 = stream.next_n(4).unwrap();
        assert!(matches!(
            f2.await,
            Err(CorrelatedStreamError::RateLimitExceeded)
        ));
        // Only the admitted first request counts.
        assert_eq!(stream.position(), 4);
    }

    #[tokio::test]
    async fn resync_drains_buffer_and_advances_position() {
        // Lazy buffer with no proactive refill, so the only generation is the explicit prefetch;
        // this isolates resync's drain from background top-ups.
        let (stream, produced, _) =
            make_stream(MockGenConfig::default(), BufferConfig::lazy(32, 0));
        stream.prefetch_n(10).await.expect("prefetch completes");
        let before = produced.load(Ordering::SeqCst);
        // Target is fully covered by the buffer: pure local drain, no extra generation.
        stream.resync(6).await.expect("resync completes");
        assert_eq!(stream.position(), 6);
        assert_eq!(
            produced.load(Ordering::SeqCst),
            before,
            "drain-only resync must not generate"
        );
        // Subsequent delivery continues from the resynced position.
        let _ = stream.next_n(2).unwrap().await.expect("served");
        assert_eq!(stream.position(), 8);
    }

    #[tokio::test]
    async fn resync_noop_when_already_at_target() {
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::eager(16));
        let _ = stream.next_n(5).unwrap().await.expect("served");
        stream.resync(5).await.expect("no-op resync completes");
        assert_eq!(stream.position(), 5);
    }

    #[tokio::test]
    async fn resync_rewind_is_rejected() {
        let (stream, _, _) = make_stream(MockGenConfig::default(), BufferConfig::eager(16));
        let _ = stream.next_n(5).unwrap().await.expect("served");
        match stream.resync(3).await {
            Err(CorrelatedStreamError::ResyncRewind {
                current: 5,
                target: 3,
            }) => {}
            other => panic!("expected ResyncRewind, got {other:?}"),
        }
        // Position is unchanged after a rejected rewind.
        assert_eq!(stream.position(), 5);
    }

    #[tokio::test]
    async fn resync_unsupported_when_deficit_and_no_skip() {
        // Lazy buffer with no proactive refill so nothing is buffered; a resync past the buffer
        // needs generator skipping, which this generator does not support.
        let gen = NoSkipMockGen {
            rng: StdRng::from_seed([0u8; 32]),
        };
        let stream = BufferedStream::<TestPB, TestErr>::new(gen, (), BufferConfig::lazy(16, 0));
        match stream.resync(10).await {
            Err(CorrelatedStreamError::ResyncUnsupported {
                generated: 0,
                target: 10,
            }) => {}
            other => panic!("expected ResyncUnsupported, got {other:?}"),
        }
        // Nothing was delivered/skipped, so the position stays put.
        assert_eq!(stream.position(), 0);
    }

    #[tokio::test]
    async fn resync_skips_at_generator_when_supported() {
        let (gen, produced, skipped) = SkipMockGen::new();
        // Lazy buffer, no proactive refill: the resync deficit must be skipped at the generator.
        let stream = BufferedStream::<TestPB, TestErr>::new(gen, (), BufferConfig::lazy(64, 0));
        stream.resync(10).await.expect("resync via skip completes");
        assert_eq!(stream.position(), 10);
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            10,
            "deficit should be skipped"
        );
        assert_eq!(
            produced.load(Ordering::SeqCst),
            10,
            "skip advances the generator without delivering items"
        );
        // Delivery resumes correctly from the skipped-to position.
        let items = stream.next_n(3).unwrap().await.expect("served");
        assert_eq!(items.len(), 3);
        assert_eq!(stream.position(), 13);
    }

    #[tokio::test]
    async fn resync_partial_buffer_then_skip_remainder() {
        let (gen, produced, skipped) = SkipMockGen::new();
        let stream = BufferedStream::<TestPB, TestErr>::new(gen, (), BufferConfig::lazy(64, 0));
        // Buffer 4 elements, then resync past them: 4 drained locally, 6 skipped at generator.
        stream.prefetch_n(4).await.expect("prefetch completes");
        let produced_after_prefetch = produced.load(Ordering::SeqCst);
        assert_eq!(produced_after_prefetch, 4);
        stream.resync(10).await.expect("resync completes");
        assert_eq!(stream.position(), 10);
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            6,
            "only the unbuffered remainder is skipped at the generator"
        );
    }

    // ----------------------------------------------------------------------
    // ===== Pipelined mock generator (dealer-client-like)
    // ----------------------------------------------------------------------

    use std::{collections::VecDeque, sync::Mutex};

    use futures::FutureExt;
    use tokio::sync::oneshot;

    use crate::correlated_randomness::generator::PipelinedCorrelationGenerator;

    /// Pending orders' release handles, in issue order. Each carries the outcome the order should
    /// resolve with, so a test can fail one nominated order and leave the rest pending.
    type GateQueue = VecDeque<oneshot::Sender<Result<(), TestErr>>>;
    type Gates = Arc<Mutex<GateQueue>>;

    /// Issues orders synchronously but resolves them only when the test releases the matching
    /// gate, so tests can observe how many orders are in flight at once and control the order in
    /// which their results become available — and whether each succeeds or fails.
    struct PipelinedMockGen {
        rng: StdRng,
        /// Sizes of all orders issued so far, in issue order.
        orders: Arc<Mutex<Vec<usize>>>,
        /// Release handles for pending orders, in issue order.
        gates: Gates,
        skipped: Arc<AtomicUsize>,
    }

    impl PipelinedMockGen {
        #[allow(clippy::type_complexity)]
        fn new() -> (Self, Arc<Mutex<Vec<usize>>>, Gates, Arc<AtomicUsize>) {
            let orders = Arc::new(Mutex::new(Vec::new()));
            let gates = Arc::new(Mutex::new(VecDeque::new()));
            let skipped = Arc::new(AtomicUsize::new(0));
            let gen = Self {
                rng: StdRng::from_seed([3u8; 32]),
                orders: orders.clone(),
                gates: gates.clone(),
                skipped: skipped.clone(),
            };
            (gen, orders, gates, skipped)
        }
    }

    /// Pops the oldest pending gate and resolves its order successfully.
    fn release_next(gates: &Mutex<GateQueue>) {
        let gate = gates
            .lock()
            .unwrap()
            .pop_front()
            .expect("an order should be pending");
        let _ = gate.send(Ok(()));
    }

    /// The error `fail_next` makes an order resolve with. Deliberately a variant no drop path
    /// produces — a dropped batch sender yields `RecvError` and a dropped prefetch sender
    /// `StreamClosed` — so a test can tell "the dispatcher forwarded the generator's error" apart
    /// from "the dispatcher just went away".
    const MOCK_ORDER_FAILURE: &str = "mock order failure";

    /// Pops the oldest pending gate and fails its order, leaving any later ones pending.
    fn fail_next(gates: &Mutex<GateQueue>) {
        let gate = gates
            .lock()
            .unwrap()
            .pop_front()
            .expect("an order should be pending");
        let _ = gate.send(Err(CorrelatedStreamError::SendError(
            MOCK_ORDER_FAILURE.to_string(),
        )));
    }

    /// Polls `cond` until it holds or ~1s of *virtual* time elapses.
    ///
    /// Every pipelined test runs with `start_paused = true`, so these sleeps (and the fixed sleeps
    /// the negative assertions use) consume no wall-clock time and are not races: tokio only
    /// auto-advances its clock once every task is idle, so a sleep resolves exactly when the
    /// dispatcher and generator have nothing left to do. "Nothing happened after 50ms" therefore
    /// means "nothing happened once the system quiesced", on a loaded CI machine as much as
    /// anywhere else.
    async fn wait_until(mut cond: impl FnMut() -> bool) -> bool {
        for _ in 0..200 {
            if cond() {
                return true;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
        cond()
    }

    impl CorrelationGenerator<TestPB> for PipelinedMockGen {
        type Net = ();
        type Error = TestErr;

        const SUPPORTS_UNILATERAL_SKIP: bool = true;

        fn run(&mut self, _net: &mut ()) -> impl TryFuture<Ok = TestPB, Error = Self::Error> {
            async move { Err(CorrelatedStreamError::StreamClosed) }
        }

        fn run_for(
            &mut self,
            n: usize,
            net: &mut (),
        ) -> impl TryFuture<Ok = Vec<TestItem>, Error = Self::Error> {
            self.issue_for(n, net)
        }

        fn skip(
            &mut self,
            n: usize,
            _net: &mut (),
        ) -> impl TryFuture<Ok = (), Error = Self::Error> {
            self.skipped.fetch_add(n, Ordering::SeqCst);
            async move { Ok(()) }
        }
    }

    impl PipelinedCorrelationGenerator<TestPB> for PipelinedMockGen {
        fn issue_for(
            &mut self,
            n: usize,
            _net: &mut (),
        ) -> futures::future::BoxFuture<'static, Result<Vec<TestItem>, Self::Error>> {
            let items: Vec<TestItem> = Singlet::<Fq>::random_n::<Vec<_>>(&mut self.rng, n);
            let (tx, rx) = oneshot::channel::<Result<(), TestErr>>();
            self.orders.lock().unwrap().push(n);
            self.gates.lock().unwrap().push_back(tx);
            async move {
                // A dropped gate resolves the order successfully: only an explicit `fail_next`
                // makes an order fail.
                match rx.await {
                    Ok(Err(e)) => Err(e),
                    _ => Ok(items),
                }
            }
            .boxed()
        }
    }

    #[allow(clippy::type_complexity)]
    fn make_pipelined_stream(
        depth: usize,
    ) -> (
        BufferedStream<TestPB, TestErr>,
        Arc<Mutex<Vec<usize>>>,
        Gates,
        Arc<AtomicUsize>,
    ) {
        let (gen, orders, gates, skipped) = PipelinedMockGen::new();
        let stream = BufferedStream::<TestPB, TestErr>::new_pipelined(
            gen,
            (),
            BufferConfig::lazy(BufferConfig::UNBOUNDED, 0),
            depth,
        );
        (stream, orders, gates, skipped)
    }

    #[tokio::test(start_paused = true)]
    async fn pipelined_orders_overlap_and_resolve_fifo() {
        let (stream, orders, gates, _) = make_pipelined_stream(4);
        let f1 = stream.next_n(2).unwrap();
        let mut f2 = stream.next_n(3).unwrap();
        // Both orders must be dispatched while neither has resolved — the stop-and-wait
        // dispatcher would only ever have one out.
        assert!(
            wait_until(|| orders.lock().unwrap().len() == 2).await,
            "expected 2 overlapping orders, got {:?}",
            orders.lock().unwrap()
        );
        assert_eq!(*orders.lock().unwrap(), vec![2, 3]);
        // Release the SECOND order first: results must still reach consumers in FIFO order,
        // so f2 cannot resolve while order 1 is still pending.
        let gate1 = gates.lock().unwrap().pop_front().unwrap();
        release_next(&gates);
        tokio::time::timeout(Duration::from_millis(100), &mut f2)
            .await
            .expect_err("f2 must not resolve before order 1");
        let _ = gate1.send(Ok(()));
        let items1 = f1.await.expect("first batch resolves");
        let items2 = f2.await.expect("second batch resolves");
        assert_eq!(items1.len(), 2);
        assert_eq!(items2.len(), 3);
        assert_eq!(stream.position(), 5);
    }

    #[tokio::test(start_paused = true)]
    async fn three_orders_overlap_and_resolve_fifo() {
        // Three orders outstanding at once, resolved back-to-front: the dispatcher must still hand
        // results to consumers in issue order, so no consumer resolves until every earlier order
        // has. Two-deep overlap cannot catch a reordering that only shows up past the second slot.
        let (stream, orders, gates, _) = make_pipelined_stream(4);
        let f1 = stream.next_n(2).unwrap();
        let mut f2 = stream.next_n(3).unwrap();
        let mut f3 = stream.next_n(1).unwrap();
        assert!(
            wait_until(|| orders.lock().unwrap().len() == 3).await,
            "expected 3 overlapping orders, got {:?}",
            orders.lock().unwrap()
        );
        assert_eq!(*orders.lock().unwrap(), vec![2, 3, 1]);
        // Resolve strictly back-to-front: orders 3 and 2 complete while order 1 is still pending.
        let mut pending: Vec<_> = gates.lock().unwrap().drain(..).collect();
        let gate1 = pending.remove(0);
        for gate in pending.into_iter().rev() {
            let _ = gate.send(Ok(()));
        }
        for (label, f) in [("f2", &mut f2), ("f3", &mut f3)] {
            assert!(
                tokio::time::timeout(Duration::from_millis(50), f)
                    .await
                    .is_err(),
                "{label} must not resolve before order 1"
            );
        }
        assert_eq!(
            stream.buffered(),
            0,
            "nothing may be delivered out of order"
        );
        let _ = gate1.send(Ok(()));
        assert_eq!(f1.await.expect("first batch resolves").len(), 2);
        assert_eq!(f2.await.expect("second batch resolves").len(), 3);
        assert_eq!(f3.await.expect("third batch resolves").len(), 1);
        assert_eq!(stream.position(), 6);
    }

    #[tokio::test(start_paused = true)]
    async fn generator_error_tears_down_with_later_orders_in_flight() {
        // An order fails while later orders are still outstanding. `FuturesOrdered` means the
        // error cannot overtake an earlier order — order 1 still resolves Ok — but once it
        // surfaces, the generator drops the remaining futures and the dispatcher must fail every
        // consumer it is holding rather than leave any of them hanging on an order that will
        // never resolve.
        let (stream, orders, gates, _) = make_pipelined_stream(4);
        let f1 = stream.next_n(1).unwrap();
        let f2 = stream.next_n(1).unwrap();
        let f3 = stream.next_n(1).unwrap();
        // A prefetch on top, so the teardown has a completion of each kind to drain.
        let prefetch = stream.prefetch_n(2);
        assert!(
            wait_until(|| orders.lock().unwrap().len() == 4).await,
            "expected 4 orders (3 batches + prefetch), got {:?}",
            orders.lock().unwrap()
        );
        // Order 1 succeeds, order 2 fails; orders 3 and 4 are never released.
        release_next(&gates);
        fail_next(&gates);
        assert_eq!(
            f1.await
                .expect("the order before the failure still resolves")
                .len(),
            1
        );
        // Every consumer must receive the *generator's* error, not the incidental one a dropped
        // completion sender would produce — i.e. the dispatcher drained them rather than merely
        // vanishing.
        for (label, f) in [("f2", f2), ("f3", f3)] {
            let result = f.await;
            assert!(
                matches!(&result, Err(CorrelatedStreamError::SendError(m)) if m == MOCK_ORDER_FAILURE),
                "{label} must receive the generator error, got {result:?}"
            );
        }
        let prefetched = prefetch.await;
        assert!(
            matches!(&prefetched, Err(CorrelatedStreamError::SendError(m)) if m == MOCK_ORDER_FAILURE),
            "the outstanding prefetch must receive the generator error, got {prefetched:?}"
        );
        // Only the delivered batch counts; the failed and abandoned orders never advanced it.
        assert_eq!(stream.position(), 3, "all three batches were admitted");
    }

    #[tokio::test(start_paused = true)]
    async fn pipeline_depth_bounds_outstanding_orders() {
        let (stream, orders, gates, _) = make_pipelined_stream(2);
        let f1 = stream.next_n(1).unwrap();
        let f2 = stream.next_n(1).unwrap();
        let f3 = stream.next_n(1).unwrap();
        assert!(wait_until(|| orders.lock().unwrap().len() == 2).await);
        // The third order must wait for a pipeline slot.
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert_eq!(orders.lock().unwrap().len(), 2, "depth 2 must cap orders");
        release_next(&gates);
        assert!(
            wait_until(|| orders.lock().unwrap().len() == 3).await,
            "freed slot must admit the third order"
        );
        release_next(&gates);
        release_next(&gates);
        for f in [f1, f2, f3] {
            assert_eq!(f.await.expect("batch resolves").len(), 1);
        }
        assert_eq!(stream.position(), 3);
    }

    #[tokio::test(start_paused = true)]
    async fn pipelined_demand_is_not_double_requested() {
        // A single unmet demand must produce exactly one order even though the pipeline window
        // has room for more — in-flight quantities are netted out of the demand computation.
        let (stream, orders, gates, _) = make_pipelined_stream(8);
        let f1 = stream.next_n(5).unwrap();
        assert!(wait_until(|| !orders.lock().unwrap().is_empty()).await);
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert_eq!(
            *orders.lock().unwrap(),
            vec![5],
            "demand ordered exactly once"
        );
        release_next(&gates);
        assert_eq!(f1.await.expect("batch resolves").len(), 5);
    }

    #[tokio::test(start_paused = true)]
    async fn pipelined_resync_burns_items_landing_during_skip_wait() {
        // A prefetch order is in flight when the resync arrives. Its items land AFTER the resync
        // drained the buffer, but they lie below the target: they must be burned against the
        // deficit (shrinking the generator skip), not buffered for later delivery.
        let (stream, orders, gates, skipped) = make_pipelined_stream(4);
        let prefetch = stream.prefetch_n(4);
        assert!(wait_until(|| !orders.lock().unwrap().is_empty()).await);
        let resync = stream.resync(6);
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            0,
            "skip must wait for the pipeline"
        );
        release_next(&gates);
        resync.await.expect("resync completes");
        prefetch.await.expect("prefetch resolves");
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            2,
            "the 4 landed items cover part of the 6-deficit; only the rest is skipped"
        );
        assert_eq!(stream.position(), 6);
        assert_eq!(
            stream.buffered(),
            0,
            "pre-target items must not be buffered"
        );
    }

    #[tokio::test(start_paused = true)]
    async fn pipelined_resync_covered_by_inflight_order_keeps_surplus() {
        // The in-flight order is larger than the resync deficit: burning covers the whole
        // deficit (no generator skip at all) and the surplus items — past the target — stay
        // buffered for delivery.
        let (stream, orders, gates, skipped) = make_pipelined_stream(4);
        let prefetch = stream.prefetch_n(5);
        assert!(wait_until(|| !orders.lock().unwrap().is_empty()).await);
        let resync = stream.resync(3);
        release_next(&gates);
        resync.await.expect("resync completes");
        prefetch.await.expect("prefetch resolves");
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            0,
            "deficit fully covered by burns"
        );
        assert_eq!(stream.position(), 3);
        assert_eq!(stream.buffered(), 2, "post-target surplus stays buffered");
        // The surplus is served without a new order.
        let items = stream.next_n(2).unwrap().await.expect("served from buffer");
        assert_eq!(items.len(), 2);
        assert_eq!(
            orders.lock().unwrap().len(),
            1,
            "no extra order for buffered items"
        );
        assert_eq!(stream.position(), 5);
    }

    #[tokio::test(start_paused = true)]
    async fn next_n_rejected_while_resync_pending() {
        // A next_n admitted mid-resync could only be filled from an in-flight pre-target order
        // (no new orders are issued during a resync) while claiming a post-target position — a
        // silent desync. Reaching the dispatcher mid-resync also means the caller raced its own
        // resync, so which side of the barrier the request belongs on is not knowable here (and
        // the other parties may have put it on the other side). It is rejected, loudly, rather
        // than served or quietly parked.
        let (stream, orders, gates, skipped) = make_pipelined_stream(4);
        let prefetch = stream.prefetch_n(2);
        assert!(wait_until(|| !orders.lock().unwrap().is_empty()).await);
        let resync = stream.resync(5);
        // Let the dispatcher register the resync before the request.
        tokio::time::sleep(Duration::from_millis(20)).await;
        let rejected = stream.next_n(3).unwrap().await;
        assert!(
            matches!(rejected, Err(CorrelatedStreamError::ResyncInProgress)),
            "next_n during a pending resync must be rejected, got {rejected:?}"
        );
        assert_eq!(
            stream.position(),
            0,
            "rejected request must not advance position"
        );
        assert_eq!(
            orders.lock().unwrap().len(),
            1,
            "a rejected request must not be ordered for"
        );
        release_next(&gates);
        resync.await.expect("resync completes");
        prefetch.await.expect("prefetch resolves");
        assert_eq!(stream.position(), 5);
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            3,
            "2 landed items burned, 3 skipped"
        );
        // Once the resync completed, requests are served again — from the target onwards.
        let f = stream.next_n(1).unwrap();
        assert!(wait_until(|| orders.lock().unwrap().len() == 2).await);
        release_next(&gates);
        assert_eq!(f.await.expect("served after resync").len(), 1);
        assert_eq!(stream.position(), 6);
    }

    #[tokio::test(start_paused = true)]
    async fn pipelined_resync_waits_for_inflight_orders() {
        let (stream, orders, gates, skipped) = make_pipelined_stream(4);
        let f1 = stream.next_n(2).unwrap();
        assert!(wait_until(|| !orders.lock().unwrap().is_empty()).await);
        // Resync past the admitted request: the deficit (3) must be skipped at the generator,
        // but only after the in-flight order resolves — skips must not overtake generates.
        let resync = stream.resync(5);
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert_eq!(
            skipped.load(Ordering::SeqCst),
            0,
            "skip must wait for the pipeline"
        );
        release_next(&gates);
        resync.await.expect("resync completes");
        assert_eq!(skipped.load(Ordering::SeqCst), 3);
        assert_eq!(f1.await.expect("batch resolves").len(), 2);
        assert_eq!(stream.position(), 5);
    }
}