use melin_pipeline::spsc;
use melin_pipeline::wait::WaitStrategy;
use melin_transport_core::pipeline::RequestId;
use melin_wire_protocol::control::RefusalKind;
use crate::batch_boundary::AtBoundary;
pub const EARLY_REPLY_QUEUE_CAPACITY: usize = 4096;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct EarlyReply<R> {
pub connection_id: u64,
pub request_id: RequestId,
pub kind: EarlyKind<R>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EarlyKind<R> {
Rejected(R),
Filtered,
Refused {
kind: RefusalKind,
reason: &'static str,
},
Busy,
}
pub fn early_reply_channel<R: Copy>(capacity: usize) -> (EarlyReplySender<R>, EarlyReplyQueue<R>) {
let (tx, rx) = spsc::channel::<Option<EarlyReply<R>>>(capacity, WaitStrategy::SpinThenYield);
(EarlyReplySender { tx }, EarlyReplyQueue { rx })
}
pub struct EarlyReplySender<R> {
tx: spsc::Producer<Option<EarlyReply<R>>>,
}
impl<R: Copy> EarlyReplySender<R> {
pub(crate) fn try_send(&mut self, reply: EarlyReply<R>) -> Result<(), spsc::Full> {
self.tx
.try_push_with(|slot| *slot = Some(reply))
.map(|_| ())
}
pub(crate) fn flush(&mut self) {
self.tx.flush();
}
}
pub struct EarlyReplyQueue<R> {
rx: spsc::Consumer<Option<EarlyReply<R>>>,
}
impl<R: Copy> EarlyReplyQueue<R> {
#[inline]
pub(crate) fn register(&mut self, drain_control: impl FnOnce()) -> Registered<'_, R> {
let mut registered = Registered { queue: self };
registered.reregister(drain_control);
registered
}
}
pub(crate) struct Registered<'q, R> {
queue: &'q mut EarlyReplyQueue<R>,
}
impl<R: Copy> Registered<'_, R> {
#[inline]
pub(crate) fn reregister(&mut self, drain_control: impl FnOnce()) {
self.queue.rx.refresh();
drain_control();
}
#[inline]
pub(crate) fn drain(
&mut self,
_between_events: &AtBoundary<'_>,
mut emit: impl FnMut(EarlyReply<R>),
) -> usize {
let mut sent = 0;
while let Some((_, Some(reply))) = self.queue.rx.try_consume_visible() {
emit(reply);
sent += 1;
}
sent
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::batch_boundary::BatchBoundary;
fn refusal(connection_id: u64) -> EarlyReply<u64> {
EarlyReply {
connection_id,
request_id: RequestId(connection_id * 10),
kind: EarlyKind::Rejected(connection_id * 100),
}
}
fn drained(registered: &mut Registered<'_, u64>) -> Vec<u64> {
let boundary = BatchBoundary::new();
let at = boundary
.at_boundary()
.expect("a fresh tracker is between events");
let mut out = Vec::new();
let count = registered.drain(&at, |reply| out.push(reply.connection_id));
assert_eq!(count, out.len(), "the count is of what was handed over");
out
}
#[test]
fn a_refusal_waits_for_both_the_flush_and_a_snapshot() {
let (mut tx, mut rx) = early_reply_channel::<u64>(EARLY_REPLY_QUEUE_CAPACITY);
tx.try_send(refusal(1)).unwrap();
let mut registered = rx.register(|| {});
assert!(drained(&mut registered).is_empty(), "not flushed");
tx.flush();
assert!(
drained(&mut registered).is_empty(),
"flushed, but after the snapshot"
);
registered.reregister(|| {});
assert_eq!(drained(&mut registered), [1]);
assert!(drained(&mut registered).is_empty(), "drained once");
registered.reregister(|| {});
assert!(drained(&mut registered).is_empty(), "the queue is empty");
}
#[test]
fn a_snapshot_is_taken_before_the_control_drain() {
let (mut tx, mut rx) = early_reply_channel::<u64>(EARLY_REPLY_QUEUE_CAPACITY);
tx.try_send(refusal(1)).unwrap();
tx.flush();
let mut drains = 0;
let mut registered = rx.register(|| {
drains += 1;
tx.try_send(refusal(2)).unwrap();
tx.flush();
});
assert_eq!(drains, 1, "registering drains the control channel once");
assert_eq!(drained(&mut registered), [1], "the snapshot predates 2");
registered.reregister(|| {
drains += 1;
tx.try_send(refusal(3)).unwrap();
tx.flush();
});
assert_eq!(drains, 2);
assert_eq!(drained(&mut registered), [2], "the snapshot predates 3");
let mut registered = rx.register(|| drains += 1);
assert_eq!(drained(&mut registered), [3]);
assert_eq!(drains, 3);
}
#[test]
fn a_full_queue_is_an_error_not_a_wait() {
let (mut tx, _rx) = early_reply_channel::<u64>(EARLY_REPLY_QUEUE_CAPACITY);
for _ in 0..EARLY_REPLY_QUEUE_CAPACITY {
tx.try_send(refusal(1)).unwrap();
}
assert!(tx.try_send(refusal(1)).is_err());
}
#[test]
fn a_queue_holds_the_capacity_it_was_built_with() {
for capacity in [2, 8] {
let (mut tx, _rx) = early_reply_channel::<u64>(capacity);
for _ in 0..capacity {
tx.try_send(refusal(1)).unwrap();
}
assert!(tx.try_send(refusal(1)).is_err(), "{capacity}");
}
}
#[test]
fn a_drained_queue_takes_more() {
let (mut tx, mut rx) = early_reply_channel::<u64>(EARLY_REPLY_QUEUE_CAPACITY);
for round in 0..3u64 {
for _ in 0..EARLY_REPLY_QUEUE_CAPACITY {
tx.try_send(refusal(round)).unwrap();
}
tx.flush();
let mut registered = rx.register(|| {});
assert_eq!(
drained(&mut registered).len(),
EARLY_REPLY_QUEUE_CAPACITY,
"round {round}"
);
}
}
#[test]
fn every_kind_travels_with_its_request_id() {
let (mut tx, mut rx) = early_reply_channel::<u64>(EARLY_REPLY_QUEUE_CAPACITY);
let kinds = [
EarlyKind::Rejected(7),
EarlyKind::Filtered,
EarlyKind::Refused {
kind: RefusalKind::Malformed,
reason: "bad",
},
EarlyKind::Busy,
];
for (n, kind) in (0u64..).zip(kinds) {
tx.try_send(EarlyReply {
connection_id: 1,
request_id: RequestId(100 + n),
kind,
})
.unwrap();
}
tx.flush();
let boundary = BatchBoundary::new();
let at = boundary.at_boundary().expect("between events");
let mut released = Vec::new();
rx.register(|| {})
.drain(&at, |r| released.push((r.request_id, r.kind)));
assert_eq!(
released,
(100..).map(RequestId).zip(kinds).collect::<Vec<_>>()
);
}
}