use tokio::sync::mpsc;
use super::frame::{Direction, Frame, FrameClass};
use super::{Envelope, FrameProcessor, Link, ProcessorSetup, StopReason};
pub const NORMAL_CHAN_CAP: usize = 64;
pub struct ProcessorRx {
pub(crate) system: mpsc::UnboundedReceiver<Envelope>,
pub(crate) normal: mpsc::Receiver<Envelope>,
}
#[derive(Clone)]
pub struct ProcessorTx {
system: mpsc::UnboundedSender<Envelope>,
normal: mpsc::Sender<Envelope>,
name: std::sync::Arc<str>,
}
impl ProcessorTx {
pub fn name(&self) -> &str {
&self.name
}
pub async fn send(&self, env: Envelope) {
match env.frame.class() {
FrameClass::System => {
let _ = self.system.send(env);
}
FrameClass::Data | FrameClass::Control => {
let _ = self.normal.send(env).await;
}
}
}
}
pub type EnvelopeSender = ProcessorTx;
pub fn channel(name: std::sync::Arc<str>, normal_cap: usize) -> (ProcessorTx, ProcessorRx) {
let (sys_tx, sys_rx) = mpsc::unbounded_channel();
let (norm_tx, norm_rx) = mpsc::channel(normal_cap);
(
ProcessorTx {
system: sys_tx,
normal: norm_tx,
name,
},
ProcessorRx {
system: sys_rx,
normal: norm_rx,
},
)
}
fn stop_reason(frame: &Frame) -> StopReason {
match frame {
Frame::End { .. } => StopReason::EndOfTask,
Frame::Stop => StopReason::Stopped,
Frame::Cancel { .. } => StopReason::Cancelled,
_ => StopReason::Stopped,
}
}
pub async fn run_processor(
mut p: Box<dyn FrameProcessor>,
mut rx: ProcessorRx,
link: Link,
setup: ProcessorSetup,
) {
loop {
let env = tokio::select! {
biased;
Some(e) = rx.system.recv() => e,
Some(e) = rx.normal.recv() => e,
else => break,
};
if let Some(o) = &setup.observer {
let ev = crate::observer::FrameEvent {
processor: &link.name,
frame: &env.frame,
meta: &env.meta,
direction: env.direction,
timestamp_ns: setup.clock.now_ns(),
};
o.on_process(&ev).await;
}
match &env.frame {
Frame::Start(params) => {
let params = params.clone();
if let Err(e) = p.start(&setup, ¶ms).await {
link.push_error(e.to_string(), false).await;
}
link.push(env.meta, env.frame, env.direction).await;
}
Frame::Interruption => {
let kept_terminal = drain_on_interruption(&mut rx, &link).await;
link.push(env.meta, env.frame, env.direction).await;
if let Some(term) = kept_terminal {
let reason = stop_reason(&term.frame);
let is_terminal = !matches!(&term.frame, Frame::Stop);
if let Err(e) = p.stop(reason).await {
link.push_error(e.to_string(), false).await;
}
link.push(term.meta, term.frame, term.direction).await;
if is_terminal {
break;
}
}
}
Frame::Cancel { .. } | Frame::End { .. } | Frame::Stop
if env.direction == Direction::Downstream =>
{
let reason = stop_reason(&env.frame);
let terminal = !matches!(&env.frame, Frame::Stop);
if let Err(e) = p.stop(reason).await {
link.push_error(e.to_string(), false).await;
}
link.push(env.meta, env.frame, env.direction).await;
if terminal {
break;
}
}
_ => {
if let Err(e) = p.process_frame(env, &link).await {
link.push_error(e.to_string(), false).await;
}
}
}
}
}
async fn drain_on_interruption(rx: &mut ProcessorRx, link: &Link) -> Option<Envelope> {
while let Ok(env) = rx.normal.try_recv() {
if env.frame.uninterruptible() {
if env.direction == Direction::Downstream && env.frame.is_terminal() {
return Some(env);
}
link.push(env.meta, env.frame, env.direction).await;
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::Result;
use crate::processor::{Clock, FrameProcessor};
use async_trait::async_trait;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
struct TapSink {
stopped: Arc<AtomicBool>,
}
#[async_trait]
impl FrameProcessor for TapSink {
fn name(&self) -> &str {
"TapSink"
}
async fn stop(&mut self, _reason: StopReason) -> Result<()> {
self.stopped.store(true, Ordering::Relaxed);
Ok(())
}
}
fn sink_link(name: Arc<str>, clock: Clock) -> Link {
Link {
next: None,
prev: None,
name,
clock,
observer: None,
enable_metrics: false,
enable_usage_metrics: false,
ttfb_start: Arc::new(std::sync::atomic::AtomicI64::new(0)),
processing_start: Arc::new(std::sync::atomic::AtomicI64::new(0)),
}
}
#[tokio::test]
async fn interruption_does_not_lose_a_buffered_downstream_end_at_a_sink() {
let name: Arc<str> = Arc::from("TapSink");
let (tx, rx) = channel(name.clone(), 8);
let clock = Clock::new();
let setup = ProcessorSetup {
clock: clock.clone(),
observer: None,
cancel: tokio_util::sync::CancellationToken::new(),
enable_metrics: false,
enable_usage_metrics: false,
};
let stopped = Arc::new(AtomicBool::new(false));
let p = Box::new(TapSink {
stopped: stopped.clone(),
});
tx.send(Envelope::new(
Frame::End { reason: None },
Direction::Downstream,
))
.await;
tx.send(Envelope::new(Frame::Interruption, Direction::Downstream))
.await;
drop(tx);
let h = tokio::spawn(run_processor(p, rx, sink_link(name, clock), setup));
let res = tokio::time::timeout(std::time::Duration::from_secs(2), h).await;
assert!(
res.is_ok(),
"run_processor hung on the interruption-vs-buffered-End race"
);
assert!(
stopped.load(Ordering::Relaxed),
"the Sink's stop() hook must run for the buffered End (else PipelineTask hangs)"
);
}
}