use super::{
Confirmation, Delivery, Depth, DepthOf, DepthReading, Feed, FeedError, FeedErrorKind,
FeedWaker, Receipt, Unit, confirmation, depth_reading,
};
use crate::config::{ConfigError, FeedSection};
use crate::hash::FastMap;
use serde::Deserialize;
use std::collections::VecDeque;
use std::num::NonZeroU32;
use std::sync::{Arc, Mutex, MutexGuard};
const ATTEMPT_CAP: NonZeroU32 = NonZeroU32::new(10).expect("non-zero");
const BACKEND: &str = "memory";
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct FeedCounts {
pub visible: u64,
pub held: u64,
pub dead_lettered: u64,
pub acknowledged: u64,
}
#[derive(Clone, Debug)]
pub struct MemoryFeed {
state: Arc<Mutex<State>>,
}
impl Default for MemoryFeed {
fn default() -> MemoryFeed {
MemoryFeed::with_attempt_cap(ATTEMPT_CAP)
}
}
#[derive(Debug)]
struct State {
visible: VecDeque<(Unit, u32)>,
held: FastMap<Receipt, (Unit, u32)>,
dead: Vec<Unit>,
acknowledged: u64,
next_receipt: u64,
waker: Option<FeedWaker>,
attempt_cap: u32,
}
impl State {
fn wake(&self) {
if let Some(waker) = &self.waker {
waker.wake();
}
}
}
impl MemoryFeed {
#[must_use]
pub fn new() -> MemoryFeed {
MemoryFeed::default()
}
#[must_use]
pub fn with_attempt_cap(cap: NonZeroU32) -> MemoryFeed {
MemoryFeed {
state: Arc::new(Mutex::new(State {
visible: VecDeque::new(),
held: FastMap::default(),
dead: Vec::new(),
acknowledged: 0,
next_receipt: 0,
waker: None,
attempt_cap: cap.get(),
})),
}
}
pub fn from_section(section: &FeedSection) -> Result<MemoryFeed, ConfigError> {
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Body {}
let backend = section.backend();
if backend.type_tag() != BACKEND {
return Err(ConfigError::Component {
context: "feed".into(),
message: format!(
"the in-process feed is built from `feed: {{ {BACKEND}: {{}} }}`, but the \
section names `{}`",
backend.type_tag()
),
});
}
let _: Option<Body> = backend.deserialize_into()?;
if section.location.is_some() {
return Err(ConfigError::Component {
context: "feed".into(),
message: "the in-process feed keeps no state outside the process, so it takes \
no `location`"
.into(),
});
}
if section.settling_window.is_some() {
return Err(ConfigError::Component {
context: "feed".into(),
message: "the in-process feed counts its units exactly, so it takes no \
`settling_window`"
.into(),
});
}
Ok(MemoryFeed::with_attempt_cap(
section.attempt_cap.unwrap_or(ATTEMPT_CAP),
))
}
#[must_use]
pub fn counts(&self) -> FeedCounts {
let state = self.lock();
FeedCounts {
visible: state.visible.len() as u64,
held: state.held.len() as u64,
dead_lettered: state.dead.len() as u64,
acknowledged: state.acknowledged,
}
}
#[must_use]
pub fn dead_letters(&self) -> Vec<Unit> {
self.lock().dead.clone()
}
fn lock(&self) -> MutexGuard<'_, State> {
self.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
fn gone(receipt: Receipt) -> FeedError {
FeedError::new(
FeedErrorKind::Gone,
format!("{receipt} is not held by the in-process feed"),
)
}
impl Feed for MemoryFeed {
fn backend(&self) -> &str {
BACKEND
}
fn set_waker(&self, waker: FeedWaker) {
let mut state = self.lock();
if !state.visible.is_empty() {
waker.wake();
}
state.waker = Some(waker);
}
fn receive(&self, max: usize) -> Result<Vec<Delivery>, FeedError> {
let mut state = self.lock();
let take = max.min(state.visible.len());
let mut out = Vec::with_capacity(take);
for _ in 0..take {
let Some((unit, deliveries)) = state.visible.pop_front() else {
break;
};
state.next_receipt += 1;
let receipt = Receipt(state.next_receipt);
let deliveries = deliveries + 1;
state.held.insert(receipt, (unit.clone(), deliveries));
out.push(Delivery::new(unit, receipt, deliveries));
}
if take > 0 && !state.visible.is_empty() {
state.wake();
}
Ok(out)
}
fn renew(&self, receipt: Receipt) -> Result<(), FeedError> {
if self.lock().held.contains_key(&receipt) {
Ok(())
} else {
Err(gone(receipt))
}
}
fn acknowledge(&self, receipt: Receipt) -> Result<(), FeedError> {
let mut state = self.lock();
if state.held.remove(&receipt).is_none() {
return Err(gone(receipt));
}
state.acknowledged += 1;
Ok(())
}
fn fail(&self, receipt: Receipt, reason: &str) -> Result<(), FeedError> {
let mut state = self.lock();
let Some((unit, deliveries)) = state.held.remove(&receipt) else {
return Err(gone(receipt));
};
if deliveries >= state.attempt_cap {
let id = unit.id.clone();
state.dead.push(unit);
drop(state);
tracing::warn!(
unit = %id,
deliveries,
reason,
"in-process feed dead-lettered a unit at its delivery cap"
);
} else {
state.visible.push_back((unit, deliveries));
state.wake();
}
Ok(())
}
fn publish(&self, unit: Unit) -> Confirmation {
let mut state = self.lock();
state.visible.push_back((unit, 0));
state.wake();
let (confirmer, confirmation) = confirmation();
confirmer.confirm();
confirmation
}
fn read_depth(&self, of: DepthOf) -> Option<DepthReading> {
let state = self.lock();
let depth = match of {
DepthOf::DeadLetters => {
let mut depth = Depth::new("the in-process dead-letter list", false);
depth.visible = state.dead.len() as u64;
depth
}
DepthOf::Feed => {
let mut depth = Depth::new("the in-process feed", false);
depth.visible = state.visible.len() as u64;
depth.in_flight = state.held.len() as u64;
depth
}
};
drop(state);
let (reporter, reading) = depth_reading();
reporter.report(Ok(depth));
Some(reading)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::feed::wake_channel;
use std::time::Duration;
fn unit(id: &str) -> Unit {
Unit::new(id, id.as_bytes().to_vec())
}
fn receive_one_of(feed: &MemoryFeed, id: &str) -> Delivery {
let delivery = feed.receive(1).unwrap().remove(0);
assert_eq!(delivery.unit.id, id);
delivery
}
fn receive_one(feed: &MemoryFeed) -> Delivery {
let mut got = feed.receive(usize::MAX).unwrap();
assert_eq!(got.len(), 1, "expected exactly one delivery, got {got:?}");
got.remove(0)
}
#[test]
fn a_unit_at_the_cap_moves_to_the_in_memory_dead_letter_list() {
let feed = MemoryFeed::new();
drop(feed.publish(unit("poison")));
for i in 1..=10 {
let delivery = receive_one(&feed);
assert_eq!(delivery.deliveries, i, "delivery count on attempt {i}");
assert_eq!(delivery.unit, unit("poison"));
feed.fail(delivery.receipt, "unreadable").unwrap();
if i < 10 {
assert_eq!(feed.counts().visible, 1, "re-queued after attempt {i}");
assert_eq!(feed.counts().dead_lettered, 0);
}
}
assert_eq!(
feed.counts(),
FeedCounts {
visible: 0,
held: 0,
dead_lettered: 1,
acknowledged: 0,
}
);
assert_eq!(feed.dead_letters(), vec![unit("poison")]);
assert!(feed.receive(usize::MAX).unwrap().is_empty());
}
#[test]
fn the_attempt_cap_sets_the_delivery_that_dead_letters() {
let section: FeedSection =
serde_yaml::from_str("memory: {}\nattempt_cap: 3\n").expect("section");
let feed = MemoryFeed::from_section(§ion).unwrap();
drop(feed.publish(unit("poison")));
for i in 1..=2 {
let delivery = receive_one(&feed);
assert_eq!(delivery.deliveries, i);
feed.fail(delivery.receipt, "unreadable").unwrap();
assert_eq!(feed.counts().visible, 1, "re-queued after attempt {i}");
}
let third = receive_one(&feed);
assert_eq!(third.deliveries, 3);
feed.fail(third.receipt, "unreadable").unwrap();
assert_eq!(feed.counts().dead_lettered, 1, "dead-lettered on attempt 3");
assert_eq!(feed.counts().visible, 0);
}
#[test]
fn from_section_refuses_a_location() {
let mut section: FeedSection = serde_yaml::from_str("memory: {}\n").expect("section");
section.location = Some("s3://u:hunter2@ops/jobs/a/?X-Amz-Signature=hunter3".to_owned());
let err = MemoryFeed::from_section(§ion).unwrap_err();
assert!(
matches!(&err, ConfigError::Component { context, message }
if context == "feed" && message.contains("`location`")),
"{err:?}"
);
for shown in [err.to_string(), format!("{err:?}")] {
assert!(
!shown.contains("hunter2") && !shown.contains("ops/jobs/a"),
"the location must not print: {shown}"
);
}
}
#[test]
fn from_section_refuses_a_settling_window() {
let mut section: FeedSection = serde_yaml::from_str("memory: {}\n").expect("section");
section.settling_window = Some(Duration::from_secs(60));
let err = MemoryFeed::from_section(§ion).unwrap_err();
assert!(
matches!(&err, ConfigError::Component { context, message }
if context == "feed" && message.contains("`settling_window`")),
"{err:?}"
);
}
#[test]
fn the_depth_a_memory_feed_reports_matches_its_counts() {
let feed = MemoryFeed::with_attempt_cap(NonZeroU32::new(1).expect("non-zero"));
let read = |of| {
feed.read_depth(of)
.expect("the in-process feed reports depth")
.try_outcome()
.expect("settled at once")
.expect("read")
};
let empty = read(DepthOf::Feed);
assert!(empty.is_empty() && !empty.approximate, "{empty}");
assert!(read(DepthOf::DeadLetters).is_empty());
for id in [
"dead-1",
"dead-2",
"dead-3",
"held",
"visible-1",
"visible-2",
] {
drop(feed.publish(unit(id)));
}
for id in ["dead-1", "dead-2", "dead-3"] {
let dead = receive_one_of(&feed, id);
feed.fail(dead.receipt, "unreadable").unwrap();
}
let _held = receive_one_of(&feed, "held");
let counts = feed.counts();
let depth = read(DepthOf::Feed);
assert_eq!(
(depth.visible, depth.in_flight, depth.delayed, depth.unsent),
(2, 1, 0, 0)
);
assert_eq!(
(depth.visible, depth.in_flight),
(counts.visible, counts.held)
);
let dead = read(DepthOf::DeadLetters);
assert_eq!(dead.visible, 3);
assert_eq!(dead.visible, counts.dead_lettered);
assert_eq!((dead.delayed, dead.in_flight, dead.unsent), (0, 0, 0));
assert!(!dead.approximate);
}
#[test]
fn from_section_refuses_another_backend_and_a_body_key() {
let section: FeedSection = serde_yaml::from_str("other: {}\n").expect("section");
let err = MemoryFeed::from_section(§ion).unwrap_err().to_string();
assert!(err.contains("names `other`"), "{err}");
let section: FeedSection = serde_yaml::from_str("memory: { depth: 4 }\n").expect("section");
let err = MemoryFeed::from_section(§ion).unwrap_err().to_string();
assert!(err.starts_with("feed.memory"), "{err}");
assert!(err.contains("depth"), "{err}");
let bare: FeedSection = serde_yaml::from_str("memory:\n").expect("section");
assert!(MemoryFeed::from_section(&bare).is_ok());
}
#[test]
fn an_acknowledged_unit_leaves_and_its_receipt_is_gone() {
let feed = MemoryFeed::new();
drop(feed.publish(unit("done")));
let delivery = receive_one(&feed);
assert_eq!(feed.counts().held, 1);
feed.renew(delivery.receipt).unwrap();
feed.acknowledge(delivery.receipt).unwrap();
assert_eq!(
feed.counts(),
FeedCounts {
visible: 0,
held: 0,
dead_lettered: 0,
acknowledged: 1,
}
);
let kind = |r: Result<(), FeedError>| r.unwrap_err().kind;
assert_eq!(
kind(feed.acknowledge(delivery.receipt)),
FeedErrorKind::Gone
);
assert_eq!(kind(feed.renew(delivery.receipt)), FeedErrorKind::Gone);
assert_eq!(kind(feed.fail(delivery.receipt, "x")), FeedErrorKind::Gone);
assert_eq!(feed.counts().acknowledged, 1);
assert!(feed.receive(usize::MAX).unwrap().is_empty());
}
#[test]
fn a_unit_published_twice_is_delivered_twice() {
let feed = MemoryFeed::new();
drop(feed.publish(unit("twice")));
drop(feed.publish(unit("twice")));
let got = feed.receive(usize::MAX).unwrap();
assert_eq!(got.len(), 2);
assert_ne!(got[0].receipt, got[1].receipt);
assert!(
got.iter()
.all(|d| d.unit == unit("twice") && d.deliveries == 1)
);
feed.acknowledge(got[0].receipt).unwrap();
assert_eq!(feed.counts().held, 1);
feed.acknowledge(got[1].receipt).unwrap();
assert_eq!(feed.counts().acknowledged, 2);
}
#[test]
fn a_publish_is_confirmed_at_once_and_wakes_the_reader() {
let feed = MemoryFeed::new();
let (waker, park) = wake_channel();
feed.set_waker(waker);
assert!(!park.park(Duration::ZERO));
let mut confirmation = feed.publish(unit("one"));
assert!(matches!(confirmation.try_outcome(), Some(Ok(()))));
assert!(park.park(Duration::ZERO), "publish did not wake the reader");
let late = MemoryFeed::new();
drop(late.publish(unit("early")));
let (waker, park) = wake_channel();
late.set_waker(waker);
assert!(park.park(Duration::ZERO), "set_waker missed a waiting unit");
}
#[test]
fn units_arrive_in_publish_order_and_a_failed_unit_goes_to_the_back() {
let feed = MemoryFeed::new();
let (waker, park) = wake_channel();
feed.set_waker(waker);
for id in ["a", "b", "c"] {
drop(feed.publish(unit(id)));
}
let first = feed.receive(1).unwrap();
assert_eq!(first.len(), 1);
assert_eq!(first[0].unit.id, "a");
assert!(park.park(Duration::ZERO));
feed.fail(first[0].receipt, "retry later").unwrap();
assert!(park.park(Duration::ZERO), "fail did not wake the reader");
let ids: Vec<_> = feed
.receive(usize::MAX)
.unwrap()
.into_iter()
.map(|d| (d.unit.id, d.deliveries))
.collect();
assert_eq!(
ids,
vec![
("b".to_owned(), 1),
("c".to_owned(), 1),
("a".to_owned(), 2)
]
);
}
#[test]
fn a_partial_receive_leaves_a_wake_pending() {
let feed = MemoryFeed::new();
let (waker, park) = wake_channel();
feed.set_waker(waker);
for id in ["a", "b"] {
drop(feed.publish(unit(id)));
}
assert!(park.park(Duration::ZERO));
assert_eq!(feed.receive(1).unwrap().len(), 1);
assert_eq!(feed.counts().visible, 1);
assert!(
park.park(Duration::ZERO),
"a unit is visible but the reader's park was not signalled"
);
}
#[test]
fn an_empty_receive_at_capacity_does_not_wake() {
let feed = MemoryFeed::new();
let (waker, park) = wake_channel();
feed.set_waker(waker);
drop(feed.publish(unit("a")));
assert!(park.park(Duration::ZERO));
assert!(feed.receive(0).unwrap().is_empty());
assert!(
!park.park(Duration::ZERO),
"receive(0) at capacity re-signalled"
);
}
#[test]
fn a_panicking_subscriber_in_fail_at_the_cap_keeps_the_unit() {
use tracing_subscriber::layer::SubscriberExt;
struct PanicOnWarn;
impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for PanicOnWarn {
fn on_event(
&self,
event: &tracing::Event<'_>,
_ctx: tracing_subscriber::layer::Context<'_, S>,
) {
if *event.metadata().level() == tracing::Level::WARN {
panic!("warning emitted");
}
}
}
let feed = MemoryFeed::new();
drop(feed.publish(unit("poison")));
for _ in 1..10 {
let delivery = receive_one(&feed);
feed.fail(delivery.receipt, "x").unwrap();
}
let last = receive_one(&feed);
let subscriber = tracing_subscriber::registry().with(PanicOnWarn);
let unwound = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
tracing::subscriber::with_default(subscriber, || feed.fail(last.receipt, "x"))
}));
assert!(unwound.is_err(), "the subscriber did not panic");
assert_eq!(
feed.counts(),
FeedCounts {
visible: 0,
held: 0,
dead_lettered: 1,
acknowledged: 0,
}
);
assert!(
!feed.state.is_poisoned(),
"the warning unwound through the lock"
);
}
#[test]
fn a_later_set_waker_replaces_the_waker() {
let feed = MemoryFeed::new();
let (first, first_park) = wake_channel();
feed.set_waker(first);
let (second, second_park) = wake_channel();
feed.set_waker(second);
drop(feed.publish(unit("a")));
assert!(
second_park.park(Duration::ZERO),
"publish signalled the replaced waker"
);
assert!(!first_park.park(Duration::ZERO));
}
}