#![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);
#[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,
}
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),
}
}
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);
backoff = (backoff * 2).min(BACKPRESSURE_MAX_INTERVAL);
}
false
}
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"
);
}
}