pub mod cascaded;
pub mod parallel;
pub mod runner;
pub mod s2s;
pub mod source_pump;
pub mod task;
#[cfg(test)]
pub(crate) mod s2s_test_mocks;
use std::sync::Arc;
use async_trait::async_trait;
use tokio::task::JoinSet;
use crate::error::Result;
use crate::processor::frame::{Direction, Frame, StartParams};
use crate::processor::runtime::{channel, ProcessorTx, NORMAL_CHAN_CAP};
use crate::processor::{
into_topology, run_processor, Envelope, FrameProcessor, Link, ProcessorSetup, StopReason,
Topology,
};
pub use cascaded::{
build_cascaded_call_with_observers, build_cascaded_pipeline, build_cascaded_task,
build_cascaded_task_with_observers, CascadedConfig, CascadedTask, ContextSummarizer,
SummarizerConfig,
};
pub use parallel::ParallelPipeline;
pub use runner::PipelineRunner;
pub use source_pump::{SourceHandle, SourcePump};
pub use task::{PipelineTask, PipelineTaskParams};
pub struct Pipeline {
processors: Vec<Box<dyn FrameProcessor>>,
}
impl Pipeline {
pub fn new(processors: Vec<Box<dyn FrameProcessor>>) -> Self {
Self { processors }
}
pub fn len(&self) -> usize {
self.processors.len()
}
pub fn is_empty(&self) -> bool {
self.processors.is_empty()
}
}
#[async_trait]
impl FrameProcessor for Pipeline {
fn name(&self) -> &str {
"Pipeline"
}
fn explode(&mut self) -> Option<Topology> {
let procs = std::mem::take(&mut self.processors);
Some(Topology::Chain(
procs.into_iter().map(into_topology).collect(),
))
}
}
pub(crate) struct Linked {
pub head: ProcessorTx,
#[allow(dead_code)]
pub tail: ProcessorTx,
}
fn make_link(
name: Arc<str>,
next: Option<ProcessorTx>,
prev: Option<ProcessorTx>,
setup: &ProcessorSetup,
) -> Link {
Link {
next,
prev,
name,
clock: setup.clock.clone(),
observer: setup.observer.clone(),
enable_metrics: setup.enable_metrics,
enable_usage_metrics: setup.enable_usage_metrics,
ttfb_start: Arc::new(std::sync::atomic::AtomicI64::new(0)),
processing_start: Arc::new(std::sync::atomic::AtomicI64::new(0)),
}
}
pub(crate) fn link_topology(
topo: Topology,
prev: Option<ProcessorTx>,
next: Option<ProcessorTx>,
setup: &ProcessorSetup,
tasks: &mut JoinSet<()>,
) -> Linked {
match topo {
Topology::Leaf(p) => {
let name: Arc<str> = Arc::from(p.name());
let (tx, rx) = channel(name.clone(), NORMAL_CHAN_CAP);
let link = make_link(name, next, prev, setup);
let setup2 = setup.clone();
tasks.spawn(async move {
run_processor(p, rx, link, setup2).await;
});
Linked {
head: tx.clone(),
tail: tx,
}
}
Topology::Chain(elems) => link_chain(elems, prev, next, setup, tasks),
Topology::Parallel(branches) => parallel::link_parallel(branches, prev, next, setup, tasks),
}
}
fn link_chain(
elems: Vec<Topology>,
prev: Option<ProcessorTx>,
next: Option<ProcessorTx>,
setup: &ProcessorSetup,
tasks: &mut JoinSet<()>,
) -> Linked {
if elems.is_empty() {
return link_topology(
Topology::Leaf(Box::new(PassThrough)),
prev,
next,
setup,
tasks,
);
}
let n = elems.len();
let mut inputs: Vec<ProcessorTx> = Vec::with_capacity(n);
let mut pending: Vec<PendingElem> = Vec::with_capacity(n);
for elem in elems {
match elem {
Topology::Leaf(p) => {
let name: Arc<str> = Arc::from(p.name());
let (tx, rx) = channel(name.clone(), NORMAL_CHAN_CAP);
inputs.push(tx);
pending.push(PendingElem::Leaf { name, p, rx });
}
other => {
let name: Arc<str> = Arc::from("nested");
let (tx, rx) = channel(name, NORMAL_CHAN_CAP);
inputs.push(tx);
pending.push(PendingElem::Composite {
topo: other,
in_rx: rx,
});
}
}
}
for (i, elem) in pending.into_iter().enumerate() {
let next_tx = if i + 1 < n {
Some(inputs[i + 1].clone())
} else {
next.clone()
};
let prev_tx = if i > 0 {
Some(inputs[i - 1].clone())
} else {
prev.clone()
};
match elem {
PendingElem::Leaf { name, p, rx } => {
let link = make_link(name, next_tx, prev_tx, setup);
let setup2 = setup.clone();
tasks.spawn(async move {
run_processor(p, rx, link, setup2).await;
});
}
PendingElem::Composite { topo, in_rx } => {
let linked = link_topology(topo, prev_tx, next_tx, setup, tasks);
let real_head = linked.head;
tasks.spawn(bridge(in_rx, real_head));
}
}
}
Linked {
head: inputs[0].clone(),
tail: inputs[n - 1].clone(),
}
}
enum PendingElem {
Leaf {
name: Arc<str>,
p: Box<dyn FrameProcessor>,
rx: crate::processor::runtime::ProcessorRx,
},
Composite {
topo: Topology,
in_rx: crate::processor::runtime::ProcessorRx,
},
}
pub(crate) async fn bridge(mut rx: crate::processor::runtime::ProcessorRx, dst: ProcessorTx) {
loop {
let env = tokio::select! {
biased;
Some(e) = rx.system.recv() => e,
Some(e) = rx.normal.recv() => e,
else => break,
};
let terminal = env.frame.is_terminal();
dst.send(env).await;
if terminal {
break;
}
}
}
struct PassThrough;
#[async_trait]
impl FrameProcessor for PassThrough {
fn name(&self) -> &str {
"PassThrough"
}
}
pub(crate) struct Source;
#[async_trait]
impl FrameProcessor for Source {
fn name(&self) -> &str {
"Source"
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
if env.direction == Direction::Upstream {
match &env.frame {
Frame::End { .. } | Frame::Stop | Frame::Cancel { .. } => {
link.push_down(env.frame).await;
return Ok(());
}
_ => {}
}
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
pub(crate) struct Sink {
pub(crate) tap: tokio::sync::mpsc::UnboundedSender<Frame>,
}
#[async_trait]
impl FrameProcessor for Sink {
fn name(&self) -> &str {
"Sink"
}
async fn start(&mut self, _setup: &ProcessorSetup, params: &StartParams) -> Result<()> {
let _ = self.tap.send(Frame::Start(params.clone()));
Ok(())
}
async fn stop(&mut self, reason: StopReason) -> Result<()> {
let frame = match reason {
StopReason::EndOfTask => Frame::End { reason: None },
StopReason::Stopped => Frame::Stop,
StopReason::Cancelled => Frame::Cancel { reason: None },
};
let _ = self.tap.send(frame);
Ok(())
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
if env.direction == Direction::Downstream {
let _ = self.tap.send(env.frame.clone());
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::processor::frame::{Frame, StartParams};
use crate::processor::Clock;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc as StdArc;
use tokio::sync::mpsc;
struct Recorder {
name: &'static str,
log: StdArc<std::sync::Mutex<Vec<String>>>,
}
#[async_trait]
impl FrameProcessor for Recorder {
fn name(&self) -> &str {
self.name
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
if let Frame::Text(t) = &env.frame {
self.log
.lock()
.unwrap()
.push(format!("{}:{}", self.name, t));
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
fn setup() -> ProcessorSetup {
ProcessorSetup {
clock: Clock::new(),
observer: None,
cancel: tokio_util::sync::CancellationToken::new(),
enable_metrics: false,
enable_usage_metrics: false,
}
}
#[tokio::test]
async fn pipeline_is_a_processor_and_nests_reaching_the_sink() {
let log = StdArc::new(std::sync::Mutex::new(Vec::new()));
let inner = Pipeline::new(vec![Box::new(Recorder {
name: "inner",
log: log.clone(),
})]);
let outer = Pipeline::new(vec![
Box::new(Recorder {
name: "a",
log: log.clone(),
}),
Box::new(inner),
Box::new(Recorder {
name: "b",
log: log.clone(),
}),
]);
let (sink_tap_tx, mut sink_tap_rx) = mpsc::unbounded_channel::<Frame>();
let topo = Topology::Chain(vec![
Topology::Leaf(Box::new(Source)),
into_topology(Box::new(outer)),
Topology::Leaf(Box::new(Sink { tap: sink_tap_tx })),
]);
let st = setup();
let mut tasks = JoinSet::new();
let linked = link_topology(topo, None, None, &st, &mut tasks);
linked
.head
.send(Envelope::new(
Frame::Text("hi".into()),
Direction::Downstream,
))
.await;
let got = tokio::time::timeout(std::time::Duration::from_secs(2), sink_tap_rx.recv())
.await
.expect("sink frame timed out")
.expect("sink closed");
assert!(matches!(got, Frame::Text(t) if t == "hi"));
linked
.head
.send(Envelope::new(
Frame::End { reason: None },
Direction::Downstream,
))
.await;
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), async {
while tasks.join_next().await.is_some() {}
})
.await;
let order = log.lock().unwrap().clone();
assert_eq!(order, vec!["a:hi", "inner:hi", "b:hi"]);
}
struct StartCounter {
seen_data_before_start: StdArc<AtomicUsize>,
started: StdArc<AtomicUsize>,
}
#[async_trait]
impl FrameProcessor for StartCounter {
fn name(&self) -> &str {
"StartCounter"
}
async fn start(&mut self, _s: &ProcessorSetup, _p: &StartParams) -> Result<()> {
self.started.fetch_add(1, Ordering::Relaxed);
Ok(())
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
if matches!(env.frame, Frame::Text(_)) && self.started.load(Ordering::Relaxed) == 0 {
self.seen_data_before_start.fetch_add(1, Ordering::Relaxed);
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
#[tokio::test]
async fn start_runs_before_data_in_a_linked_chain() {
let started = StdArc::new(AtomicUsize::new(0));
let bad = StdArc::new(AtomicUsize::new(0));
let (sink_tap_tx, _rx) = mpsc::unbounded_channel::<Frame>();
let topo = Topology::Chain(vec![
Topology::Leaf(Box::new(Source)),
Topology::Leaf(Box::new(StartCounter {
seen_data_before_start: bad.clone(),
started: started.clone(),
})),
Topology::Leaf(Box::new(Sink { tap: sink_tap_tx })),
]);
let st = setup();
let mut tasks = JoinSet::new();
let linked = link_topology(topo, None, None, &st, &mut tasks);
linked
.head
.send(Envelope::new(
Frame::Start(StartParams::default()),
Direction::Downstream,
))
.await;
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
linked
.head
.send(Envelope::new(
Frame::Text("data".into()),
Direction::Downstream,
))
.await;
linked
.head
.send(Envelope::new(
Frame::End { reason: None },
Direction::Downstream,
))
.await;
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), async {
while tasks.join_next().await.is_some() {}
})
.await;
assert_eq!(started.load(Ordering::Relaxed), 1);
assert_eq!(
bad.load(Ordering::Relaxed),
0,
"data must not precede start"
);
}
}