use std::time::{Duration, Instant};
use tokio::sync::mpsc::error::TrySendError;
use crate::worker::registry::{WorkerMessage, WorkerTaskSender};
pub const PARK_POLL_INTERVAL: Duration = Duration::from_millis(500);
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DeliveryRefusal {
Saturated {
waited: Duration,
},
Full,
Closed,
NotAccepting {
reason: String,
},
}
pub fn deliver_within_schedule_to_start(
sender: &WorkerTaskSender,
message: WorkerMessage,
schedule_to_start: Option<Duration>,
accepting: &mut dyn FnMut() -> Result<(), String>,
) -> Result<(), DeliveryRefusal> {
let started_at = Instant::now();
let mut pending = message;
loop {
match sender.try_send(pending) {
Ok(()) => return Ok(()),
Err(TrySendError::Closed(_)) => return Err(DeliveryRefusal::Closed),
Err(TrySendError::Full(returned)) => {
let Some(budget) = schedule_to_start else {
return Err(DeliveryRefusal::Full);
};
let waited = started_at.elapsed();
let Some(remaining) = budget.checked_sub(waited) else {
return Err(DeliveryRefusal::Saturated { waited });
};
if remaining.is_zero() {
return Err(DeliveryRefusal::Saturated { waited });
}
if let Err(reason) = accepting() {
return Err(DeliveryRefusal::NotAccepting { reason });
}
std::thread::sleep(remaining.min(PARK_POLL_INTERVAL));
pending = returned;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn task() -> WorkerMessage {
WorkerMessage::ActivityTask(Box::default())
}
fn always_accepting() -> impl FnMut() -> Result<(), String> {
|| Ok(())
}
#[test]
fn a_free_intake_accepts_the_task_immediately() -> Result<(), String> {
let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
let mut accepting = always_accepting();
deliver_within_schedule_to_start(&sender, task(), None, &mut accepting)
.map_err(|refusal| format!("unexpected refusal: {refusal:?}"))?;
receiver
.try_recv()
.map_err(|error| format!("task never reached the worker: {error}"))?;
Ok(())
}
#[test]
fn a_full_intake_with_no_clock_fails_at_once() {
let (sender, _receiver) = tokio::sync::mpsc::channel(1);
let mut accepting = always_accepting();
assert!(
sender.try_send(task()).is_ok(),
"the fixture must fill the single intake slot"
);
let started = Instant::now();
let refusal = deliver_within_schedule_to_start(&sender, task(), None, &mut accepting);
assert_eq!(refusal, Err(DeliveryRefusal::Full));
assert!(
started.elapsed() < PARK_POLL_INTERVAL,
"the unclocked path must not wait"
);
}
#[test]
fn a_closed_intake_is_never_reported_as_saturation() {
let (sender, receiver) = tokio::sync::mpsc::channel(1);
drop(receiver);
let mut accepting = always_accepting();
let refusal = deliver_within_schedule_to_start(
&sender,
task(),
Some(Duration::from_secs(30)),
&mut accepting,
);
assert_eq!(refusal, Err(DeliveryRefusal::Closed));
}
#[test]
fn an_intake_that_stays_full_past_the_clock_is_saturated() {
let (sender, _receiver) = tokio::sync::mpsc::channel(1);
let mut accepting = always_accepting();
assert!(sender.try_send(task()).is_ok(), "fixture must fill intake");
let budget = Duration::from_millis(300);
let started = Instant::now();
let refusal =
deliver_within_schedule_to_start(&sender, task(), Some(budget), &mut accepting);
let elapsed = started.elapsed();
assert!(
matches!(refusal, Err(DeliveryRefusal::Saturated { waited }) if waited >= budget),
"expected saturation after at least {budget:?}, got {refusal:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"saturation must be bounded by the clock, took {elapsed:?}"
);
}
#[test]
fn an_intake_that_frees_up_within_the_clock_is_served() -> Result<(), String> {
let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
let mut accepting = always_accepting();
sender
.try_send(task())
.map_err(|error| format!("fixture fill failed: {error}"))?;
let drainer = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(100));
let taken = receiver.try_recv().is_ok();
(receiver, taken)
});
let outcome = deliver_within_schedule_to_start(
&sender,
task(),
Some(Duration::from_secs(5)),
&mut accepting,
);
let (_receiver, taken) = drainer.join().map_err(|_| "drainer thread panicked")?;
assert!(taken, "the fixture drainer never freed the intake slot");
assert_eq!(outcome, Ok(()));
Ok(())
}
#[test]
fn a_drain_during_the_retry_window_fails_with_the_drain_reason() {
let (sender, _receiver) = tokio::sync::mpsc::channel(1);
assert!(sender.try_send(task()).is_ok(), "fixture must fill intake");
let mut accepting = || Err("server draining".to_owned());
let refusal = deliver_within_schedule_to_start(
&sender,
task(),
Some(Duration::from_secs(30)),
&mut accepting,
);
assert_eq!(
refusal,
Err(DeliveryRefusal::NotAccepting {
reason: "server draining".to_owned()
})
);
}
}