use std::cell::OnceCell;
use std::future::Future;
use std::pin::Pin;
use bytes::BytesMut;
use crate::{Field, FieldMut, HeaderMap, IncomingMessage};
use super::dispatch::Delivery;
use super::failure::FailurePolicy;
use super::handler::{HandlerOutcome, HandlerResult};
use super::main_runtime::MainRuntime;
use super::shutdown::Shutdown;
pub(crate) trait HeaderSource {
fn headers(&self) -> &HeaderMap;
}
impl HeaderSource for HeaderMap {
fn headers(&self) -> &HeaderMap {
self
}
}
#[derive(Clone, Copy)]
pub(crate) enum DeliveryHeaders<'a> {
Ready(&'a HeaderMap),
Source(&'a (dyn HeaderSource + Sync)),
}
impl<'a> DeliveryHeaders<'a> {
pub(crate) fn get(self) -> &'a HeaderMap {
match self {
Self::Ready(headers) => headers,
Self::Source(source) => source.headers(),
}
}
}
impl<'a> From<&'a HeaderMap> for DeliveryHeaders<'a> {
fn from(headers: &'a HeaderMap) -> Self {
Self::Ready(headers)
}
}
pub(crate) struct FromDelivery<'m, M>(pub(crate) &'m M);
impl<M: IncomingMessage> HeaderSource for FromDelivery<'_, M> {
fn headers(&self) -> &HeaderMap {
self.0.headers()
}
}
type Continuation = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum OutcomeKind {
Ack,
Drop,
Retry,
RetryAfter,
}
impl OutcomeKind {
fn of(outcome: HandlerResult) -> Self {
match outcome {
HandlerResult::Ack => Self::Ack,
HandlerResult::Nack { requeue: false } => Self::Drop,
HandlerResult::Nack { requeue: true } => Self::Retry,
HandlerResult::NackAfter { .. } => Self::RetryAfter,
}
}
}
struct AfterHook {
gate: Option<OutcomeKind>,
fut: Continuation,
}
pub struct Context<'a, C = (), S = ()> {
name: &'a str,
source: &'a (dyn HeaderSource + Sync),
original: OnceCell<&'a HeaderMap>,
modified: Option<HeaderMap>,
state: &'a S,
cx: C,
delivery: &'a Delivery<C>,
after: Vec<AfterHook>,
failfast: Option<&'a Shutdown>,
encode: Option<&'a mut BytesMut>,
decode: FailurePolicy,
#[cfg(any(feature = "testing", feature = "otel"))]
decode_failed: bool,
}
impl<C, S> std::fmt::Debug for Context<'_, C, S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Context")
.field("name", &self.name)
.field("after_hooks", &self.after.len())
.finish_non_exhaustive()
}
}
impl<'a, C, S> Context<'a, C, S> {
pub(crate) fn new(
name: &'a str,
source: &'a (dyn HeaderSource + Sync),
state: &'a S,
cx: C,
delivery: &'a Delivery<C>,
) -> Self {
Self {
name,
source,
original: OnceCell::new(),
modified: None,
state,
cx,
delivery,
after: Vec::new(),
failfast: None,
encode: None,
decode: FailurePolicy::Drop,
#[cfg(any(feature = "testing", feature = "otel"))]
decode_failed: false,
}
}
#[cfg(any(feature = "testing", feature = "otel"))]
pub(crate) fn mark_decode_failed(&mut self) {
self.decode_failed = true;
}
#[cfg(feature = "otel")]
pub(crate) fn decode_failed(&self) -> bool {
self.decode_failed
}
#[cfg(feature = "testing")]
pub(crate) fn took_decode_failed(&mut self) -> bool {
std::mem::take(&mut self.decode_failed)
}
#[must_use]
pub(crate) fn with_failfast(mut self, failfast: &'a Shutdown) -> Self {
self.failfast = Some(failfast);
self
}
pub(crate) fn with_encode_buffer(mut self, buf: &'a mut BytesMut) -> Self {
self.encode = Some(buf);
self
}
pub(crate) fn take_encode_buffer(&mut self) -> Option<&'a mut BytesMut> {
self.encode.take()
}
pub(crate) fn with_decode_policy(mut self, decode: FailurePolicy) -> Self {
self.decode = decode;
self
}
#[must_use]
pub fn decode_policy(&self) -> FailurePolicy {
self.decode
}
pub(crate) fn fail_fast(&self, reason: &str) {
if let Some(failfast) = self.failfast {
failfast.signal(self.name, reason);
}
}
#[must_use]
pub fn name(&self) -> &str {
self.name
}
#[must_use]
pub fn main_runtime(&self) -> &MainRuntime {
&self.delivery.runtime
}
fn original(&self) -> &'a HeaderMap {
self.original.get_or_init(|| self.source.headers())
}
#[must_use]
pub fn headers(&self) -> &HeaderMap {
self.modified
.as_ref()
.map_or_else(|| self.original(), |modified| modified)
}
pub(crate) fn delivery_headers(&self) -> DeliveryHeaders<'_> {
if let Some(modified) = &self.modified {
return DeliveryHeaders::Ready(modified);
}
self.original
.get()
.map_or(DeliveryHeaders::Source(self.source), |original| {
DeliveryHeaders::Ready(original)
})
}
pub fn headers_mut(&mut self) -> &mut HeaderMap {
let original = self.original();
self.modified.get_or_insert_with(|| original.clone())
}
#[must_use]
pub fn state(&self) -> &S {
self.state
}
pub fn context<K: Field<C>>(&self, key: K) -> K::Value<'_> {
key.get(&self.cx)
}
pub(crate) fn cx_ref(&self) -> &C {
&self.cx
}
pub fn set<K: FieldMut<C>>(&mut self, key: K, value: K::Owned) {
key.set(&mut self.cx, value);
}
#[allow(clippy::needless_pass_by_value)]
pub fn after(&mut self, outcome: HandlerOutcome) -> After<'_, 'a, C, S> {
After {
ctx: self,
gate: Some(OutcomeKind::of(outcome.outcome())),
}
}
pub fn after_ack(&mut self, fut: impl Future<Output = ()> + Send + 'static) {
self.after.push(AfterHook {
gate: Some(OutcomeKind::Ack),
fut: Box::pin(fut),
});
}
pub fn after_settle(&mut self, fut: impl Future<Output = ()> + Send + 'static) {
self.after.push(AfterHook {
gate: None,
fut: Box::pin(fut),
});
}
pub(crate) fn has_hooks(&self) -> bool {
!self.after.is_empty()
}
pub(crate) fn take_hooks_for(&mut self, outcome: HandlerResult) -> Vec<Continuation> {
let kind = OutcomeKind::of(outcome);
let mut runnable = Vec::new();
let mut kept = Vec::new();
for hook in self.after.drain(..) {
if hook.gate.is_none_or(|gate| gate == kind) {
runnable.push(hook.fut);
} else {
kept.push(hook);
}
}
self.after = kept;
runnable
}
pub(crate) fn take_settle_hooks(&mut self) -> Vec<Continuation> {
let mut runnable = Vec::new();
let mut kept = Vec::new();
for hook in self.after.drain(..) {
if hook.gate.is_none() {
runnable.push(hook.fut);
} else {
kept.push(hook);
}
}
self.after = kept;
runnable
}
pub(crate) fn delivery(&self) -> &'a Delivery<C> {
self.delivery
}
pub(crate) fn subscription(&self) -> &'a str {
self.name
}
}
#[must_use = "call `.then(fut)` to register the post-settle hook"]
pub struct After<'ctx, 'a, C = (), S = ()> {
ctx: &'ctx mut Context<'a, C, S>,
gate: Option<OutcomeKind>,
}
impl<C, S> std::fmt::Debug for After<'_, '_, C, S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("After").field("gate", &self.gate).finish()
}
}
impl<C, S> After<'_, '_, C, S> {
pub fn then(self, fut: impl Future<Output = ()> + Send + 'static) {
self.ctx.after.push(AfterHook {
gate: self.gate,
fut: Box::pin(fut),
});
}
}
#[cfg(test)]
mod tests;