#![allow(clippy::type_complexity, clippy::should_implement_trait)]
use alloc::boxed::Box;
use core::marker::PhantomData;
use crate::batch::{Batch, BatchPolicy, BatchStage, BatchStageBytes, ByteSize};
use crate::driver::{RunStats, SyncDriver};
use crate::emit::{Emit, EmitError};
use crate::error::{Error, ErrorPolicy, Result, StageError, StageFailure};
use crate::sink::Sink;
use crate::source::{IterSource, Source};
use crate::stage::Stage;
use crate::stage_id::StageId;
#[doc(hidden)]
pub enum StageOp<T> {
Process(T),
Flush,
Close,
}
#[doc(hidden)]
pub type BoxedStageFn<T> = Box<dyn FnMut(StageOp<T>) -> Result<()> + Send + 'static>;
struct StageEmit<'a, U> {
next_fn: &'a mut BoxedStageFn<U>,
cached_err: Option<Error>,
}
#[doc(hidden)]
#[cfg(feature = "std")]
pub enum DeadLetterOp {
Send(StageFailure),
Flush,
Close,
}
#[doc(hidden)]
#[cfg(feature = "std")]
pub type DeadLetterFn = Box<dyn FnMut(DeadLetterOp) -> Result<()> + Send + 'static>;
#[cfg(feature = "std")]
#[derive(Clone, Default)]
pub(crate) struct DeadLetter {
inner: std::sync::Arc<std::sync::Mutex<Option<DeadLetterFn>>>,
}
#[cfg(feature = "std")]
impl DeadLetter {
fn new() -> Self {
Self::default()
}
fn install(&self, f: DeadLetterFn) {
*self.inner.lock().expect("dead-letter mutex poisoned") = Some(f);
}
fn route(&self, failure: StageFailure) -> Result<()> {
let mut guard = self.inner.lock().expect("dead-letter mutex poisoned");
match guard.as_mut() {
Some(f) => f(DeadLetterOp::Send(failure)),
None => Ok(()),
}
}
fn finish(&self) -> Result<()> {
let mut guard = self.inner.lock().expect("dead-letter mutex poisoned");
if let Some(f) = guard.as_mut() {
f(DeadLetterOp::Flush)?;
f(DeadLetterOp::Close)?;
}
Ok(())
}
}
#[cfg(not(feature = "std"))]
#[derive(Clone, Default)]
pub(crate) struct DeadLetter;
#[cfg(not(feature = "std"))]
impl DeadLetter {
fn new() -> Self {
Self
}
fn route(&self, _failure: StageFailure) -> Result<()> {
Ok(())
}
fn finish(&self) -> Result<()> {
Ok(())
}
}
impl<'a, U> Emit for StageEmit<'a, U> {
type Item = U;
fn emit(&mut self, item: U) -> core::result::Result<(), EmitError> {
if self.cached_err.is_some() {
return Err(EmitError::Closed);
}
match (self.next_fn)(StageOp::Process(item)) {
Ok(()) => Ok(()),
Err(e) => {
self.cached_err = Some(e);
Err(EmitError::Closed)
}
}
}
}
pub struct Pipeline<S>
where
S: Source,
{
pub(crate) source: S,
pub(crate) source_id: StageId,
pub(crate) stage_fn: BoxedStageFn<S::Item>,
pub(crate) dead_letter: DeadLetter,
}
impl<S> Pipeline<S>
where
S: Source + 'static,
S::Item: Send + 'static,
S::Error: Send + 'static,
{
pub fn from_source(
source: S,
) -> PipelineBuilder<
S::Item,
S,
impl FnOnce(BoxedStageFn<S::Item>) -> BoxedStageFn<S::Item> + Send + 'static,
> {
PipelineBuilder {
source,
source_id: StageId::new("source"),
finalize: identity_finalize::<S::Item>,
error_policy: ErrorPolicy::FailFast,
pending_stage_id: None,
dead_letter: DeadLetter::new(),
_marker: PhantomData,
}
}
pub fn run(self) -> Result<RunStats> {
SyncDriver::new().run(self)
}
pub fn run_with<D>(self, driver: D) -> Result<RunStats>
where
D: crate::driver::Driver,
S: Send,
S::Item: Send,
S::Error: Send,
{
driver.run(self)
}
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
pub fn run_threaded(self) -> Result<RunStats>
where
S: Send,
S::Item: Send,
S::Error: Send,
{
crate::driver::ThreadedDriver::new().run(self)
}
}
impl Pipeline<IterSource<core::iter::Empty<()>>> {
pub fn from_iter<II>(
iter: II,
) -> PipelineBuilder<
II::Item,
IterSource<II::IntoIter>,
impl FnOnce(BoxedStageFn<II::Item>) -> BoxedStageFn<II::Item> + Send + 'static,
>
where
II: IntoIterator,
II::Item: Send + 'static,
II::IntoIter: Send + 'static,
{
Pipeline::from_source(IterSource::new(iter))
}
}
pub struct PipelineBuilder<T, S, Acc>
where
S: Source,
Acc: FnOnce(BoxedStageFn<T>) -> BoxedStageFn<S::Item> + Send + 'static,
{
source: S,
source_id: StageId,
finalize: Acc,
error_policy: ErrorPolicy,
pending_stage_id: Option<StageId>,
dead_letter: DeadLetter,
_marker: PhantomData<fn() -> T>,
}
fn identity_finalize<T: 'static + Send>(f: BoxedStageFn<T>) -> BoxedStageFn<T> {
f
}
fn handle_stage_error<E: StageError>(
policy: ErrorPolicy,
stage_id: StageId,
err: E,
dead_letter: &DeadLetter,
) -> Result<()> {
match policy {
ErrorPolicy::FailFast => Err(Error::Stage {
stage: stage_id,
source: Box::new(err),
}),
ErrorPolicy::Continue => Ok(()),
ErrorPolicy::DeadLetter => dead_letter.route(StageFailure::new(stage_id, Box::new(err))),
}
}
impl<T, S, Acc> PipelineBuilder<T, S, Acc>
where
S: Source + 'static,
S::Item: Send + 'static,
S::Error: Send + 'static,
T: Send + 'static,
Acc: FnOnce(BoxedStageFn<T>) -> BoxedStageFn<S::Item> + Send + 'static,
{
#[must_use]
pub fn stage_id<I: Into<StageId>>(mut self, id: I) -> Self {
self.pending_stage_id = Some(id.into());
self
}
#[must_use]
pub fn on_error(mut self, policy: ErrorPolicy) -> Self {
self.error_policy = policy;
self
}
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
pub fn dead_letter<Sk>(self, sink: Sk) -> Self
where
Sk: Sink<Item = StageFailure> + Send + 'static,
Sk::Error: 'static,
{
let mut sink = sink;
let sink_id = StageId::new("dead_letter");
let f: DeadLetterFn = Box::new(move |op| match op {
DeadLetterOp::Send(failure) => sink.write(failure).map_err(|e| Error::Sink {
stage: sink_id,
source: Box::new(e),
}),
DeadLetterOp::Flush => sink.flush().map_err(|e| Error::Sink {
stage: sink_id,
source: Box::new(e),
}),
DeadLetterOp::Close => sink.close().map_err(|e| Error::Sink {
stage: sink_id,
source: Box::new(e),
}),
});
self.dead_letter.install(f);
self
}
pub fn map<U, F>(
self,
mut f: F,
) -> PipelineBuilder<U, S, impl FnOnce(BoxedStageFn<U>) -> BoxedStageFn<S::Item> + Send + 'static>
where
U: Send + 'static,
F: FnMut(T) -> U + Send + 'static,
{
let old_finalize = self.finalize;
let new_finalize = move |next: BoxedStageFn<U>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => next(StageOp::Process(f(item))),
StageOp::Flush => next(StageOp::Flush),
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn filter<F>(
self,
mut pred: F,
) -> PipelineBuilder<T, S, impl FnOnce(BoxedStageFn<T>) -> BoxedStageFn<S::Item> + Send + 'static>
where
F: FnMut(&T) -> bool + Send + 'static,
{
let old_finalize = self.finalize;
let new_finalize = move |next: BoxedStageFn<T>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => {
if pred(&item) {
next(StageOp::Process(item))
} else {
Ok(())
}
}
StageOp::Flush => next(StageOp::Flush),
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn filter_map<U, F>(
self,
mut f: F,
) -> PipelineBuilder<U, S, impl FnOnce(BoxedStageFn<U>) -> BoxedStageFn<S::Item> + Send + 'static>
where
U: Send + 'static,
F: FnMut(T) -> Option<U> + Send + 'static,
{
let old_finalize = self.finalize;
let new_finalize = move |next: BoxedStageFn<U>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => match f(item) {
Some(out) => next(StageOp::Process(out)),
None => Ok(()),
},
StageOp::Flush => next(StageOp::Flush),
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn flat_map<U, F, II>(
self,
mut f: F,
) -> PipelineBuilder<U, S, impl FnOnce(BoxedStageFn<U>) -> BoxedStageFn<S::Item> + Send + 'static>
where
U: Send + 'static,
II: IntoIterator<Item = U>,
F: FnMut(T) -> II + Send + 'static,
{
let old_finalize = self.finalize;
let new_finalize = move |next: BoxedStageFn<U>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => {
for out in f(item) {
next(StageOp::Process(out))?;
}
Ok(())
}
StageOp::Flush => next(StageOp::Flush),
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn inspect<F>(
self,
mut f: F,
) -> PipelineBuilder<T, S, impl FnOnce(BoxedStageFn<T>) -> BoxedStageFn<S::Item> + Send + 'static>
where
F: FnMut(&T) + Send + 'static,
{
let old_finalize = self.finalize;
let new_finalize = move |next: BoxedStageFn<T>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => {
f(&item);
next(StageOp::Process(item))
}
StageOp::Flush => next(StageOp::Flush),
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn try_map<U, F, E>(
self,
mut f: F,
) -> PipelineBuilder<U, S, impl FnOnce(BoxedStageFn<U>) -> BoxedStageFn<S::Item> + Send + 'static>
where
U: Send + 'static,
E: StageError,
F: FnMut(T) -> core::result::Result<U, E> + Send + 'static,
{
let stage_id = self.pending_stage_id.unwrap_or(StageId::new("try_map"));
let policy = self.error_policy;
let old_finalize = self.finalize;
let dead_letter = self.dead_letter.clone();
let new_finalize = move |next: BoxedStageFn<U>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => match f(item) {
Ok(out) => next(StageOp::Process(out)),
Err(e) => handle_stage_error(policy, stage_id, e, &dead_letter),
},
StageOp::Flush => next(StageOp::Flush),
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn stage<St>(
self,
mut stage: St,
) -> PipelineBuilder<
St::Output,
S,
impl FnOnce(BoxedStageFn<St::Output>) -> BoxedStageFn<S::Item> + Send + 'static,
>
where
St: Stage<Input = T> + Send + 'static,
St::Output: Send + 'static,
St::Error: 'static,
{
let stage_id = self.pending_stage_id.unwrap_or(StageId::new("stage"));
let policy = self.error_policy;
let old_finalize = self.finalize;
let dead_letter = self.dead_letter.clone();
let new_finalize = move |next: BoxedStageFn<St::Output>| -> BoxedStageFn<S::Item> {
let mut next = next;
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => {
let mut adapter = StageEmit {
next_fn: &mut next,
cached_err: None,
};
let stage_result = stage.process(item, &mut adapter);
if let Some(err) = adapter.cached_err {
return Err(err);
}
match stage_result {
Ok(()) => Ok(()),
Err(e) => handle_stage_error(policy, stage_id, e, &dead_letter),
}
}
StageOp::Flush => {
let mut adapter = StageEmit {
next_fn: &mut next,
cached_err: None,
};
let stage_result = stage.flush(&mut adapter);
if let Some(err) = adapter.cached_err {
return Err(err);
}
match stage_result {
Ok(()) => next(StageOp::Flush),
Err(e) => handle_stage_error(policy, stage_id, e, &dead_letter),
}
}
StageOp::Close => next(StageOp::Close),
});
old_finalize(t_fn)
};
PipelineBuilder {
source: self.source,
source_id: self.source_id,
finalize: new_finalize,
error_policy: self.error_policy,
pending_stage_id: None,
dead_letter: self.dead_letter,
_marker: PhantomData,
}
}
pub fn batch(
mut self,
policy: BatchPolicy,
) -> PipelineBuilder<
Batch<T>,
S,
impl FnOnce(BoxedStageFn<Batch<T>>) -> BoxedStageFn<S::Item> + Send + 'static,
> {
assert!(
policy.has_trigger(),
"BatchPolicy must have at least one trigger configured"
);
assert!(
policy.bytes_limit().is_none(),
"BatchPolicy::max_bytes requires PipelineBuilder::batch_bytes (T: ByteSize)"
);
let id = self
.pending_stage_id
.take()
.unwrap_or(StageId::new("batch"));
self.stage_id(id).stage(BatchStage::<T>::new(policy))
}
pub fn batch_bytes(
mut self,
policy: BatchPolicy,
) -> PipelineBuilder<
Batch<T>,
S,
impl FnOnce(BoxedStageFn<Batch<T>>) -> BoxedStageFn<S::Item> + Send + 'static,
>
where
T: ByteSize,
{
assert!(
policy.has_trigger(),
"BatchPolicy must have at least one trigger configured"
);
let id = self
.pending_stage_id
.take()
.unwrap_or(StageId::new("batch"));
self.stage_id(id).stage(BatchStageBytes::<T>::new(policy))
}
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
pub fn window(
mut self,
policy: crate::window::WindowPolicy,
) -> PipelineBuilder<
crate::window::Window<T>,
S,
impl FnOnce(BoxedStageFn<crate::window::Window<T>>) -> BoxedStageFn<S::Item> + Send + 'static,
>
where
T: Clone,
{
let id = self
.pending_stage_id
.take()
.unwrap_or(StageId::new("window"));
self.stage_id(id).stage(
crate::window::WindowStage::<T, crate::window::SystemClock>::new(
policy,
crate::window::SystemClock,
),
)
}
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
pub fn window_with<C>(
mut self,
policy: crate::window::WindowPolicy,
clock: C,
) -> PipelineBuilder<
crate::window::Window<T>,
S,
impl FnOnce(BoxedStageFn<crate::window::Window<T>>) -> BoxedStageFn<S::Item> + Send + 'static,
>
where
T: Clone,
C: crate::window::Clock + 'static,
{
let id = self
.pending_stage_id
.take()
.unwrap_or(StageId::new("window"));
self.stage_id(id)
.stage(crate::window::WindowStage::<T, C>::new(policy, clock))
}
pub fn sink<Sk>(self, sink: Sk) -> Pipeline<S>
where
Sk: Sink<Item = T> + Send + 'static,
Sk::Error: 'static,
{
let mut sink = sink;
let sink_id = self.pending_stage_id.unwrap_or(StageId::new("sink"));
let t_fn: BoxedStageFn<T> = Box::new(move |op| match op {
StageOp::Process(item) => sink.write(item).map_err(|e| Error::Sink {
stage: sink_id,
source: Box::new(e),
}),
StageOp::Flush => sink.flush().map_err(|e| Error::Sink {
stage: sink_id,
source: Box::new(e),
}),
StageOp::Close => sink.close().map_err(|e| Error::Sink {
stage: sink_id,
source: Box::new(e),
}),
});
let item_fn: BoxedStageFn<S::Item> = (self.finalize)(t_fn);
Pipeline {
source: self.source,
source_id: self.source_id,
stage_fn: item_fn,
dead_letter: self.dead_letter,
}
}
}
pub(crate) fn run_sync<S>(mut pipeline: Pipeline<S>) -> Result<RunStats>
where
S: Source + 'static,
S::Item: 'static,
S::Error: 'static,
{
#[cfg(feature = "std")]
let start = std::time::Instant::now();
let mut stats = RunStats::default();
loop {
match pipeline.source.pull() {
Ok(Some(item)) => {
stats.items_in = stats.items_in.saturating_add(1);
(pipeline.stage_fn)(StageOp::Process(item))?;
}
Ok(None) => break,
Err(e) => {
return Err(Error::Source {
stage: pipeline.source_id,
source: Box::new(e),
});
}
}
}
(pipeline.stage_fn)(StageOp::Flush)?;
(pipeline.stage_fn)(StageOp::Close)?;
let _ = pipeline.source.close();
pipeline.dead_letter.finish()?;
#[cfg(feature = "std")]
{
stats.duration = start.elapsed();
}
Ok(stats)
}