use std::{
collections::BTreeMap, error::Error as StdError, fmt, future::Future, marker::PhantomData,
sync::Arc, time::Duration,
};
use crate::codec::Codec;
use crate::runtime::publish_source::BrokerRegistration;
use crate::{Broker, DescribeServer, ServerSpec};
use tokio_util::task::TaskTracker;
use crate::runtime::dispatch::Delivery;
use crate::runtime::lifecycle::{BoxError, BrokerCell, BrokerLifecycle, ConnectedSlot};
use crate::runtime::metadata::HandlerMetadata;
use crate::runtime::middleware::{Identity, Stack};
use crate::runtime::publish::{PublishIdentity, PublishLayer, PublishStack};
use crate::runtime::router::RouterSink;
#[cfg(feature = "testing")]
use crate::testing::coordinator::TestHooks;
use super::scope::BrokerScope;
use super::{AppInfo, LifecycleHook, LifecyclePhase, Starter, StateInit};
pub struct RustStream<Layers = Identity, State = (), Pipeline = PublishIdentity, Phase = Wired> {
pub(super) info: AppInfo,
pub(super) brokers: Vec<RegisteredBroker>,
pub(super) starters: Vec<Starter<State>>,
pub(super) handlers: Vec<HandlerMetadata>,
pub(super) servers: BTreeMap<String, ServerSpec>,
pub(super) publish_pipeline: Pipeline,
pub(super) state_init: StateInit<State>,
pub(super) after_startup: Vec<LifecycleHook<State>>,
pub(super) on_shutdown: Vec<LifecycleHook<State>>,
pub(super) after_shutdown: Vec<LifecycleHook<State>>,
pub(super) shutdown_timeout: Option<Duration>,
pub(super) continuations: TaskTracker,
#[cfg(feature = "testing")]
pub(super) test_hooks: Arc<TestHooks>,
pub(super) global: Layers,
pub(super) phase: PhantomData<fn() -> Phase>,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Setup;
#[derive(Debug, Clone, Copy, Default)]
pub struct Wired;
pub(crate) struct RegisteredBroker {
pub(crate) lifecycle: Box<dyn BrokerLifecycle>,
pub(crate) label: Option<String>,
}
#[cfg(feature = "testing")]
pub(crate) struct TestParts<State> {
pub(crate) brokers: Vec<RegisteredBroker>,
pub(crate) starters: Vec<Starter<State>>,
pub(crate) state_init: StateInit<State>,
pub(crate) after_startup: Vec<LifecycleHook<State>>,
pub(crate) shutdown_timeout: Option<Duration>,
pub(crate) continuations: TaskTracker,
pub(crate) test_hooks: Arc<TestHooks>,
}
impl<Layers, State, Pipeline, Phase> fmt::Debug for RustStream<Layers, State, Pipeline, Phase> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RustStream")
.field("info", &self.info)
.field("brokers", &self.brokers.len())
.field("handlers", &self.handlers.len())
.finish_non_exhaustive()
}
}
impl RustStream<Identity, (), PublishIdentity, Setup> {
#[must_use]
pub fn new(info: AppInfo) -> Self {
Self {
info,
brokers: Vec::new(),
starters: Vec::new(),
handlers: Vec::new(),
servers: BTreeMap::new(),
publish_pipeline: PublishIdentity,
state_init: Box::new(|| Box::pin(async { Ok(()) })),
after_startup: Vec::new(),
on_shutdown: Vec::new(),
after_shutdown: Vec::new(),
shutdown_timeout: None,
continuations: TaskTracker::new(),
#[cfg(feature = "testing")]
test_hooks: Arc::new(TestHooks::detached()),
global: Identity,
phase: PhantomData,
}
}
}
impl<Layers, State, Pipeline> RustStream<Layers, State, Pipeline, Setup> {
#[must_use]
pub fn layer<N>(self, layer: N) -> RustStream<Stack<N, Layers>, State, Pipeline, Setup> {
RustStream {
info: self.info,
brokers: self.brokers,
starters: self.starters,
handlers: self.handlers,
servers: self.servers,
publish_pipeline: self.publish_pipeline,
state_init: self.state_init,
after_startup: self.after_startup,
on_shutdown: self.on_shutdown,
after_shutdown: self.after_shutdown,
shutdown_timeout: self.shutdown_timeout,
continuations: self.continuations,
#[cfg(feature = "testing")]
test_hooks: self.test_hooks,
global: Stack::new(layer, self.global),
phase: PhantomData,
}
}
#[must_use]
pub fn on_startup<F, Fut, State2, E>(
self,
hook: F,
) -> RustStream<Layers, State2, Pipeline, Setup>
where
F: FnOnce(State) -> Fut + Send + 'static,
Fut: Future<Output = Result<State2, E>> + Send,
State: Send + 'static,
State2: Send + Sync + 'static,
E: StdError + Send + Sync + 'static,
{
assert!(
self.after_startup.is_empty()
&& self.on_shutdown.is_empty()
&& self.after_shutdown.is_empty(),
"on_startup must be called before lifecycle hooks are registered: hooks close over \
the state type and cannot be carried across the state change"
);
let prev = self.state_init;
RustStream {
info: self.info,
brokers: self.brokers,
starters: Vec::new(),
handlers: self.handlers,
servers: self.servers,
publish_pipeline: self.publish_pipeline,
state_init: Box::new(move || {
Box::pin(async move {
let prev_state = prev().await?;
hook(prev_state).await.map_err(|e| Box::new(e) as BoxError)
})
}),
after_startup: Vec::new(),
on_shutdown: Vec::new(),
after_shutdown: Vec::new(),
shutdown_timeout: self.shutdown_timeout,
continuations: self.continuations,
#[cfg(feature = "testing")]
test_hooks: self.test_hooks,
global: self.global,
phase: PhantomData,
}
}
#[must_use]
pub fn publish_layer<M>(
self,
middleware: M,
) -> RustStream<Layers, State, PublishStack<M, Pipeline>, Setup>
where
M: PublishLayer + Clone + 'static,
{
RustStream {
info: self.info,
brokers: self.brokers,
starters: self.starters,
handlers: self.handlers,
servers: self.servers,
publish_pipeline: PublishStack::new(middleware, self.publish_pipeline),
state_init: self.state_init,
after_startup: self.after_startup,
on_shutdown: self.on_shutdown,
after_shutdown: self.after_shutdown,
shutdown_timeout: self.shutdown_timeout,
continuations: self.continuations,
#[cfg(feature = "testing")]
test_hooks: self.test_hooks,
global: self.global,
phase: PhantomData,
}
}
}
impl<Layers, State, Pipeline, Phase> RustStream<Layers, State, Pipeline, Phase> {
fn into_phase<Q>(self) -> RustStream<Layers, State, Pipeline, Q> {
RustStream {
info: self.info,
brokers: self.brokers,
starters: self.starters,
handlers: self.handlers,
servers: self.servers,
publish_pipeline: self.publish_pipeline,
state_init: self.state_init,
after_startup: self.after_startup,
on_shutdown: self.on_shutdown,
after_shutdown: self.after_shutdown,
shutdown_timeout: self.shutdown_timeout,
continuations: self.continuations,
#[cfg(feature = "testing")]
test_hooks: self.test_hooks,
global: self.global,
phase: PhantomData,
}
}
#[must_use]
pub fn after_startup<F, Fut, E>(self, hook: F) -> Self
where
State: Send + Sync + 'static,
F: FnOnce(Arc<State>) -> Fut + Send + 'static,
Fut: Future<Output = Result<(), E>> + Send,
E: StdError + Send + Sync + 'static,
{
self.push_lifecycle_hook(LifecyclePhase::AfterStartup, hook)
}
#[must_use]
pub fn on_shutdown<F, Fut, E>(self, hook: F) -> Self
where
State: Send + Sync + 'static,
F: FnOnce(Arc<State>) -> Fut + Send + 'static,
Fut: Future<Output = Result<(), E>> + Send,
E: StdError + Send + Sync + 'static,
{
self.push_lifecycle_hook(LifecyclePhase::OnShutdown, hook)
}
#[must_use]
pub fn after_shutdown<F, Fut, E>(self, hook: F) -> Self
where
State: Send + Sync + 'static,
F: FnOnce(Arc<State>) -> Fut + Send + 'static,
Fut: Future<Output = Result<(), E>> + Send,
E: StdError + Send + Sync + 'static,
{
self.push_lifecycle_hook(LifecyclePhase::AfterShutdown, hook)
}
fn push_lifecycle_hook<F, Fut, E>(mut self, phase: LifecyclePhase, hook: F) -> Self
where
State: Send + Sync + 'static,
F: FnOnce(Arc<State>) -> Fut + Send + 'static,
Fut: Future<Output = Result<(), E>> + Send,
E: StdError + Send + Sync + 'static,
{
let boxed: LifecycleHook<State> = Box::new(move |state| {
Box::pin(async move { hook(state).await.map_err(|e| Box::new(e) as BoxError) })
});
match phase {
LifecyclePhase::AfterStartup => self.after_startup.push(boxed),
LifecyclePhase::OnShutdown => self.on_shutdown.push(boxed),
LifecyclePhase::AfterShutdown => self.after_shutdown.push(boxed),
}
self
}
#[must_use]
pub fn shutdown_timeout(mut self, timeout: Duration) -> Self {
self.shutdown_timeout = Some(timeout);
self
}
#[must_use]
pub fn register_broker<R>(mut self, broker: R) -> Self
where
R: BrokerRegistration,
{
let (broker, slot) = broker.into_parts();
self.brokers.push(RegisteredBroker {
lifecycle: Box::new(BrokerCell { broker, slot }),
label: None,
});
self
}
#[must_use]
pub fn server(mut self, name: impl Into<String>, spec: ServerSpec) -> Self {
self.servers.insert(name.into(), spec);
self
}
#[must_use]
pub fn with_broker<R, F>(
self,
broker: R,
build: F,
) -> RustStream<Layers, State, Pipeline, Wired>
where
R: BrokerRegistration,
Layers: Clone,
Pipeline: Clone,
State: Send + Sync + 'static,
F: FnOnce(&mut BrokerScope<R::Broker, Layers, (), State, Pipeline>),
{
let mut this = self.into_phase::<Wired>();
let (broker, slot) = broker.into_parts();
let mut scope = this.new_scope(broker, slot, ());
build(&mut scope);
this.collect_scope(scope, None);
this
}
#[must_use]
pub fn with_broker_codec<R, C, F>(
self,
broker: R,
codec: C,
build: F,
) -> RustStream<Layers, State, Pipeline, Wired>
where
R: BrokerRegistration,
C: Codec + Clone + 'static,
Layers: Clone,
Pipeline: Clone,
State: Send + Sync + 'static,
F: FnOnce(&mut BrokerScope<R::Broker, Layers, C, State, Pipeline>),
{
let mut this = self.into_phase::<Wired>();
let (broker, slot) = broker.into_parts();
let mut scope = this.new_scope(broker, slot, codec);
build(&mut scope);
this.collect_scope(scope, None);
this
}
#[must_use]
pub fn with_broker_labeled<R, F>(
self,
label: impl Into<String>,
broker: R,
build: F,
) -> RustStream<Layers, State, Pipeline, Wired>
where
R: BrokerRegistration,
R::Broker: DescribeServer,
Layers: Clone,
Pipeline: Clone,
State: Send + Sync + 'static,
F: FnOnce(&mut BrokerScope<R::Broker, Layers, (), State, Pipeline>),
{
let mut this = self.into_phase::<Wired>();
let (broker, slot) = broker.into_parts();
let label = this.record_server(label, &broker);
let mut scope = this.new_scope(broker, slot, ());
build(&mut scope);
this.collect_scope(scope, Some(label));
this
}
#[must_use]
pub fn with_broker_labeled_codec<R, C, F>(
self,
label: impl Into<String>,
broker: R,
codec: C,
build: F,
) -> RustStream<Layers, State, Pipeline, Wired>
where
R: BrokerRegistration,
R::Broker: DescribeServer,
C: Codec + Clone + 'static,
Layers: Clone,
Pipeline: Clone,
State: Send + Sync + 'static,
F: FnOnce(&mut BrokerScope<R::Broker, Layers, C, State, Pipeline>),
{
let mut this = self.into_phase::<Wired>();
let (broker, slot) = broker.into_parts();
let label = this.record_server(label, &broker);
let mut scope = this.new_scope(broker, slot, codec);
build(&mut scope);
this.collect_scope(scope, Some(label));
this
}
fn record_server<B: DescribeServer>(&mut self, label: impl Into<String>, broker: &B) -> String {
let label = label.into();
self.servers
.entry(label.clone())
.or_insert_with(|| broker.describe_server());
label
}
fn new_scope<B, C>(
&self,
broker: B,
slot: ConnectedSlot<B>,
codec: C,
) -> BrokerScope<B, Layers, C, State, Pipeline>
where
B: Broker + 'static,
Layers: Clone,
Pipeline: Clone,
State: Send + Sync + 'static,
{
BrokerScope {
broker,
slot,
startup_hooks: Vec::new(),
sink: RouterSink::new(),
pipeline: self.publish_pipeline.clone(),
retry_publisher: None,
global: self.global.clone(),
codec,
}
}
fn collect_scope<B, C>(
&mut self,
scope: BrokerScope<B, Layers, C, State, Pipeline>,
label: Option<String>,
) where
B: Broker + 'static,
State: Send + Sync + 'static,
{
let BrokerScope {
broker,
slot,
startup_hooks,
sink,
retry_publisher,
..
} = scope;
self.after_startup.extend(startup_hooks);
#[cfg(feature = "testing")]
let delivery = Arc::new(Delivery::instrumented(
retry_publisher,
self.continuations.clone(),
self.test_hooks.clone(),
self.brokers.len(),
));
#[cfg(not(feature = "testing"))]
let delivery = Arc::new(Delivery::detached(
retry_publisher,
self.continuations.clone(),
));
let (starters, handlers) = sink.into_parts();
for (bound, meta) in starters.into_iter().zip(handlers) {
let slot = Arc::clone(&slot);
let delivery = delivery.clone();
self.starters.push(Box::new(move |state, shutdown, token| {
let connected = slot
.lock()
.expect("connected slot mutex poisoned")
.clone()
.expect("brokers connect before subscriptions open");
bound(connected, state, delivery, shutdown, token)
}));
self.handlers.push(meta);
}
self.brokers.push(RegisteredBroker {
lifecycle: Box::new(BrokerCell { broker, slot }),
label,
});
}
#[must_use]
pub fn handlers(&self) -> &[HandlerMetadata] {
&self.handlers
}
#[must_use]
pub fn info(&self) -> &AppInfo {
&self.info
}
#[must_use]
pub fn servers(&self) -> &BTreeMap<String, ServerSpec> {
&self.servers
}
#[cfg(feature = "testing")]
pub(crate) fn into_test_parts(self) -> TestParts<State> {
TestParts {
brokers: self.brokers,
starters: self.starters,
state_init: self.state_init,
after_startup: self.after_startup,
shutdown_timeout: self.shutdown_timeout,
continuations: self.continuations,
test_hooks: self.test_hooks,
}
}
}