use crate::error::ErrorClass;
use crate::metrics::Meter;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::sync::oneshot;
mod end_check;
pub mod memory;
mod metered;
pub use end_check::{END_CHECK_INTERVAL, EndAction, EndCheck, EndOutcome, MarkerRead};
pub(crate) use metered::{DEAD_LETTER_READ_INTERVAL, DeadLetterSampler, MeteredFeed};
pub const DEFAULT_SETTLING_WINDOW: Duration = Duration::from_secs(5 * 60);
pub const STOP_RESERVE: Duration = Duration::from_secs(5);
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct Unit {
pub id: String,
pub body: Vec<u8>,
}
impl Unit {
pub fn new(id: impl Into<String>, body: impl Into<Vec<u8>>) -> Unit {
Unit {
id: id.into(),
body: body.into(),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Receipt(pub u64);
impl fmt::Display for Receipt {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "receipt {}", self.0)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct Delivery {
pub unit: Unit,
pub receipt: Receipt,
pub deliveries: u32,
}
impl Delivery {
#[must_use]
pub fn new(unit: Unit, receipt: Receipt, deliveries: u32) -> Delivery {
Delivery {
unit,
receipt,
deliveries,
}
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct FeedMetricsCtx {
pub meter: Option<Meter>,
pub per_partition_detail: bool,
}
impl FeedMetricsCtx {
#[must_use]
pub fn new(meter: Option<Meter>) -> Self {
FeedMetricsCtx {
meter,
per_partition_detail: false,
}
}
#[must_use]
pub fn with_partition_detail(mut self, enabled: bool) -> Self {
self.per_partition_detail = enabled;
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum FeedErrorKind {
Gone,
Refused,
Retryable,
Fatal,
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("feed error ({kind:?}): {reason}")]
#[non_exhaustive]
pub struct FeedError {
pub kind: FeedErrorKind,
pub reason: String,
}
impl FeedError {
pub fn new(kind: FeedErrorKind, reason: impl Into<String>) -> FeedError {
FeedError {
kind,
reason: reason.into(),
}
}
#[must_use]
pub fn class(&self) -> ErrorClass {
match self.kind {
FeedErrorKind::Retryable => ErrorClass::Retryable,
FeedErrorKind::Gone | FeedErrorKind::Refused | FeedErrorKind::Fatal => {
ErrorClass::Fatal
}
}
}
}
pub trait Feed: Send + Sync {
fn backend(&self) -> &str;
fn set_waker(&self, waker: FeedWaker);
fn attach_metrics(&self, _cx: FeedMetricsCtx) {}
fn receive(&self, max: usize) -> Result<Vec<Delivery>, FeedError>;
fn renew(&self, receipt: Receipt) -> Result<(), FeedError>;
fn acknowledge(&self, receipt: Receipt) -> Result<(), FeedError>;
fn fail(&self, receipt: Receipt, reason: &str) -> Result<(), FeedError>;
fn publish(&self, unit: Unit) -> Confirmation;
fn read_depth(&self, _of: DepthOf) -> Option<DepthReading> {
None
}
fn hands_back(&self) -> bool {
false
}
fn stop_receiving(&self) {}
fn hand_back(&self, _receipt: Receipt) -> Result<(), FeedError> {
Ok(())
}
fn close(&self, _within: Duration) -> Option<Confirmation> {
None
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum FeedOrigin {
Section,
Given,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum DepthOf {
Feed,
DeadLetters,
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct Depth {
pub what: String,
pub visible: u64,
pub delayed: u64,
pub in_flight: u64,
pub unsent: u64,
pub approximate: bool,
}
impl Depth {
#[must_use]
pub fn new(what: impl Into<String>, approximate: bool) -> Depth {
Depth {
what: what.into(),
visible: 0,
delayed: 0,
in_flight: 0,
unsent: 0,
approximate,
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.visible == 0 && self.delayed == 0 && self.in_flight == 0 && self.unsent == 0
}
}
impl fmt::Display for Depth {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"{}: {} visible, {} delayed, {} in flight, {} unsent",
self.what, self.visible, self.delayed, self.in_flight, self.unsent
)?;
if self.approximate {
f.write_str(" (approximate)")?;
}
Ok(())
}
}
#[derive(Debug)]
pub struct DepthReporter(oneshot::Sender<Result<Depth, FeedError>>);
impl DepthReporter {
pub fn report(self, outcome: Result<Depth, FeedError>) {
let _ = self.0.send(outcome);
}
}
#[derive(Debug)]
pub struct DepthReading {
rx: oneshot::Receiver<Result<Depth, FeedError>>,
outcome: Option<Result<Depth, FeedError>>,
}
impl DepthReading {
pub fn try_outcome(&mut self) -> Option<Result<Depth, FeedError>> {
if self.outcome.is_none() {
use oneshot::error::TryRecvError;
self.outcome = match self.rx.try_recv() {
Ok(outcome) => Some(outcome),
Err(TryRecvError::Empty) => None,
Err(TryRecvError::Closed) => Some(Err(FeedError::new(
FeedErrorKind::Retryable,
"the feed dropped a depth read before settling it",
))),
};
}
self.outcome.clone()
}
}
#[must_use]
pub fn depth_reading() -> (DepthReporter, DepthReading) {
let (tx, rx) = oneshot::channel();
(DepthReporter(tx), DepthReading { rx, outcome: None })
}
#[derive(Clone, Debug)]
pub struct FeedWaker(crossbeam_channel::Sender<()>);
impl FeedWaker {
pub fn wake(&self) {
let _ = self.0.try_send(());
}
}
#[derive(Debug)]
pub struct FeedPark {
rx: crossbeam_channel::Receiver<()>,
_tx: crossbeam_channel::Sender<()>,
}
impl FeedPark {
pub fn park(&self, timeout: Duration) -> bool {
self.rx.recv_timeout(timeout).is_ok()
}
}
#[must_use]
pub fn wake_channel() -> (FeedWaker, FeedPark) {
let (tx, rx) = crossbeam_channel::bounded(1);
(FeedWaker(tx.clone()), FeedPark { rx, _tx: tx })
}
#[derive(Debug)]
pub struct Confirmer(oneshot::Sender<Result<(), FeedError>>);
impl Confirmer {
pub fn confirm(self) {
let _ = self.0.send(Ok(()));
}
pub fn reject(self, error: FeedError) {
let _ = self.0.send(Err(error));
}
}
#[derive(Debug)]
pub struct Confirmation {
rx: oneshot::Receiver<Result<(), FeedError>>,
outcome: Option<Result<(), FeedError>>,
on_settle: Option<OnSettle>,
}
struct OnSettle(Box<SettleFn>);
type SettleFn = dyn FnOnce(&Result<(), FeedError>) + Send + Sync;
impl fmt::Debug for OnSettle {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("OnSettle")
}
}
impl Confirmation {
pub fn try_outcome(&mut self) -> Option<Result<(), FeedError>> {
if self.outcome.is_none() {
use oneshot::error::TryRecvError;
let outcome = match self.rx.try_recv() {
Ok(outcome) => outcome,
Err(TryRecvError::Empty) => return None,
Err(TryRecvError::Closed) => Err(unsettled()),
};
self.settle(outcome);
}
self.outcome.clone()
}
pub(crate) fn on_settle(
mut self,
f: impl FnOnce(&Result<(), FeedError>) + Send + Sync + 'static,
) -> Confirmation {
self.on_settle = Some(OnSettle(Box::new(f)));
self
}
fn settle(&mut self, outcome: Result<(), FeedError>) {
if let Some(OnSettle(f)) = self.on_settle.take() {
f(&outcome);
}
self.outcome = Some(outcome);
}
}
impl Future for Confirmation {
type Output = Result<(), FeedError>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if let Some(outcome) = &self.outcome {
return Poll::Ready(outcome.clone());
}
let outcome = match Pin::new(&mut self.rx).poll(cx) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Ok(outcome)) => outcome,
Poll::Ready(Err(_)) => Err(unsettled()),
};
self.settle(outcome.clone());
Poll::Ready(outcome)
}
}
#[must_use]
pub fn confirmation() -> (Confirmer, Confirmation) {
let (tx, rx) = oneshot::channel();
(
Confirmer(tx),
Confirmation {
rx,
outcome: None,
on_settle: None,
},
)
}
fn unsettled() -> FeedError {
FeedError::new(
FeedErrorKind::Retryable,
"the feed dropped the confirmation before settling it; the publish or close may or may \
not have finished",
)
}
#[cfg(test)]
mod tests {
use super::*;
struct EmptyFeed;
impl Feed for EmptyFeed {
fn backend(&self) -> &str {
"empty"
}
fn set_waker(&self, _waker: FeedWaker) {}
fn receive(&self, _max: usize) -> Result<Vec<Delivery>, FeedError> {
Ok(vec![])
}
fn renew(&self, receipt: Receipt) -> Result<(), FeedError> {
Err(FeedError::new(FeedErrorKind::Gone, receipt.to_string()))
}
fn acknowledge(&self, receipt: Receipt) -> Result<(), FeedError> {
Err(FeedError::new(FeedErrorKind::Gone, receipt.to_string()))
}
fn fail(&self, receipt: Receipt, _reason: &str) -> Result<(), FeedError> {
Err(FeedError::new(FeedErrorKind::Gone, receipt.to_string()))
}
fn publish(&self, _unit: Unit) -> Confirmation {
let (confirmer, confirmation) = confirmation();
confirmer.reject(FeedError::new(FeedErrorKind::Refused, "empty feed"));
confirmation
}
}
#[test]
fn feed_metrics_ctx_partition_detail_defaults_off_and_round_trips() {
let cx = FeedMetricsCtx::new(None);
assert!(cx.meter.is_none());
assert!(!cx.per_partition_detail);
let cx = FeedMetricsCtx::new(None).with_partition_detail(true);
assert!(cx.per_partition_detail);
}
#[test]
fn feed_is_dyn_compatible() {
fn shared<T: Send + Sync>(_: &T) {}
let feed: std::sync::Arc<dyn Feed> = std::sync::Arc::new(EmptyFeed);
shared(&feed);
assert!(feed.receive(1).unwrap().is_empty());
let err = feed.acknowledge(Receipt(7)).unwrap_err();
assert_eq!(err.kind, FeedErrorKind::Gone);
let mut confirmation = feed.publish(Unit::new("u", Vec::new()));
assert_eq!(
confirmation.try_outcome().unwrap().unwrap_err().kind,
FeedErrorKind::Refused
);
assert!(feed.read_depth(DepthOf::Feed).is_none());
assert!(!feed.hands_back());
feed.stop_receiving();
assert!(feed.hand_back(Receipt(7)).is_ok());
assert!(feed.close(Duration::from_secs(1)).is_none());
}
#[test]
fn a_feed_neither_hands_back_nor_closes_by_default() {
assert!(!EmptyFeed.hands_back());
assert!(EmptyFeed.hand_back(Receipt(1)).is_ok());
assert!(EmptyFeed.close(Duration::ZERO).is_none());
}
#[test]
fn a_feed_reports_no_depth_by_default() {
assert!(EmptyFeed.read_depth(DepthOf::Feed).is_none());
assert!(EmptyFeed.read_depth(DepthOf::DeadLetters).is_none());
}
#[test]
fn a_depth_reading_settles_once_and_repeats() {
let (reporter, mut reading) = depth_reading();
assert!(reading.try_outcome().is_none());
drop(reporter);
let kind = reading.try_outcome().unwrap().unwrap_err().kind;
assert_eq!(kind, FeedErrorKind::Retryable);
let (reporter, mut reading) = depth_reading();
let mut depth = Depth::new("q", true);
depth.delayed = 2;
reporter.report(Ok(depth.clone()));
assert_eq!(reading.try_outcome().unwrap().unwrap(), depth);
assert_eq!(reading.try_outcome().unwrap().unwrap(), depth);
}
#[test]
fn a_depth_is_empty_only_with_every_count_zero() {
let mut depth = Depth::new("q", false);
assert!(depth.is_empty());
depth.visible = u64::MAX;
depth.delayed = 1;
assert!(!depth.is_empty());
let mut depth = Depth::new("q", true);
depth.unsent = 1;
assert!(!depth.is_empty());
assert_eq!(
depth.to_string(),
"q: 0 visible, 0 delayed, 0 in flight, 1 unsent (approximate)"
);
assert_eq!(
Depth::new("q", false).to_string(),
"q: 0 visible, 0 delayed, 0 in flight, 0 unsent"
);
}
#[test]
fn a_dropped_confirmer_reads_as_retryable() {
let (confirmer, mut checked) = confirmation();
assert!(checked.try_outcome().is_none());
drop(confirmer);
let kind = checked.try_outcome().unwrap().unwrap_err().kind;
assert_eq!(kind, FeedErrorKind::Retryable);
let (confirmer, awaited) = confirmation();
drop(confirmer);
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
let kind = rt.block_on(awaited).unwrap_err().kind;
assert_eq!(kind, FeedErrorKind::Retryable);
}
#[test]
fn a_settled_confirmation_repeats_its_outcome() {
let (confirmer, mut confirmation) = confirmation();
assert!(confirmation.try_outcome().is_none());
confirmer.confirm();
assert!(matches!(confirmation.try_outcome(), Some(Ok(()))));
assert!(matches!(confirmation.try_outcome(), Some(Ok(()))));
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
assert!(rt.block_on(confirmation).is_ok());
}
#[test]
fn a_wake_before_the_park_ends_it_once() {
let (waker, park) = wake_channel();
assert!(!park.park(Duration::ZERO));
waker.wake();
waker.clone().wake();
assert!(park.park(Duration::ZERO));
assert!(!park.park(Duration::ZERO));
}
#[test]
fn a_park_with_no_waker_left_waits_out_its_timeout() {
let (waker, park) = wake_channel();
drop(waker);
let timeout = Duration::from_millis(20);
let start = std::time::Instant::now();
assert!(!park.park(timeout));
assert!(
start.elapsed() >= timeout,
"park returned after {:?}, before its timeout",
start.elapsed()
);
}
#[test]
fn only_a_retryable_feed_error_classes_as_retryable() {
let class = |kind| FeedError::new(kind, "x").class();
assert_eq!(class(FeedErrorKind::Retryable), ErrorClass::Retryable);
assert_eq!(class(FeedErrorKind::Gone), ErrorClass::Fatal);
assert_eq!(class(FeedErrorKind::Refused), ErrorClass::Fatal);
assert_eq!(class(FeedErrorKind::Fatal), ErrorClass::Fatal);
}
#[test]
fn a_settle_hook_shows_in_the_debug_output() {
let (_confirmer, confirmation) = confirmation();
let hooked = confirmation.on_settle(|_| {});
assert!(format!("{hooked:?}").contains("OnSettle"), "{hooked:?}");
}
#[test]
fn an_awaited_confirmation_repeats_its_outcome_to_a_check() {
let (confirmer, mut confirmation) = confirmation();
confirmer.confirm();
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
assert!(rt.block_on(&mut confirmation).is_ok());
assert!(matches!(confirmation.try_outcome(), Some(Ok(()))));
}
}