primitives/correlated_randomness/stream/buffered/

config.rs

// ================================
// ===== Buffer Configuration =====
// ================================

use std::{sync::Arc, time::Duration};

use blanket::blanket;
use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};

use crate::correlated_randomness::stream::CorrelatedStreamError;

/// Configuration for the `PreprocessingStream`.
#[derive(Debug, Clone)]
pub struct BufferConfig {
    /// Maximum number of elements the buffer can hold.
    capacity: usize,
    /// Maximum number of elements that can be requested at once.
    max_request_size: usize,
    /// Threshold to trigger generation of elements.
    refill_threshold: usize,
}

impl BufferConfig {
    /// Capacity sentinel for an effectively unbounded buffer: admission never trips and the
    /// dispatcher does not pre-allocate it. Used by on-demand streams (see `new_on_demand`).
    pub const UNBOUNDED: usize = usize::MAX;

    /// Creates a new `BufferConfig` with the specified capacity, and sets:
    /// - `max_request_size = capacity` (allowing requests up to the full buffer size).
    /// - `refill_threshold = capacity` (always refill when any elements are consumed).
    pub fn eager(capacity: usize) -> Self {
        assert!(capacity > 0, "capacity must be greater than 0");
        Self {
            capacity,
            max_request_size: capacity,
            refill_threshold: capacity,
        }
    }

    /// Creates a new `BufferConfig` with the specified capacity and max_request_size, and sets:
    /// - `refill_threshold = capacity` (always refill when any elements are consumed).
    pub fn eager_with(capacity: usize, max_request_size: usize) -> Self {
        assert!(capacity > 0, "capacity must be greater than 0");
        Self {
            capacity,
            max_request_size: max_request_size.min(capacity),
            refill_threshold: capacity,
        }
    }

    /// Creates a new `BufferConfig` with the specified capacity and refill threshold, and sets:
    /// - `max_request_size = capacity` (allowing requests up to the full buffer size).
    ///
    /// If `refill_threshold` > `capacity`, it is adjusted down to match `capacity`.
    pub fn lazy(capacity: usize, refill_threshold: usize) -> Self {
        assert!(capacity > 0, "capacity must be greater than 0");
        Self {
            capacity,
            max_request_size: capacity,
            refill_threshold: refill_threshold.min(capacity),
        }
    }

    /// Creates a new `BufferConfig` with the specified capacity, max_request_size, and
    /// refill_threshold.
    ///
    /// If `max_request_size` > `capacity` or `refill_threshold` > `capacity`, they are adjusted
    /// down to match `capacity`.
    pub fn lazy_with(capacity: usize, refill_threshold: usize, max_request_size: usize) -> Self {
        assert!(capacity > 0, "capacity must be greater than 0");
        Self {
            capacity,
            max_request_size: max_request_size.min(capacity),
            refill_threshold: refill_threshold.min(capacity),
        }
    }

    /// Wraps this config in a shared handle.
    pub fn shared(self) -> SharedBufferConfig {
        Arc::new(RwLock::new(self))
    }

    /// The maximum number of elements the buffer can hold (validated to be >= max_request_size).
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// The maximum number of elements that can be requested at once (validated to be <= capacity).
    #[inline]
    pub fn max_request_size(&self) -> usize {
        self.max_request_size
    }

    /// The threshold to trigger generation of elements (validated to be <= capacity).
    #[inline]
    pub fn refill_threshold(&self) -> usize {
        self.refill_threshold
    }

    // `capacity` is the ceiling: `max_request_size` and `refill_threshold` are always clamped to
    // it, matching the constructors. Setting the capacity down pulls the other two down with it;
    // setting either of the others clamps to the current capacity (never grows it).

    /// Sets the capacity, clamping `max_request_size` and `refill_threshold` down to fit.
    pub fn set_capacity(&mut self, capacity: usize) {
        assert!(capacity > 0, "capacity must be greater than 0");
        self.capacity = capacity;
        self.max_request_size = self.max_request_size.min(capacity);
        self.refill_threshold = self.refill_threshold.min(capacity);
    }

    /// Sets the max_request_size, clamped to the current capacity.
    pub fn set_max_request_size(&mut self, max_request_size: usize) {
        self.max_request_size = max_request_size.min(self.capacity);
    }

    /// Sets the refill_threshold, clamped to the current capacity.
    pub fn set_refill_threshold(&mut self, refill_threshold: usize) {
        self.refill_threshold = refill_threshold.min(self.capacity);
    }
}

// ================================
// ===== Buffer trait =============
// ================================

/// Default timeout when acquiring the config lock. Critical sections are tiny CPU work
/// (a few `usize` reads/writes) and never held across `.await`, so any contention longer
/// than this indicates a logic bug or stuck thread.
pub const LOCK_TIMEOUT: Duration = Duration::from_millis(500);

/// Acquires a shared **read** lock (concurrent reads don't block) for at most `LOCK_TIMEOUT`,
/// returning `LockTimeout` on failure.
#[inline]
pub fn try_read_config(
    m: &RwLock<BufferConfig>,
) -> Result<RwLockReadGuard<'_, BufferConfig>, CorrelatedStreamError> {
    m.try_read_for(LOCK_TIMEOUT)
        .ok_or(CorrelatedStreamError::LockTimeout {
            timeout_ms: LOCK_TIMEOUT.as_millis() as u64,
        })
}

/// Acquires an exclusive **write** lock for at most `LOCK_TIMEOUT`, returning `LockTimeout` on
/// failure.
#[inline]
pub fn try_write_config(
    m: &RwLock<BufferConfig>,
) -> Result<RwLockWriteGuard<'_, BufferConfig>, CorrelatedStreamError> {
    m.try_write_for(LOCK_TIMEOUT)
        .ok_or(CorrelatedStreamError::LockTimeout {
            timeout_ms: LOCK_TIMEOUT.as_millis() as u64,
        })
}

/// A shared, thread-safe handle to a [`BufferConfig`].
pub type SharedBufferConfig = Arc<RwLock<BufferConfig>>;

/// A buffer-backed entity exposing a shared `BufferConfig`.
///
/// Default-impl methods provide thread-safe getters and setters through an
/// `Arc<parking_lot::RwLock<_>>`, with a bounded timeout (`LOCK_TIMEOUT`) on acquisition. Getters
/// take a shared read lock (concurrent reads don't block); setters take an exclusive write lock.
/// Critical sections are tiny and never held across `.await`.
#[blanket(derive(Arc, Ref, Mut))]
pub trait Buffer {
    /// The shared configuration handle.
    fn config(&self) -> &SharedBufferConfig;

    #[inline]
    fn capacity(&self) -> Result<usize, CorrelatedStreamError> {
        Ok(try_read_config(self.config())?.capacity())
    }
    #[inline]
    fn max_request_size(&self) -> Result<usize, CorrelatedStreamError> {
        Ok(try_read_config(self.config())?.max_request_size())
    }
    #[inline]
    fn refill_threshold(&self) -> Result<usize, CorrelatedStreamError> {
        Ok(try_read_config(self.config())?.refill_threshold())
    }

    #[inline]
    fn set_capacity(&self, capacity: usize) -> Result<(), CorrelatedStreamError> {
        try_write_config(self.config())?.set_capacity(capacity);
        Ok(())
    }
    #[inline]
    fn set_max_request_size(&self, n: usize) -> Result<(), CorrelatedStreamError> {
        try_write_config(self.config())?.set_max_request_size(n);
        Ok(())
    }
    #[inline]
    fn set_refill_threshold(&self, n: usize) -> Result<(), CorrelatedStreamError> {
        try_write_config(self.config())?.set_refill_threshold(n);
        Ok(())
    }
}

/// Configures exactly one stream.
impl Buffer for SharedBufferConfig {
    fn config(&self) -> &SharedBufferConfig {
        self
    }
}

/// A slice of config handles is itself a [`Buffer`], so a whole bundle of streams is tuned in one
/// call. Writes (`set_*`) apply to every entry; reads (`capacity`, ...) return the first as a
/// representative, assuming the bundle is configured uniformly.
impl Buffer for [SharedBufferConfig] {
    fn config(&self) -> &SharedBufferConfig {
        self.first().expect("Buffer config slice must be non-empty")
    }
    fn set_capacity(&self, capacity: usize) -> Result<(), CorrelatedStreamError> {
        self.iter().try_for_each(|c| {
            try_write_config(c)?.set_capacity(capacity);
            Ok(())
        })
    }
    fn set_max_request_size(&self, n: usize) -> Result<(), CorrelatedStreamError> {
        self.iter().try_for_each(|c| {
            try_write_config(c)?.set_max_request_size(n);
            Ok(())
        })
    }
    fn set_refill_threshold(&self, n: usize) -> Result<(), CorrelatedStreamError> {
        self.iter().try_for_each(|c| {
            try_write_config(c)?.set_refill_threshold(n);
            Ok(())
        })
    }
}