tenshift-core 0.1.3

Thread-safe, backpressure-aware data loading pipeline for iterative processing
Documentation
//! Internal flow-control helpers for tenshift pipelines.
//!
//! This module centralizes channel sizing and explicit backpressure signaling so
//! the executor can swap those policies without rewriting stage orchestration.

#![allow(clippy::module_name_repetitions)]

use crossbeam_channel::{Sender, TrySendError};
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread;
use std::time::Duration;

const BACKPRESSURE_MIN_INTERVAL: Duration = Duration::from_micros(10);
const BACKPRESSURE_MAX_INTERVAL: Duration = Duration::from_millis(50);

/// Internal bounded-channel capacities derived from runtime configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ChannelCapacities {
    pub(crate) raw_chunks: usize,
    pub(crate) processed_chunks: usize,
    pub(crate) output_batches: usize,
}

/// Derive channel capacities from worker count and configured prefetch depth.
///
/// The output channel must be large enough for short consumer jitter, while the
/// raw and processed channels need enough headroom for each worker to keep one
/// chunk in flight without turning small `prefetch` settings into thread
/// starvation under moderate contention.
pub(crate) fn tuned_channel_capacities(
    prefetch_size: usize,
    num_workers: usize,
) -> ChannelCapacities {
    let workers = num_workers.max(1);
    let prefetch = prefetch_size.max(1);
    let worker_headroom = workers.saturating_mul(2);

    ChannelCapacities {
        raw_chunks: prefetch.max(worker_headroom),
        processed_chunks: prefetch.saturating_add(worker_headroom).max(2),
        output_batches: prefetch.max(workers),
    }
}

/// Block the caller while downstream output pressure asks the source to pause.
pub(crate) fn wait_if_paused(output_paused: &AtomicBool, shutdown: &AtomicBool) -> bool {
    let mut backoff = BACKPRESSURE_MIN_INTERVAL;
    while output_paused.load(Ordering::Relaxed) {
        if shutdown.load(Ordering::Relaxed) {
            output_paused.store(false, Ordering::Relaxed);
            return true;
        }
        thread::sleep(backoff);
        // Exponential backoff: double each time, cap at max.
        // Avoids thundering herd when many threads wake simultaneously.
        backoff = (backoff * 2).min(BACKPRESSURE_MAX_INTERVAL);
    }
    false
}

/// Send to the consumer-facing output channel while toggling the shared pause
/// flag whenever the channel is saturated.
pub(crate) fn send_with_backpressure<T>(
    out_tx: &Sender<T>,
    mut value: T,
    output_paused: &AtomicBool,
    shutdown: &AtomicBool,
) -> bool {
    let mut backoff = BACKPRESSURE_MIN_INTERVAL;
    loop {
        if shutdown.load(Ordering::Relaxed) {
            output_paused.store(false, Ordering::Relaxed);
            return true;
        }

        match out_tx.try_send(value) {
            Ok(()) => {
                output_paused.store(false, Ordering::Relaxed);
                return false;
            }
            Err(TrySendError::Full(returned)) => {
                output_paused.store(true, Ordering::Relaxed);
                value = returned;
                thread::sleep(backoff);
                backoff = (backoff * 2).min(BACKPRESSURE_MAX_INTERVAL);
            }
            Err(TrySendError::Disconnected(_returned)) => {
                output_paused.store(false, Ordering::Relaxed);
                return true;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::send_with_backpressure;
    use crossbeam_channel::bounded;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn output_full_activates_and_clears_backpressure_flag() {
        let (tx, rx) = bounded::<usize>(1);
        tx.send(1).expect("prime bounded channel");

        let paused = Arc::new(AtomicBool::new(false));
        let shutdown = Arc::new(AtomicBool::new(false));

        let send_paused = Arc::clone(&paused);
        let send_shutdown = Arc::clone(&shutdown);
        let sender = tx.clone();

        let send_handle = thread::spawn(move || {
            assert!(!send_with_backpressure(
                &sender,
                2,
                send_paused.as_ref(),
                send_shutdown.as_ref()
            ));
        });

        let mut observed = false;
        let deadline = std::time::Instant::now() + Duration::from_secs(1);
        while std::time::Instant::now() < deadline {
            if paused.load(Ordering::Relaxed) {
                observed = true;
                break;
            }
            thread::sleep(Duration::from_millis(1));
        }

        assert!(observed, "full output channel must activate backpressure");
        assert_eq!(rx.recv().expect("drain first item"), 1);
        assert_eq!(rx.recv().expect("drain second item"), 2);

        send_handle
            .join()
            .expect("sender thread should exit cleanly");
        assert!(
            !paused.load(Ordering::Relaxed),
            "backpressure flag must clear after output drain"
        );
    }
}