//! the order a shard takes its work in: live messages first, backfill steps in between
use std::collections::VecDeque;
use std::time::{Duration, Instant};
use tokio::runtime::Handle as TokioHandle;
use tokio::sync::mpsc::Receiver;
use tokio::sync::mpsc::error::TryRecvError;
use tokio::sync::{Semaphore, SemaphorePermit};
use crate::util::stop::Stop;
/// how long a backfill step waits behind live messages before it goes first, so a shard that
/// never runs out of live work still gets its backfills written
pub(crate) const BULK_FLOOR: Duration = Duration::from_millis(200);
pub(crate) enum Next<'p, L, B> {
Live(L),
/// a backfill step, with the permit that lets this shard write it
Bulk(B, SemaphorePermit<'p>),
}
enum Woke<'p, L> {
Live(L),
Permit(SemaphorePermit<'p>),
Other,
}
/// `C` is what the bulk channel carries, and `B` a queued step, which is either one of those or
/// a step the shard queued for itself with [`Self::push_local`]
pub(crate) struct Inbox<L, C, B = C> {
live: Receiver<L>,
bulk: Receiver<C>,
/// steps waiting their turn, with when they were queued
queued: VecDeque<(Instant, B)>,
live_open: bool,
bulk_open: bool,
floor: Duration,
/// turns of [`Self::next`]'s loop, so a test can tell waiting from polling
#[cfg(test)]
turns: usize,
}
impl<L, C, B: From<C>> Inbox<L, C, B> {
pub(crate) fn new(live: Receiver<L>, bulk: Receiver<C>, floor: Duration) -> Self {
Self {
live,
bulk,
queued: VecDeque::new(),
live_open: true,
bulk_open: true,
floor,
#[cfg(test)]
turns: 0,
}
}
/// the next piece of work, or none once hydrant stops or both channels close with nothing
/// left. a live message goes first unless the oldest backfill step has waited out the floor,
/// and a step only runs with one of `permits`, which never blocks: a shard that can't get one
/// keeps serving live messages
pub(crate) fn next<'p>(
&mut self,
stop: &Stop,
permits: &'p Semaphore,
handle: &TokioHandle,
) -> Option<Next<'p, L, B>> {
loop {
#[cfg(test)]
{
self.turns += 1;
}
if stop.is_set() {
return None;
}
self.take_bulk();
let bulk_due = self
.queued
.front()
.is_some_and(|(arrived, _)| arrived.elapsed() >= self.floor);
if !bulk_due && let Some(msg) = self.try_live() {
return Some(Next::Live(msg));
}
if !self.queued.is_empty() {
if let Ok(permit) = permits.try_acquire()
&& let Some((_, step)) = self.queued.pop_front()
{
return Some(Next::Bulk(step, permit));
}
match self.wait_for_permit(permits, handle) {
Woke::Live(msg) => return Some(Next::Live(msg)),
Woke::Permit(permit) => {
if let Some((_, step)) = self.queued.pop_front() {
return Some(Next::Bulk(step, permit));
}
}
Woke::Other => {}
}
continue;
}
if !self.live_open && !self.bulk_open {
return None;
}
if let Some(msg) = self.wait(handle) {
return Some(Next::Live(msg));
}
}
}
/// queues a step of the shard's own, which takes its turn like a step off the channel
pub(crate) fn push_local(&mut self, step: B) {
self.queued.push_back((Instant::now(), step));
}
fn take_bulk(&mut self) {
while self.bulk_open {
match self.bulk.try_recv() {
Ok(step) => self.queued.push_back((Instant::now(), step.into())),
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => self.bulk_open = false,
}
}
}
fn try_live(&mut self) -> Option<L> {
if !self.live_open {
return None;
}
match self.live.try_recv() {
Ok(msg) => Some(msg),
Err(TryRecvError::Empty) => None,
Err(TryRecvError::Disconnected) => {
self.live_open = false;
None
}
}
}
/// blocks until a live message comes, the permit frees up or another step arrives, for a
/// shard whose step waits on a permit another shard holds. the permit goes to waiters in
/// the order they asked, so shards take turns
fn wait_for_permit<'p>(&mut self, permits: &'p Semaphore, handle: &TokioHandle) -> Woke<'p, L> {
let (live_open, bulk_open) = (self.live_open, self.bulk_open);
let (live, bulk) = (&mut self.live, &mut self.bulk);
let woke = handle.block_on(async {
tokio::select! {
msg = live.recv(), if live_open => Ok(Ok(msg)),
permit = permits.acquire() => Ok(Err(permit)),
step = bulk.recv(), if bulk_open => Err(step),
}
});
match woke {
Ok(Ok(Some(msg))) => return Woke::Live(msg),
Ok(Ok(None)) => self.live_open = false,
// the permits are never closed
Ok(Err(permit)) => return permit.map_or(Woke::Other, Woke::Permit),
Err(Some(step)) => self.queued.push_back((Instant::now(), step.into())),
Err(None) => self.bulk_open = false,
}
Woke::Other
}
/// blocks until either channel has something. a backfill step goes onto the queue, a live
/// message comes back
fn wait(&mut self, handle: &TokioHandle) -> Option<L> {
let (live_open, bulk_open) = (self.live_open, self.bulk_open);
let (live, bulk) = (&mut self.live, &mut self.bulk);
let woke = handle.block_on(async {
tokio::select! {
msg = live.recv(), if live_open => Ok(msg),
step = bulk.recv(), if bulk_open => Err(step),
}
});
match woke {
Ok(Some(msg)) => return Some(msg),
Ok(None) => self.live_open = false,
Err(Some(step)) => self.queued.push_back((Instant::now(), step.into())),
Err(None) => self.bulk_open = false,
}
None
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::mpsc;
use super::*;
fn inbox(floor: Duration) -> (mpsc::Sender<u32>, mpsc::Sender<u32>, Inbox<u32, u32>) {
let (live_tx, live) = mpsc::channel(64);
let (bulk_tx, bulk) = mpsc::channel(64);
(live_tx, bulk_tx, Inbox::new(live, bulk, floor))
}
fn take(next: Option<Next<'_, u32, u32>>) -> (&'static str, u32) {
match next {
Some(Next::Live(msg)) => ("live", msg),
Some(Next::Bulk(step, _)) => ("bulk", step),
None => ("none", 0),
}
}
#[tokio::test(flavor = "multi_thread")]
async fn live_goes_first_until_a_step_has_waited_out_the_floor() {
let handle = tokio::runtime::Handle::current();
let (stop, permits) = (Stop::default(), Semaphore::new(1));
let (live, bulk, mut inbox) = inbox(Duration::from_millis(100));
for msg in 0..3 {
live.try_send(msg).unwrap();
}
bulk.try_send(10).unwrap();
let order = tokio::task::spawn_blocking(move || {
let mut order = Vec::new();
// the step just arrived, so live goes first
order.push(take(inbox.next(&stop, &permits, &handle)));
std::thread::sleep(Duration::from_millis(150));
// the step has waited past the floor, so it goes before the live messages left
order.push(take(inbox.next(&stop, &permits, &handle)));
order.push(take(inbox.next(&stop, &permits, &handle)));
// nothing waits, so the next step runs as soon as it comes
bulk.try_send(11).unwrap();
order.push(take(inbox.next(&stop, &permits, &handle)));
order.push(take(inbox.next(&stop, &permits, &handle)));
order
})
.await
.unwrap();
assert_eq!(
order,
[
("live", 0),
("bulk", 10),
("live", 1),
("live", 2),
("bulk", 11)
]
);
}
#[tokio::test(flavor = "multi_thread")]
async fn shards_without_the_permit_keep_serving_live_and_never_write_steps_at_once() {
const STEPS: u32 = 20;
let permits = Arc::new(Semaphore::new(1));
let writing = Arc::new(AtomicUsize::new(0));
let mut shards = Vec::new();
for _ in 0..2 {
let (live, bulk, mut inbox) = inbox(BULK_FLOOR);
for step in 0..STEPS {
bulk.try_send(step).unwrap();
live.try_send(step).unwrap();
}
drop((live, bulk));
let (permits, writing) = (permits.clone(), writing.clone());
let handle = tokio::runtime::Handle::current();
shards.push(tokio::task::spawn_blocking(move || {
let stop = Stop::default();
let (mut steps, mut lives) = (0, 0);
while let Some(next) = inbox.next(&stop, &permits, &handle) {
match next {
Next::Bulk(_, _permit) => {
assert_eq!(writing.fetch_add(1, Ordering::SeqCst), 0);
std::thread::sleep(Duration::from_millis(2));
writing.fetch_sub(1, Ordering::SeqCst);
steps += 1;
}
Next::Live(_) => lives += 1,
}
}
(steps, lives)
}));
}
for shard in shards {
assert_eq!(shard.await.unwrap(), (STEPS, STEPS));
}
}
#[tokio::test(flavor = "multi_thread")]
async fn a_stop_ends_the_inbox_before_the_next_step() {
let handle = tokio::runtime::Handle::current();
let (stop, permits) = (Stop::default(), Semaphore::new(1));
let (_live, bulk, mut inbox) = inbox(BULK_FLOOR);
bulk.try_send(1).unwrap();
stop.trigger();
let next = tokio::task::spawn_blocking(move || take(inbox.next(&stop, &permits, &handle)))
.await
.unwrap();
assert_eq!(next, ("none", 0));
}
#[tokio::test(flavor = "multi_thread")]
async fn a_step_waiting_on_the_permit_starts_once_it_frees_without_polling() {
let handle = tokio::runtime::Handle::current();
let permits = Arc::new(Semaphore::new(1));
let held = permits.clone().try_acquire_owned().unwrap();
let (live, bulk, mut inbox) = inbox(BULK_FLOOR);
bulk.try_send(1).unwrap();
let shard = {
let permits = permits.clone();
tokio::task::spawn_blocking(move || {
let stop = Stop::default();
let mut served = 0;
loop {
match inbox.next(&stop, &permits, &handle) {
Some(Next::Live(_)) => served += 1,
Some(Next::Bulk(_, _)) => return (Instant::now(), served, inbox.turns),
None => panic!("the inbox ended"),
}
}
})
};
// live keeps flowing while the step waits
for msg in 0..3 {
tokio::time::sleep(Duration::from_millis(30)).await;
live.send(msg).await.unwrap();
}
// and the shard sleeps through a quiet stretch instead of checking back
tokio::time::sleep(Duration::from_millis(300)).await;
let freed = Instant::now();
drop(held);
let (started, served, turns) = shard.await.unwrap();
assert_eq!(served, 3);
assert!(
started.duration_since(freed) < Duration::from_millis(20),
"{:?}",
started.duration_since(freed)
);
// a couple of turns per message and the permit, not one every few ms
assert!(turns <= 8, "{turns} turns");
}
}