use crate::cancel::CancellationToken;
use crate::error::CanoError;
use crate::resource::Resources;
use crate::task::{TaskConfig, TaskResult};
use futures_util::Stream;
use serde::Serialize;
use serde::de::DeserializeOwned;
use std::borrow::Cow;
use std::fmt;
use std::future::Future;
use std::hash::Hash;
use std::pin::Pin;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum StreamErrorPolicy {
#[default]
FailFast,
SkipAndContinue,
RetryOnError {
max_errors: u32,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamWindow {
Count(usize),
Duration(std::time::Duration),
}
pub type StreamBatch = StreamWindow;
#[derive(Debug)]
pub enum WindowSignal<TState> {
Continue,
Stop(TaskResult<TState>),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CloseReason {
Exhausted,
Cancelled,
}
#[crate::task::stream]
pub trait StreamTask<TState, TResourceKey = Cow<'static, str>>: Send + Sync
where
TState: Clone + fmt::Debug + Send + Sync + 'static,
TResourceKey: Hash + Eq + Send + Sync + 'static,
{
type Item: Send + 'static;
type Output: Send + 'static;
type Cursor: Serialize + DeserializeOwned + Send + Sync + 'static;
fn window(&self) -> StreamWindow {
StreamWindow::Count(1)
}
fn on_item_error(&self) -> StreamErrorPolicy {
StreamErrorPolicy::FailFast
}
fn config(&self) -> TaskConfig {
crate::task::minimal_task_config()
}
fn name(&self) -> Cow<'static, str> {
crate::task::default_task_name::<Self>()
}
async fn open(
&self,
res: &Resources<TResourceKey>,
cursor: Option<Self::Cursor>,
) -> Result<Pin<Box<dyn Stream<Item = Self::Item> + Send>>, CanoError>;
async fn process_item(
&self,
res: &Resources<TResourceKey>,
item: Self::Item,
) -> Result<(Self::Output, Self::Cursor), CanoError>;
async fn flush_window(
&self,
res: &Resources<TResourceKey>,
outputs: Vec<Self::Output>,
) -> Result<WindowSignal<TState>, CanoError>;
async fn on_close(
&self,
res: &Resources<TResourceKey>,
reason: CloseReason,
) -> Result<TaskResult<TState>, CanoError>;
#[doc(hidden)]
fn run_in_memory<'life0, 'life1, 'async_trait>(
&'life0 self,
res: &'life1 Resources<TResourceKey>,
) -> Pin<Box<dyn Future<Output = Result<TaskResult<TState>, CanoError>> + Send + 'async_trait>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: Sync + 'async_trait + Sized,
{
Box::pin(run_stream_in_memory(self, res))
}
}
pub(crate) enum WindowStep<TCursor, TState> {
Window { cursor: TCursor },
Done {
final_cursor: Option<TCursor>,
result: TaskResult<TState>,
},
Cancelled { final_cursor: Option<TCursor> },
}
const MIN_DURATION_WINDOW: std::time::Duration = std::time::Duration::from_millis(1);
#[allow(clippy::too_many_arguments)]
async fn drive_window<T, S, K>(
task: &T,
res: &Resources<K>,
stream: &mut Pin<Box<dyn Stream<Item = T::Item> + Send>>,
consecutive_errors: &mut u32,
window: &StreamWindow,
policy: &StreamErrorPolicy,
attempt_timeout: Option<std::time::Duration>,
token: &CancellationToken,
) -> Result<WindowStep<T::Cursor, S>, CanoError>
where
T: StreamTask<S, K> + ?Sized,
S: Clone + fmt::Debug + Send + Sync + 'static,
K: Hash + Eq + Send + Sync + 'static,
{
use futures_util::StreamExt as _;
let (count_limit, duration_len) = match window {
StreamWindow::Count(n) => (Some((*n).max(1)), None),
StreamWindow::Duration(d) => (None, Some((*d).max(MIN_DURATION_WINDOW))),
};
let mut buf: Vec<T::Output> = Vec::new();
let mut last_cursor: Option<T::Cursor> = None;
let mut tick: Option<Pin<Box<tokio::time::Sleep>>> = duration_len
.map(|len| Box::pin(tokio::time::sleep_until(tokio::time::Instant::now() + len)));
loop {
if let Some(limit) = count_limit
&& buf.len() >= limit
{
return flush_full_window(task, res, std::mem::take(&mut buf), last_cursor).await;
}
tokio::select! {
biased;
_ = token.cancelled() => {
if !buf.is_empty() {
let drain_flush = flush_partial_window(task, res, std::mem::take(&mut buf)).await;
if let Err(e) = drain_flush {
let _ = task.on_close(res, CloseReason::Cancelled).await;
return Err(e);
}
}
let _ = task.on_close(res, CloseReason::Cancelled).await?;
return Ok(WindowStep::Cancelled { final_cursor: last_cursor });
}
_ = async {
match tick.as_mut() {
Some(sleep) => sleep.as_mut().await,
None => std::future::pending::<()>().await,
}
} => {
if buf.is_empty() {
tick = duration_len
.map(|len| Box::pin(tokio::time::sleep_until(tokio::time::Instant::now() + len)));
continue;
}
return flush_full_window(task, res, std::mem::take(&mut buf), last_cursor).await;
}
item = stream.next() => {
match item {
Some(item) => {
let processed = match attempt_timeout {
Some(d) => match tokio::time::timeout(d, task.process_item(res, item)).await {
Ok(inner) => inner,
Err(_elapsed) => Err(CanoError::timeout(
"stream process_item exceeded attempt_timeout",
)),
},
None => task.process_item(res, item).await,
};
match processed {
Ok((out, cursor)) => {
*consecutive_errors = 0;
buf.push(out);
last_cursor = Some(cursor);
#[cfg(feature = "metrics")]
crate::metrics::stream_items(1, 0);
}
Err(e) => {
#[cfg(feature = "metrics")]
crate::metrics::stream_items(0, 1);
match policy {
StreamErrorPolicy::FailFast => return Err(e),
StreamErrorPolicy::SkipAndContinue => {}
StreamErrorPolicy::RetryOnError { max_errors } => {
*consecutive_errors += 1;
if *consecutive_errors > *max_errors {
return Err(e);
}
}
}
}
}
}
None => {
if !buf.is_empty() {
match flush_partial_window(task, res, std::mem::take(&mut buf)).await? {
WindowSignal::Stop(result) => {
return Ok(WindowStep::Done {
final_cursor: last_cursor,
result,
});
}
WindowSignal::Continue => {}
}
}
let result = task.on_close(res, CloseReason::Exhausted).await?;
return Ok(WindowStep::Done { final_cursor: last_cursor, result });
}
}
}
}
}
}
async fn flush_partial_window<T, S, K>(
task: &T,
res: &Resources<K>,
outputs: Vec<T::Output>,
) -> Result<WindowSignal<S>, CanoError>
where
T: StreamTask<S, K> + ?Sized,
S: Clone + fmt::Debug + Send + Sync + 'static,
K: Hash + Eq + Send + Sync + 'static,
{
#[cfg(feature = "metrics")]
crate::metrics::stream_window();
task.flush_window(res, outputs).await
}
async fn flush_full_window<T, S, K>(
task: &T,
res: &Resources<K>,
outputs: Vec<T::Output>,
last_cursor: Option<T::Cursor>,
) -> Result<WindowStep<T::Cursor, S>, CanoError>
where
T: StreamTask<S, K> + ?Sized,
S: Clone + fmt::Debug + Send + Sync + 'static,
K: Hash + Eq + Send + Sync + 'static,
{
#[cfg(feature = "metrics")]
crate::metrics::stream_window();
Ok(match task.flush_window(res, outputs).await? {
WindowSignal::Continue => WindowStep::Window {
cursor: last_cursor.expect("a non-empty window always has a cursor"),
},
WindowSignal::Stop(result) => WindowStep::Done {
final_cursor: last_cursor,
result,
},
})
}
async fn run_stream_in_memory<T, S, K>(
task: &T,
res: &Resources<K>,
) -> Result<TaskResult<S>, CanoError>
where
T: StreamTask<S, K> + ?Sized,
S: Clone + fmt::Debug + Send + Sync + 'static,
K: Hash + Eq + Send + Sync + 'static,
{
let token = CancellationToken::disabled();
let window = task.window();
let policy = task.on_item_error();
let attempt_timeout = task.config().attempt_timeout;
let mut consecutive_errors: u32 = 0;
let result: Result<TaskResult<S>, CanoError> = async {
let mut stream = task.open(res, None).await?;
loop {
match drive_window(
task,
res,
&mut stream,
&mut consecutive_errors,
&window,
&policy,
attempt_timeout,
&token,
)
.await?
{
WindowStep::Window { .. } => continue,
WindowStep::Done { result, .. } => return Ok(result),
WindowStep::Cancelled { .. } => return Err(CanoError::cancelled()),
}
}
}
.await;
#[cfg(feature = "metrics")]
crate::metrics::stream_run(if result.is_ok() {
"completed"
} else {
"failed"
});
result
}
pub enum ErasedWindowStep<TState> {
Window { cursor: Vec<u8> },
Done {
final_cursor: Option<Vec<u8>>,
result: TaskResult<TState>,
},
Cancelled { final_cursor: Option<Vec<u8>> },
}
pub type WindowFuture<'a, TState> =
Pin<Box<dyn Future<Output = Result<ErasedWindowStep<TState>, CanoError>> + Send + 'a>>;
pub trait ErasedStreamSession<TState, TResourceKey>: Send
where
TState: Clone + Send + Sync + 'static,
TResourceKey: Hash + Eq + Send + Sync + 'static,
{
fn next_window<'a>(
&'a mut self,
res: &'a Resources<TResourceKey>,
token: &'a CancellationToken,
) -> WindowFuture<'a, TState>;
}
pub type OpenSessionFuture<'a, TState, TResourceKey> = Pin<
Box<
dyn Future<Output = Result<Box<dyn ErasedStreamSession<TState, TResourceKey>>, CanoError>>
+ Send
+ 'a,
>,
>;
pub trait ErasedStreamTask<TState, TResourceKey>: Send + Sync
where
TState: Clone + Send + Sync + 'static,
TResourceKey: Hash + Eq + Send + Sync + 'static,
{
fn name(&self) -> Cow<'static, str>;
fn open_session<'a>(
&'a self,
res: &'a Resources<TResourceKey>,
cursor_bytes: Option<Vec<u8>>,
attempt_timeout: Option<std::time::Duration>,
) -> OpenSessionFuture<'a, TState, TResourceKey>;
}
struct StreamSession<T, S, K>
where
T: StreamTask<S, K> + 'static,
S: Clone + fmt::Debug + Send + Sync + 'static,
K: Hash + Eq + Send + Sync + 'static,
{
task: Arc<T>,
stream: Pin<Box<dyn Stream<Item = T::Item> + Send>>,
window: StreamWindow,
policy: StreamErrorPolicy,
attempt_timeout: Option<std::time::Duration>,
consecutive_errors: u32,
}
impl<T, S, K> ErasedStreamSession<S, K> for StreamSession<T, S, K>
where
T: StreamTask<S, K> + 'static,
S: Clone + fmt::Debug + Send + Sync + 'static,
K: Hash + Eq + Send + Sync + 'static,
{
fn next_window<'a>(
&'a mut self,
res: &'a Resources<K>,
token: &'a CancellationToken,
) -> WindowFuture<'a, S> {
Box::pin(async move {
let step = drive_window(
&*self.task,
res,
&mut self.stream,
&mut self.consecutive_errors,
&self.window,
&self.policy,
self.attempt_timeout,
token,
)
.await?;
Ok(match step {
WindowStep::Window { cursor } => ErasedWindowStep::Window {
cursor: encode_cursor(&cursor, &self.task.name())?,
},
WindowStep::Done {
final_cursor,
result,
} => ErasedWindowStep::Done {
final_cursor: final_cursor
.map(|c| encode_cursor(&c, &self.task.name()))
.transpose()?,
result,
},
WindowStep::Cancelled { final_cursor } => ErasedWindowStep::Cancelled {
final_cursor: final_cursor
.map(|c| encode_cursor(&c, &self.task.name()))
.transpose()?,
},
})
})
}
}
pub(crate) struct StreamAdapter<T>(pub Arc<T>);
impl<TState, TResourceKey, T> ErasedStreamTask<TState, TResourceKey> for StreamAdapter<T>
where
TState: Clone + fmt::Debug + Send + Sync + 'static,
TResourceKey: Hash + Eq + Send + Sync + 'static,
T: StreamTask<TState, TResourceKey> + 'static,
{
fn name(&self) -> Cow<'static, str> {
self.0.name()
}
fn open_session<'a>(
&'a self,
res: &'a Resources<TResourceKey>,
cursor_bytes: Option<Vec<u8>>,
attempt_timeout: Option<std::time::Duration>,
) -> OpenSessionFuture<'a, TState, TResourceKey> {
Box::pin(async move {
let cursor: Option<T::Cursor> = match cursor_bytes {
None => None,
Some(ref b) => Some(serde_json::from_slice(b).map_err(|e| {
CanoError::task_execution(format!(
"deserialize stream cursor for `{}`: {e}",
self.0.name()
))
})?),
};
let stream = self.0.open(res, cursor).await?;
let session = StreamSession {
task: Arc::clone(&self.0),
stream,
window: self.0.window(),
policy: self.0.on_item_error(),
attempt_timeout,
consecutive_errors: 0,
};
Ok(Box::new(session) as Box<dyn ErasedStreamSession<TState, TResourceKey>>)
})
}
}
fn encode_cursor<C: Serialize>(cursor: &C, task_name: &str) -> Result<Vec<u8>, CanoError> {
serde_json::to_vec(cursor).map_err(|e| {
CanoError::task_execution(format!("serialize stream cursor for `{task_name}`: {e}"))
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::task;
use crate::task::Task;
use futures_util::stream;
use std::sync::Mutex;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum Step {
Consume,
Done,
}
#[test]
fn value_type_defaults() {
assert_eq!(StreamErrorPolicy::default(), StreamErrorPolicy::FailFast);
let _ = StreamWindow::Count(8);
let _ = StreamWindow::Duration(std::time::Duration::from_millis(5));
assert_eq!(CloseReason::Exhausted, CloseReason::Exhausted);
}
#[derive(Default)]
struct Collector {
seen: Mutex<Vec<u32>>,
windows: Mutex<Vec<Vec<u32>>>,
}
#[task::stream]
impl StreamTask<Step> for Collector {
type Item = u32;
type Output = u32;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(2)
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![10u32, 20, 30, 40, 50]))
as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
self.seen.lock().unwrap().push(item);
Ok((item * 2, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u32>,
) -> Result<WindowSignal<Step>, CanoError> {
self.windows.lock().unwrap().push(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn in_memory_windows_and_order() {
let task = Collector::default();
let res = Resources::new();
let result = Task::run(&task, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
assert_eq!(*task.seen.lock().unwrap(), vec![10, 20, 30, 40, 50]);
assert_eq!(
*task.windows.lock().unwrap(),
vec![vec![20u32, 40], vec![60, 80], vec![100]]
);
}
struct FailOnSecond {
policy: StreamErrorPolicy,
flushed: Mutex<Vec<u32>>,
}
#[task::stream]
impl StreamTask<Step> for FailOnSecond {
type Item = u32;
type Output = u32;
type Cursor = u64;
fn on_item_error(&self) -> StreamErrorPolicy {
self.policy.clone()
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32, 2, 3])) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
if item == 2 {
Err(CanoError::task_execution("item 2 failed"))
} else {
Ok((item, item as u64))
}
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u32>,
) -> Result<WindowSignal<Step>, CanoError> {
self.flushed.lock().unwrap().extend(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn fail_fast_propagates() {
let task = FailOnSecond {
policy: StreamErrorPolicy::FailFast,
flushed: Mutex::new(Vec::new()),
};
let res = Resources::new();
let err = Task::run(&task, &res).await.unwrap_err();
assert!(matches!(err, CanoError::TaskExecution(_)));
}
#[tokio::test]
async fn skip_and_continue_drops_bad_item() {
let task = FailOnSecond {
policy: StreamErrorPolicy::SkipAndContinue,
flushed: Mutex::new(Vec::new()),
};
let res = Resources::new();
let result = Task::run(&task, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
assert_eq!(*task.flushed.lock().unwrap(), vec![1u32, 3]);
}
struct StopAfterFirst;
#[task::stream]
impl StreamTask<Step> for StopAfterFirst {
type Item = u32;
type Output = u32;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32, 2, 3, 4]))
as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Stop(TaskResult::Single(Step::Done)))
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
panic!("on_close must not run when a window returns Stop");
}
}
#[tokio::test]
async fn window_stop_short_circuits() {
let res = Resources::new();
let result = Task::run(&StopAfterFirst, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
}
#[tokio::test]
async fn integrates_with_workflow_via_register() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let workflow = Workflow::bare()
.register(Step::Consume, Collector::default())
.add_exit_state(Step::Done);
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
}
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
#[derive(Default)]
struct InMemoryStore {
rows: Mutex<HashMap<String, Vec<crate::recovery::CheckpointRow>>>,
committed: Mutex<Vec<Vec<u8>>>,
}
#[crate::checkpoint_store]
impl crate::recovery::CheckpointStore for InMemoryStore {
async fn append(
&self,
workflow_id: &str,
row: crate::recovery::CheckpointRow,
) -> Result<(), CanoError> {
if row.kind == crate::recovery::RowKind::StepCursor
&& let Some(blob) = &row.output_blob
{
self.committed.lock().unwrap().push(blob.clone());
}
let mut g = self.rows.lock().unwrap();
let v = g.entry(workflow_id.to_string()).or_default();
if v.iter().any(|r| r.sequence == row.sequence) {
return Err(CanoError::checkpoint_store("duplicate sequence"));
}
v.push(row);
Ok(())
}
async fn load_run(
&self,
workflow_id: &str,
) -> Result<Vec<crate::recovery::CheckpointRow>, CanoError> {
let g = self.rows.lock().unwrap();
let mut v = g.get(workflow_id).cloned().unwrap_or_default();
v.sort_by_key(|r| r.sequence);
Ok(v)
}
async fn clear(&self, workflow_id: &str) -> Result<(), CanoError> {
self.rows.lock().unwrap().remove(workflow_id);
Ok(())
}
}
struct Forever {
closed_cancelled: Arc<AtomicBool>,
flushed_windows: Arc<AtomicU32>,
}
#[task::stream]
impl StreamTask<Step> for Forever {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(2)
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(0u64..)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.flushed_windows.fetch_add(1, Ordering::SeqCst);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
if reason == CloseReason::Cancelled {
self.closed_cancelled.store(true, Ordering::SeqCst);
}
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn cancel_drains_and_surfaces_cancelled() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let closed = Arc::new(AtomicBool::new(false));
let task = Forever {
closed_cancelled: Arc::clone(&closed),
flushed_windows: Arc::new(AtomicU32::new(0)),
};
let (handle, token) = CancellationToken::new();
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
handle.cancel();
});
let result = workflow.orchestrate(Step::Consume, token).await;
assert!(
matches!(&result, Err(e) if e.category() == "cancelled"),
"a cancelled stream must surface as cancelled, got {result:?}"
);
assert!(
closed.load(Ordering::SeqCst),
"on_close(Cancelled) must run (cooperative drain reached the close hook)"
);
}
struct Resumable {
opened: Arc<Mutex<Vec<Option<u64>>>>,
processed: Arc<Mutex<Vec<u64>>>,
fail_third: Arc<AtomicBool>,
}
#[task::stream]
impl StreamTask<Step> for Resumable {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(2)
}
async fn open(
&self,
_res: &Resources,
cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
self.opened.lock().unwrap().push(cursor);
let start = cursor.map(|c| c + 1).unwrap_or(1);
let items: Vec<u64> = (start..=6).collect();
Ok(Box::pin(stream::iter(items)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
self.processed.lock().unwrap().push(item);
Ok((item, item)) }
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
if outputs == vec![5u64, 6] && self.fail_third.swap(false, Ordering::SeqCst) {
return Err(CanoError::task_execution("simulated crash in window [5,6]"));
}
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn persists_cursor_and_resumes() {
use crate::cancel::CancellationToken;
use crate::recovery::{CheckpointStore, RowKind};
use crate::workflow::Workflow;
let opened = Arc::new(Mutex::new(Vec::new()));
let processed = Arc::new(Mutex::new(Vec::new()));
let task = Resumable {
opened: Arc::clone(&opened),
processed: Arc::clone(&processed),
fail_third: Arc::new(AtomicBool::new(true)),
};
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("resume-test");
let r1 = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await;
assert!(r1.is_err(), "run 1 should fail mid-stream: {r1:?}");
let rows = store.load_run("resume-test").await.unwrap();
let cursors: Vec<u64> = rows
.iter()
.filter(|r| r.kind == RowKind::StepCursor)
.map(|r| serde_json::from_slice::<u64>(r.output_blob.as_ref().unwrap()).unwrap())
.collect();
assert_eq!(
cursors,
vec![2, 4],
"only fully-flushed windows commit a cursor"
);
let r2 = workflow
.resume_from("resume-test", CancellationToken::disabled())
.await
.unwrap();
assert_eq!(r2, Step::Done);
assert_eq!(*opened.lock().unwrap(), vec![None, Some(4)]);
assert_eq!(
*processed.lock().unwrap(),
vec![1u64, 2, 3, 4, 5, 6, 5, 6],
"resume reprocesses only the items after the committed cursor"
);
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum S3 {
Consume,
ViaStop,
ViaClose,
}
struct StopOnFinalWindow;
#[task::stream]
impl StreamTask<S3> for StopOnFinalWindow {
type Item = u32;
type Output = u32;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(3) }
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32, 2])) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<S3>, CanoError> {
Ok(WindowSignal::Stop(TaskResult::Single(S3::ViaStop)))
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<S3>, CanoError> {
Ok(TaskResult::Single(S3::ViaClose))
}
}
#[tokio::test]
async fn terminal_flush_stop_is_honored() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let workflow = Workflow::bare()
.register_stream(S3::Consume, StopOnFinalWindow)
.add_exit_states([S3::ViaStop, S3::ViaClose]);
let result = workflow
.orchestrate(S3::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(
result,
S3::ViaStop,
"a Stop from the final partial flush must win over on_close(Exhausted)"
);
}
#[derive(Default)]
struct CancelCounter {
cancels: AtomicU32,
}
impl crate::observer::WorkflowObserver for CancelCounter {
fn on_cancelled(&self, _state: &str) {
self.cancels.fetch_add(1, Ordering::SeqCst);
}
}
#[tokio::test]
async fn cancel_fires_on_cancelled_once() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let task = Forever {
closed_cancelled: Arc::new(AtomicBool::new(false)),
flushed_windows: Arc::new(AtomicU32::new(0)),
};
let counter = Arc::new(CancelCounter::default());
let (handle, token) = CancellationToken::new();
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_observer(counter.clone());
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
handle.cancel();
});
let result = workflow.orchestrate(Step::Consume, token).await;
assert!(matches!(&result, Err(e) if e.category() == "cancelled"));
assert_eq!(
counter.cancels.load(Ordering::SeqCst),
1,
"on_cancelled must fire exactly once on a stream cancel"
);
}
struct ZeroWindow {
seen: Arc<parking_lot::Mutex<Vec<u32>>>,
}
#[task::stream]
impl StreamTask<Step> for ZeroWindow {
type Item = u32;
type Output = u32;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Duration(std::time::Duration::ZERO)
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32, 2, 3])) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
self.seen.lock().push(item);
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn duration_window_zero_is_clamped_not_livelocked() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let seen = Arc::new(parking_lot::Mutex::new(Vec::new()));
let workflow = Workflow::bare()
.register_stream(Step::Consume, ZeroWindow { seen: seen.clone() })
.add_exit_state(Step::Done);
let result = tokio::time::timeout(
std::time::Duration::from_secs(5),
workflow.orchestrate(Step::Consume, CancellationToken::disabled()),
)
.await
.expect("must not livelock on a zero-duration window");
assert!(matches!(result, Ok(Step::Done)), "got {result:?}");
assert_eq!(&*seen.lock(), &[1, 2, 3]);
}
struct FlushFailsOnDrain {
on_close_reason: Arc<parking_lot::Mutex<Option<CloseReason>>>,
}
#[task::stream]
impl StreamTask<Step> for FlushFailsOnDrain {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(1000)
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::unfold(0u64, |n| async move {
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
Some((n, n + 1))
})) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Err(CanoError::task_execution("flush failed during drain"))
}
async fn on_close(
&self,
_res: &Resources,
reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
*self.on_close_reason.lock() = Some(reason);
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn cancel_drain_flush_error_still_runs_on_close() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let on_close_reason = Arc::new(parking_lot::Mutex::new(None));
let task = FlushFailsOnDrain {
on_close_reason: on_close_reason.clone(),
};
let (handle, token) = CancellationToken::new();
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
handle.cancel();
});
let result = workflow.orchestrate(Step::Consume, token).await;
assert!(
matches!(&result, Err(e) if e.category() == "task_execution"),
"the drain flush error must be surfaced, not swallowed into `Cancelled`; got \
{result:?}"
);
assert_eq!(
*on_close_reason.lock(),
Some(CloseReason::Cancelled),
"on_close(Cancelled) must still run as best-effort cleanup even though the \
drain flush failed"
);
}
#[tokio::test]
async fn total_timeout_flushes_commits_and_reclassifies() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let task = Forever {
closed_cancelled: Arc::new(AtomicBool::new(false)),
flushed_windows: Arc::new(AtomicU32::new(0)),
};
let closed = Arc::clone(&task.closed_cancelled);
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("total-timeout")
.with_total_timeout(std::time::Duration::from_millis(30));
let result = tokio::time::timeout(
std::time::Duration::from_secs(5),
workflow.orchestrate(Step::Consume, CancellationToken::disabled()),
)
.await
.expect("a Stream state must not silently defeat with_total_timeout");
assert!(
matches!(&result, Err(e) if e.category() == "workflow_timeout"),
"a total-timeout budget must be honoured for a Stream state, not silently \
ignored; got {result:?}"
);
assert!(
closed.load(Ordering::SeqCst),
"on_close(Cancelled) must still run — the timeout is delivered through the \
same cooperative drain as a real cancel"
);
assert!(
!store.committed.lock().unwrap().is_empty(),
"the in-flight window's cursor must be committed before the timeout surfaces"
);
}
struct HangingOpen {
on_close_called: Arc<AtomicBool>,
}
#[task::stream]
impl StreamTask<Step> for HangingOpen {
type Item = u64;
type Output = u64;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
std::future::pending::<()>().await;
unreachable!("cancellation must interrupt a hung open() before this resolves")
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
self.on_close_called.store(true, Ordering::SeqCst);
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn cancel_during_open_session_is_cancellable_without_on_close() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let on_close_called = Arc::new(AtomicBool::new(false));
let (handle, token) = CancellationToken::new();
let workflow = Workflow::bare()
.register_stream(
Step::Consume,
HangingOpen {
on_close_called: on_close_called.clone(),
},
)
.add_exit_state(Step::Done);
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
handle.cancel();
});
let result = tokio::time::timeout(
std::time::Duration::from_secs(5),
workflow.orchestrate(Step::Consume, token),
)
.await
.expect("a hung open() must still be cancellable");
assert!(
matches!(&result, Err(e) if e.category() == "cancelled"),
"got {result:?}"
);
assert!(
!on_close_called.load(Ordering::SeqCst),
"on_close must not run — open() never produced a session to close"
);
}
struct SlowItem;
#[task::stream]
impl StreamTask<Step> for SlowItem {
type Item = u32;
type Output = u32;
type Cursor = u64;
fn config(&self) -> TaskConfig {
TaskConfig::minimal().with_attempt_timeout(std::time::Duration::from_millis(10))
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32])) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn attempt_timeout_bounds_process_item() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let workflow = Workflow::bare()
.register_stream(Step::Consume, SlowItem)
.add_exit_state(Step::Done);
let result = tokio::time::timeout(
std::time::Duration::from_secs(5),
workflow.orchestrate(Step::Consume, CancellationToken::disabled()),
)
.await
.expect("attempt_timeout must bound the hung process_item well under 5s");
assert!(
matches!(&result, Err(e) if e.category() == "timeout"),
"a process_item exceeding attempt_timeout must surface a timeout error, got {result:?}"
);
}
struct RecvStream(tokio::sync::mpsc::UnboundedReceiver<u64>);
impl Stream for RecvStream {
type Item = u64;
fn poll_next(
self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<u64>> {
self.get_mut().0.poll_recv(cx)
}
}
fn step_cursors(store: &InMemoryStore) -> Vec<u64> {
store
.committed
.lock()
.unwrap()
.iter()
.map(|blob| serde_json::from_slice::<u64>(blob).unwrap())
.collect()
}
struct DurationSource {
rx: Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<u64>>>,
windows: Arc<Mutex<Vec<Vec<u64>>>>,
}
#[task::stream]
impl StreamTask<Step> for DurationSource {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Duration(std::time::Duration::from_millis(50))
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
let rx = self.rx.lock().unwrap().take().expect("open called once");
Ok(Box::pin(RecvStream(rx)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.windows.lock().unwrap().push(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test(start_paused = true)]
async fn duration_window_flushes_on_deadline_and_rearms() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<u64>();
let windows = Arc::new(Mutex::new(Vec::new()));
let store = Arc::new(InMemoryStore::default());
let task = DurationSource {
rx: Mutex::new(Some(rx)),
windows: Arc::clone(&windows),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("dur-rearm");
tokio::spawn(async move {
for v in 1u64..=4 {
tokio::time::sleep(std::time::Duration::from_millis(30)).await;
let _ = tx.send(v);
}
});
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
assert_eq!(
*windows.lock().unwrap(),
vec![vec![1u64], vec![2, 3], vec![4]],
"tumbling duration windows re-arm after each Continue flush"
);
assert_eq!(
step_cursors(&store),
vec![1u64, 3, 4],
"each duration flush commits its last item's cursor"
);
}
#[tokio::test(start_paused = true)]
async fn duration_window_skips_empty_intervals() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<u64>();
let windows = Arc::new(Mutex::new(Vec::new()));
let task = DurationSource {
rx: Mutex::new(Some(rx)),
windows: Arc::clone(&windows),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(120)).await;
let _ = tx.send(1);
});
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
let w = windows.lock().unwrap();
assert!(
w.iter().all(|win| !win.is_empty()),
"an idle duration window must never flush an empty buffer: {w:?}"
);
assert_eq!(
*w,
vec![vec![1u64]],
"exactly one real window despite two elapsed-but-empty deadlines"
);
}
struct DurationStop {
rx: Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<u64>>>,
on_close_ran: Arc<AtomicBool>,
}
#[task::stream]
impl StreamTask<S3> for DurationStop {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Duration(std::time::Duration::from_millis(50))
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
let rx = self.rx.lock().unwrap().take().expect("open called once");
Ok(Box::pin(RecvStream(rx)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<S3>, CanoError> {
Ok(WindowSignal::Stop(TaskResult::Single(S3::ViaStop)))
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<S3>, CanoError> {
self.on_close_ran.store(true, Ordering::SeqCst);
Ok(TaskResult::Single(S3::ViaClose))
}
}
#[tokio::test(start_paused = true)]
async fn duration_window_stop_transitions_without_close() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<u64>();
let on_close_ran = Arc::new(AtomicBool::new(false));
let task = DurationStop {
rx: Mutex::new(Some(rx)),
on_close_ran: Arc::clone(&on_close_ran),
};
let workflow = Workflow::bare()
.register_stream(S3::Consume, task)
.add_exit_states([S3::ViaStop, S3::ViaClose]);
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(30)).await;
let _ = tx.send(1);
tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
drop(tx);
});
let result = workflow
.orchestrate(S3::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(
result,
S3::ViaStop,
"a duration-window Stop wins over on_close"
);
assert!(
!on_close_ran.load(Ordering::SeqCst),
"on_close must not run when a duration window returns Stop"
);
}
struct ScriptedErrors {
len: u64,
fail: Vec<u64>,
policy: StreamErrorPolicy,
flushed: Arc<Mutex<Vec<u64>>>,
}
#[task::stream]
impl StreamTask<Step> for ScriptedErrors {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn on_item_error(&self) -> StreamErrorPolicy {
self.policy.clone()
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
let items: Vec<u64> = (1..=self.len).collect();
Ok(Box::pin(stream::iter(items)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
if self.fail.contains(&item) {
Err(CanoError::task_execution(format!("item {item} failed")))
} else {
Ok((item, item))
}
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.flushed.lock().unwrap().extend(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn retry_on_error_tolerates_consecutive_within_max() {
let flushed = Arc::new(Mutex::new(Vec::new()));
let task = ScriptedErrors {
len: 4,
fail: vec![2, 3],
policy: StreamErrorPolicy::RetryOnError { max_errors: 2 },
flushed: Arc::clone(&flushed),
};
let res = Resources::new();
let result = Task::run(&task, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
assert_eq!(
*flushed.lock().unwrap(),
vec![1u64, 4],
"only ok items flush"
);
}
#[tokio::test]
async fn retry_on_error_fails_past_max() {
let task = ScriptedErrors {
len: 3,
fail: vec![2, 3],
policy: StreamErrorPolicy::RetryOnError { max_errors: 1 },
flushed: Arc::new(Mutex::new(Vec::new())),
};
let res = Resources::new();
let err = Task::run(&task, &res).await.unwrap_err();
assert_eq!(err.category(), "task_execution");
}
#[tokio::test]
async fn retry_on_error_counter_resets_on_success() {
let task = ScriptedErrors {
len: 6,
fail: vec![2, 3, 5, 6],
policy: StreamErrorPolicy::RetryOnError { max_errors: 2 },
flushed: Arc::new(Mutex::new(Vec::new())),
};
let res = Resources::new();
let result = Task::run(&task, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
}
#[tokio::test]
async fn retry_on_error_max_zero_fails_on_first() {
let task = ScriptedErrors {
len: 3,
fail: vec![2],
policy: StreamErrorPolicy::RetryOnError { max_errors: 0 },
flushed: Arc::new(Mutex::new(Vec::new())),
};
let res = Resources::new();
let err = Task::run(&task, &res).await.unwrap_err();
assert_eq!(err.category(), "task_execution");
}
struct SlowUnderPolicy {
policy: StreamErrorPolicy,
flushed: Arc<Mutex<Vec<u64>>>,
}
#[task::stream]
impl StreamTask<Step> for SlowUnderPolicy {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn config(&self) -> TaskConfig {
TaskConfig::minimal().with_attempt_timeout(std::time::Duration::from_millis(10))
}
fn on_item_error(&self) -> StreamErrorPolicy {
self.policy.clone()
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u64, 2, 3])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
if item == 1 {
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
}
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.flushed.lock().unwrap().extend(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn timeout_item_skipped_under_skip_and_continue() {
let flushed = Arc::new(Mutex::new(Vec::new()));
let task = SlowUnderPolicy {
policy: StreamErrorPolicy::SkipAndContinue,
flushed: Arc::clone(&flushed),
};
let res = Resources::new();
let result = Task::run(&task, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
assert_eq!(
*flushed.lock().unwrap(),
vec![2u64, 3],
"the timed-out item is skipped, not fatal"
);
}
#[tokio::test]
async fn timeout_item_counts_under_retry_on_error() {
let task = SlowUnderPolicy {
policy: StreamErrorPolicy::RetryOnError { max_errors: 0 },
flushed: Arc::new(Mutex::new(Vec::new())),
};
let res = Resources::new();
let err = Task::run(&task, &res).await.unwrap_err();
assert_eq!(err.category(), "timeout");
}
#[tokio::test]
async fn skip_does_not_commit_bad_item_cursor() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let task = ScriptedErrors {
len: 3,
fail: vec![2],
policy: StreamErrorPolicy::SkipAndContinue,
flushed: Arc::new(Mutex::new(Vec::new())),
};
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("skip-cursor");
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
assert_eq!(
step_cursors(&store),
vec![1u64, 3],
"the skipped item's cursor (2) is never committed"
);
}
struct CountWindowSource {
window: StreamWindow,
windows: Arc<Mutex<Vec<Vec<u64>>>>,
}
#[task::stream]
impl StreamTask<Step> for CountWindowSource {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
self.window.clone()
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u64, 2, 3])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.windows.lock().unwrap().push(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn count_zero_window_clamps_to_per_item() {
let windows = Arc::new(Mutex::new(Vec::new()));
let task = CountWindowSource {
window: StreamWindow::Count(0),
windows: Arc::clone(&windows),
};
let res = Resources::new();
let result = Task::run(&task, &res).await.unwrap();
assert_eq!(result, TaskResult::Single(Step::Done));
assert_eq!(
*windows.lock().unwrap(),
vec![vec![1u64], vec![2], vec![3]],
"Count(0) behaves like Count(1)"
);
}
struct CountingFlush {
len: u64,
window: StreamWindow,
flushes: Arc<Mutex<Vec<Vec<u64>>>>,
}
#[task::stream]
impl StreamTask<Step> for CountingFlush {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
self.window.clone()
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
let items: Vec<u64> = (1..=self.len).collect();
Ok(Box::pin(stream::iter(items)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.flushes.lock().unwrap().push(outputs);
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn empty_stream_closes_without_flush_or_cursor() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let flushes = Arc::new(Mutex::new(Vec::new()));
let task = CountingFlush {
len: 0,
window: StreamWindow::Count(2),
flushes: Arc::clone(&flushes),
};
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("empty");
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
assert!(
flushes.lock().unwrap().is_empty(),
"no flush for an empty source"
);
assert!(
step_cursors(&store).is_empty(),
"no cursor committed when nothing is processed"
);
}
#[tokio::test]
async fn exhaust_exact_divide_commits_only_full_window_cursors() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let flushes = Arc::new(Mutex::new(Vec::new()));
let task = CountingFlush {
len: 4,
window: StreamWindow::Count(2),
flushes: Arc::clone(&flushes),
};
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("exact");
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
assert_eq!(*flushes.lock().unwrap(), vec![vec![1u64, 2], vec![3, 4]]);
assert_eq!(step_cursors(&store), vec![2u64, 4]);
}
#[tokio::test]
async fn exhaust_partial_window_commits_its_cursor() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let flushes = Arc::new(Mutex::new(Vec::new()));
let task = CountingFlush {
len: 5,
window: StreamWindow::Count(2),
flushes: Arc::clone(&flushes),
};
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("partial");
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
assert_eq!(
*flushes.lock().unwrap(),
vec![vec![1u64, 2], vec![3, 4], vec![5]]
);
assert_eq!(step_cursors(&store), vec![2u64, 4, 5]);
}
struct CancelMidWindow {
handle: crate::cancel::CancellationHandle,
cancel_after: u32,
seen: AtomicU32,
window: StreamWindow,
stop_on_flush: bool,
flushed: Arc<Mutex<Vec<Vec<u64>>>>,
close_reason: Arc<Mutex<Option<CloseReason>>>,
}
#[task::stream]
impl StreamTask<S3> for CancelMidWindow {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
self.window.clone()
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(0u64..)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
let n = self.seen.fetch_add(1, Ordering::SeqCst) + 1;
if n == self.cancel_after {
self.handle.cancel();
}
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<S3>, CanoError> {
self.flushed.lock().unwrap().push(outputs);
if self.stop_on_flush {
Ok(WindowSignal::Stop(TaskResult::Single(S3::ViaStop)))
} else {
Ok(WindowSignal::Continue)
}
}
async fn on_close(
&self,
_res: &Resources,
reason: CloseReason,
) -> Result<TaskResult<S3>, CanoError> {
*self.close_reason.lock().unwrap() = Some(reason);
Ok(TaskResult::Single(S3::ViaClose))
}
}
#[tokio::test]
async fn cancel_flushes_partial_in_flight_window() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let flushed = Arc::new(Mutex::new(Vec::new()));
let close_reason = Arc::new(Mutex::new(None));
let (handle, token) = CancellationToken::new();
let task = CancelMidWindow {
handle,
cancel_after: 2,
seen: AtomicU32::new(0),
window: StreamWindow::Count(3),
stop_on_flush: false,
flushed: Arc::clone(&flushed),
close_reason: Arc::clone(&close_reason),
};
let workflow = Workflow::bare()
.register_stream(S3::Consume, task)
.add_exit_states([S3::ViaStop, S3::ViaClose]);
let result = workflow.orchestrate(S3::Consume, token).await;
assert!(
matches!(&result, Err(e) if e.category() == "cancelled"),
"got {result:?}"
);
assert_eq!(
*flushed.lock().unwrap(),
vec![vec![0u64, 1]],
"the partial window flushes once on cancel"
);
assert_eq!(*close_reason.lock().unwrap(), Some(CloseReason::Cancelled));
}
#[tokio::test]
async fn cancel_ignores_stop_from_partial_flush() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (handle, token) = CancellationToken::new();
let task = CancelMidWindow {
handle,
cancel_after: 2,
seen: AtomicU32::new(0),
window: StreamWindow::Count(3),
stop_on_flush: true,
flushed: Arc::new(Mutex::new(Vec::new())),
close_reason: Arc::new(Mutex::new(None)),
};
let workflow = Workflow::bare()
.register_stream(S3::Consume, task)
.add_exit_states([S3::ViaStop, S3::ViaClose]);
let result = workflow.orchestrate(S3::Consume, token).await;
assert!(
matches!(&result, Err(e) if e.category() == "cancelled"),
"Stop from the cancel-drain flush must not transition, got {result:?}"
);
}
struct CancelAfterWindow {
handle: crate::cancel::CancellationHandle,
flushes: AtomicU32,
close_errors: bool,
closed_cancelled: Arc<AtomicBool>,
}
#[task::stream]
impl StreamTask<Step> for CancelAfterWindow {
type Item = u64;
type Output = u64;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(0u64..)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
self.flushes.fetch_add(1, Ordering::SeqCst);
self.handle.cancel();
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
if reason == CloseReason::Cancelled {
self.closed_cancelled.store(true, Ordering::SeqCst);
if self.close_errors {
return Err(CanoError::task_execution("cleanup failed"));
}
}
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn cancel_with_empty_buffer_skips_flush_but_runs_close() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (handle, token) = CancellationToken::new();
let closed = Arc::new(AtomicBool::new(false));
let task = CancelAfterWindow {
handle,
flushes: AtomicU32::new(0),
close_errors: false,
closed_cancelled: Arc::clone(&closed),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
let result = workflow.orchestrate(Step::Consume, token).await;
assert!(
matches!(&result, Err(e) if e.category() == "cancelled"),
"got {result:?}"
);
assert!(
closed.load(Ordering::SeqCst),
"on_close(Cancelled) still runs with an empty buffer"
);
}
#[tokio::test]
async fn cancel_propagates_on_close_error() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (handle, token) = CancellationToken::new();
let task = CancelAfterWindow {
handle,
flushes: AtomicU32::new(0),
close_errors: true,
closed_cancelled: Arc::new(AtomicBool::new(false)),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
let err = workflow
.orchestrate(Step::Consume, token)
.await
.unwrap_err();
assert_eq!(err.category(), "task_execution");
assert!(err.to_string().contains("cleanup failed"), "got {err}");
}
struct CancelThenResume {
handle: crate::cancel::CancellationHandle,
seen: AtomicU32,
opened_cursors: Arc<Mutex<Vec<Option<u64>>>>,
}
#[task::stream]
impl StreamTask<Step> for CancelThenResume {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(3)
}
async fn open(
&self,
_res: &Resources,
cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
self.opened_cursors.lock().unwrap().push(cursor);
let items: Vec<u64> = match cursor {
None => (0u64..1000).collect(),
Some(_) => Vec::new(),
};
Ok(Box::pin(stream::iter(items)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
let n = self.seen.fetch_add(1, Ordering::SeqCst) + 1;
if n == 2 {
self.handle.cancel();
}
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn cancel_commits_partial_cursor_and_resumes() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (handle, token) = CancellationToken::new();
let opened = Arc::new(Mutex::new(Vec::new()));
let task = CancelThenResume {
handle,
seen: AtomicU32::new(0),
opened_cursors: Arc::clone(&opened),
};
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("cancel-resume");
let r1 = workflow.orchestrate(Step::Consume, token).await;
assert!(
matches!(&r1, Err(e) if e.category() == "cancelled"),
"got {r1:?}"
);
assert_eq!(
step_cursors(&store),
vec![1u64],
"the cancelled run commits the in-flight window's final cursor"
);
let r2 = workflow
.resume_from("cancel-resume", CancellationToken::disabled())
.await
.unwrap();
assert_eq!(r2, Step::Done);
assert_eq!(*opened.lock().unwrap(), vec![None, Some(1)]);
}
struct SplitOnClose;
#[task::stream]
impl StreamTask<Step> for SplitOnClose {
type Item = u64;
type Output = u64;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u64])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Split(vec![Step::Done, Step::Done]))
}
}
#[tokio::test]
async fn stream_split_result_is_rejected() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let workflow = Workflow::bare()
.register_stream(Step::Consume, SplitOnClose)
.add_exit_state(Step::Done);
let err = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap_err();
assert_eq!(err.category(), "workflow");
assert!(err.to_string().contains("split"), "got {err}");
}
struct CrashAfterFirstWindow;
#[task::stream]
impl StreamTask<Step> for CrashAfterFirstWindow {
type Item = u64;
type Output = u64;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u64, 2])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
if outputs == vec![2u64] {
return Err(CanoError::task_execution("crash"));
}
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn corrupt_cursor_fails_to_deserialize_on_resume() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let store = Arc::new(InMemoryStore::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, CrashAfterFirstWindow)
.add_exit_state(Step::Done)
.with_checkpoint_store(store.clone())
.with_workflow_id("corrupt");
let r1 = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await;
assert!(r1.is_err(), "run 1 should crash: {r1:?}");
{
let mut g = store.rows.lock().unwrap();
for row in g.get_mut("corrupt").unwrap().iter_mut() {
if row.kind == crate::recovery::RowKind::StepCursor {
row.output_blob = Some(b"not-json".to_vec());
}
}
}
let err = workflow
.resume_from("corrupt", CancellationToken::disabled())
.await
.unwrap_err();
assert_eq!(err.category(), "task_execution");
assert!(
err.to_string().contains("deserialize stream cursor"),
"got {err}"
);
}
#[derive(Default)]
struct CursorAppendFails;
#[crate::checkpoint_store]
impl crate::recovery::CheckpointStore for CursorAppendFails {
async fn append(
&self,
_workflow_id: &str,
row: crate::recovery::CheckpointRow,
) -> Result<(), CanoError> {
if row.kind == crate::recovery::RowKind::StepCursor {
Err(CanoError::checkpoint_store("disk full"))
} else {
Ok(())
}
}
async fn load_run(
&self,
_workflow_id: &str,
) -> Result<Vec<crate::recovery::CheckpointRow>, CanoError> {
Ok(Vec::new())
}
async fn clear(&self, _workflow_id: &str) -> Result<(), CanoError> {
Ok(())
}
}
#[tokio::test]
async fn checkpoint_append_failure_surfaces() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let task = CountingFlush {
len: 3,
window: StreamWindow::Count(1),
flushes: Arc::new(Mutex::new(Vec::new())),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(Arc::new(CursorAppendFails))
.with_workflow_id("append-fail");
let err = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap_err();
assert_eq!(err.category(), "checkpoint_store");
assert!(
err.to_string().contains("append stream cursor checkpoint"),
"got {err}"
);
}
struct PanicInFlush;
#[task::stream]
impl StreamTask<Step> for PanicInFlush {
type Item = u64;
type Output = u64;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u64])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
panic!("boom in flush_window");
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn panic_in_callback_becomes_error() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let workflow = Workflow::bare()
.register_stream(Step::Consume, PanicInFlush)
.add_exit_state(Step::Done);
let err = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap_err();
assert_eq!(err.category(), "task_execution");
assert!(err.to_string().contains("panic"), "got {err}");
}
struct AlwaysFails {
opened: Arc<AtomicU32>,
}
#[task::stream]
impl StreamTask<Step> for AlwaysFails {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn config(&self) -> TaskConfig {
TaskConfig::minimal().with_fixed_retry(2, std::time::Duration::from_millis(1))
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
self.opened.fetch_add(1, Ordering::SeqCst);
Ok(Box::pin(stream::iter(vec![1u64])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(
&self,
_res: &Resources,
_item: u64,
) -> Result<(u64, u64), CanoError> {
Err(CanoError::task_execution("always fails"))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn outer_retry_not_applied_open_called_once() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let opened = Arc::new(AtomicU32::new(0));
let task = AlwaysFails {
opened: Arc::clone(&opened),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await;
assert!(result.is_err(), "FailFast item error fails the run");
assert_eq!(
opened.load(Ordering::SeqCst),
1,
"config().max_attempts must not re-open/re-consume the stream"
);
}
#[derive(Default)]
struct EventLog {
events: Mutex<Vec<String>>,
}
impl crate::observer::WorkflowObserver for EventLog {
fn on_task_start(&self, task_id: &str) {
self.events.lock().unwrap().push(format!("start:{task_id}"));
}
fn on_task_success(&self, task_id: &str) {
self.events
.lock()
.unwrap()
.push(format!("success:{task_id}"));
}
fn on_task_failure(&self, _task_id: &str, _err: &CanoError) {
self.events.lock().unwrap().push("failure".to_string());
}
fn on_cancelled(&self, _state: &str) {
self.events.lock().unwrap().push("cancelled".to_string());
}
}
struct NamedExhaust;
#[task::stream]
impl StreamTask<Step> for NamedExhaust {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn name(&self) -> Cow<'static, str> {
Cow::Borrowed("my-custom-stream")
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u64])) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn name_override_forwarded_to_observer() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let log = Arc::new(EventLog::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, NamedExhaust)
.add_exit_state(Step::Done)
.with_observer(log.clone());
workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
let events = log.events.lock().unwrap();
assert_eq!(
*events,
vec![
"start:my-custom-stream".to_string(),
"success:my-custom-stream".to_string()
],
"the StreamTask name() override reaches observer hooks"
);
}
#[tokio::test]
async fn cancel_fires_full_observer_sequence() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let (handle, token) = CancellationToken::new();
let task = CancelAfterWindow {
handle,
flushes: AtomicU32::new(0),
close_errors: false,
closed_cancelled: Arc::new(AtomicBool::new(false)),
};
let log = Arc::new(EventLog::default());
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_observer(log.clone());
let result = workflow.orchestrate(Step::Consume, token).await;
assert!(matches!(&result, Err(e) if e.category() == "cancelled"));
let events = log.events.lock().unwrap();
assert_eq!(events.len(), 3, "exactly three hooks fire, got {events:?}");
assert!(
events[0].starts_with("start:") && events[0].contains("CancelAfterWindow"),
"first hook is on_task_start, got {events:?}"
);
assert_eq!(
&events[1..],
&["failure".to_string(), "cancelled".to_string()],
"cancel fires start → failure → cancelled, got {events:?}"
);
}
#[tokio::test]
async fn register_stream_without_store_completes() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
let task = CountingFlush {
len: 3,
window: StreamWindow::Count(2),
flushes: Arc::new(Mutex::new(Vec::new())),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
}
struct DelayedSource {
emit_delay: std::time::Duration,
items: Vec<u64>,
}
#[task::stream]
impl StreamTask<Step> for DelayedSource {
type Item = u64;
type Output = u64;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Duration(std::time::Duration::from_millis(50))
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
let items = self.items.clone();
let delay = self.emit_delay;
Ok(Box::pin(stream::unfold(
(items.into_iter(), delay),
|(mut iter, d)| async move {
let item = iter.next()?;
Some((item, (iter, d)))
},
)) as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u64) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
#[tokio::test]
async fn duration_window_flushes_on_elapsed_time() {
use crate::cancel::CancellationToken;
use crate::workflow::Workflow;
tokio::time::pause();
let task = DelayedSource {
emit_delay: std::time::Duration::from_millis(10),
items: vec![1, 2, 3],
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done);
let result = workflow
.orchestrate(Step::Consume, CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
}
#[tokio::test]
async fn crash_resume_commits_cursor_between_windows() {
use crate::cancel::CancellationToken;
use crate::recovery::{CheckpointRow, CheckpointStore};
use crate::workflow::Workflow;
#[derive(Default)]
struct PreExistingCursor;
#[crate::checkpoint_store]
impl CheckpointStore for PreExistingCursor {
async fn append(
&self,
_workflow_id: &str,
_row: CheckpointRow,
) -> Result<(), CanoError> {
Ok(())
}
async fn load_run(&self, _workflow_id: &str) -> Result<Vec<CheckpointRow>, CanoError> {
let cursor_bytes = serde_json::to_vec(&2u64).unwrap();
Ok(vec![
CheckpointRow::new(1, "Consume", "ResumeFromCursor").with_workflow_version(0),
CheckpointRow::new(2, "Consume", "ResumeFromCursor")
.with_cursor(cursor_bytes)
.with_workflow_version(0),
])
}
async fn clear(&self, _workflow_id: &str) -> Result<(), CanoError> {
Ok(())
}
}
struct ResumeFromCursor {
opened_cursor: std::sync::Arc<std::sync::Mutex<Option<u64>>>,
}
#[task::stream]
impl StreamTask<Step> for ResumeFromCursor {
type Item = u64;
type Output = u64;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u64> + Send>>, CanoError> {
*self.opened_cursor.lock().unwrap() = cursor;
let start = cursor.unwrap_or(0) + 1;
Ok(Box::pin(stream::iter(vec![start, start + 1, start + 2]))
as Pin<Box<dyn Stream<Item = u64> + Send>>)
}
async fn process_item(
&self,
_res: &Resources,
item: u64,
) -> Result<(u64, u64), CanoError> {
Ok((item, item))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u64>,
) -> Result<WindowSignal<Step>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<Step>, CanoError> {
Ok(TaskResult::Single(Step::Done))
}
}
let opened_cursor = std::sync::Arc::new(std::sync::Mutex::new(None::<u64>));
let task = ResumeFromCursor {
opened_cursor: std::sync::Arc::clone(&opened_cursor),
};
let workflow = Workflow::bare()
.register_stream(Step::Consume, task)
.add_exit_state(Step::Done)
.with_checkpoint_store(Arc::new(PreExistingCursor))
.with_workflow_id("crash-resume");
let result = workflow
.resume_from("crash-resume", CancellationToken::disabled())
.await
.unwrap();
assert_eq!(result, Step::Done);
assert_eq!(*opened_cursor.lock().unwrap(), Some(2));
}
}
#[cfg(all(test, feature = "metrics"))]
mod metrics_tests {
use super::*;
use crate::cancel::CancellationToken;
use crate::metrics::test_support::*;
use crate::task;
use crate::task::Task;
use crate::workflow::Workflow;
use futures_util::stream;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum St {
Consume,
Done,
}
struct FiveItems;
#[task::stream]
impl StreamTask<St> for FiveItems {
type Item = u32;
type Output = u32;
type Cursor = u64;
fn window(&self) -> StreamWindow {
StreamWindow::Count(2)
}
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32, 2, 3, 4, 5]))
as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<St>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<St>, CanoError> {
Ok(TaskResult::Single(St::Done))
}
}
#[test]
fn stream_metrics_counted_correctly() {
let (result, rows) = run_with_recorder(|| async {
let workflow = Workflow::bare()
.register_stream(St::Consume, FiveItems)
.add_exit_state(St::Done);
workflow
.orchestrate(St::Consume, CancellationToken::disabled())
.await
});
assert!(result.is_ok(), "workflow should succeed: {result:?}");
assert_eq!(
counter(&rows, "cano_stream_runs_total", &[("outcome", "completed")]),
1,
"one completed stream run"
);
assert_eq!(
counter(&rows, "cano_stream_windows_total", &[]),
3,
"three windows flushed"
);
assert_eq!(
counter(&rows, "cano_stream_items_total", &[("result", "ok")]),
5,
"five ok items"
);
}
struct SelfCancel {
handle: crate::cancel::CancellationHandle,
}
#[task::stream]
impl StreamTask<St> for SelfCancel {
type Item = u32;
type Output = u32;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(0u32..)) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<St>, CanoError> {
self.handle.cancel();
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<St>, CanoError> {
Ok(TaskResult::Single(St::Done))
}
}
#[test]
fn cancelled_stream_records_cancelled_outcome() {
let (handle, token) = CancellationToken::new();
let (result, rows) = run_with_recorder(|| async move {
let workflow = Workflow::bare()
.register_stream(St::Consume, SelfCancel { handle })
.add_exit_state(St::Done);
workflow.orchestrate(St::Consume, token).await
});
assert!(result.is_err(), "a cancelled run is Err: {result:?}");
assert_eq!(
counter(&rows, "cano_stream_runs_total", &[("outcome", "cancelled")]),
1,
"a cooperative cancel is recorded as cancelled, not failed"
);
}
struct NeverStops;
#[task::stream]
impl StreamTask<St> for NeverStops {
type Item = u32;
type Output = u32;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(0u32..)) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
tokio::task::yield_now().await;
Ok((item, item as u64))
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<St>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<St>, CanoError> {
Ok(TaskResult::Single(St::Done))
}
}
#[test]
fn total_timeout_stream_records_failed_not_cancelled() {
let (result, rows) = run_with_recorder(|| async {
let workflow = Workflow::bare()
.register_stream(St::Consume, NeverStops)
.add_exit_state(St::Done)
.with_total_timeout(std::time::Duration::from_nanos(1));
workflow
.orchestrate(St::Consume, CancellationToken::disabled())
.await
});
assert!(
matches!(&result, Err(e) if e.category() == "workflow_timeout"),
"a tripped total-timeout budget must surface as WorkflowTimeout: {result:?}"
);
assert_eq!(
counter(&rows, "cano_stream_runs_total", &[("outcome", "failed")]),
1,
"a total-timeout trip is recorded as failed, not cancelled — `cancelled` is \
reserved for a real CancellationToken firing"
);
assert_eq!(
counter_opt(&rows, "cano_stream_runs_total", &[("outcome", "cancelled")]),
None,
"must not also be recorded as cancelled"
);
}
struct FailSecond;
#[task::stream]
impl StreamTask<St> for FailSecond {
type Item = u32;
type Output = u32;
type Cursor = u64;
async fn open(
&self,
_res: &Resources,
_cursor: Option<u64>,
) -> Result<Pin<Box<dyn Stream<Item = u32> + Send>>, CanoError> {
Ok(Box::pin(stream::iter(vec![1u32, 2])) as Pin<Box<dyn Stream<Item = u32> + Send>>)
}
async fn process_item(&self, _res: &Resources, item: u32) -> Result<(u32, u64), CanoError> {
if item == 2 {
Err(CanoError::task_execution("boom"))
} else {
Ok((item, item as u64))
}
}
async fn flush_window(
&self,
_res: &Resources,
_outputs: Vec<u32>,
) -> Result<WindowSignal<St>, CanoError> {
Ok(WindowSignal::Continue)
}
async fn on_close(
&self,
_res: &Resources,
_reason: CloseReason,
) -> Result<TaskResult<St>, CanoError> {
Ok(TaskResult::Single(St::Done))
}
}
#[test]
fn failed_stream_records_failed_outcome_and_err_item() {
let (result, rows) = run_with_recorder(|| async {
let workflow = Workflow::bare()
.register_stream(St::Consume, FailSecond)
.add_exit_state(St::Done);
workflow
.orchestrate(St::Consume, CancellationToken::disabled())
.await
});
assert!(
result.is_err(),
"FailFast item error fails the run: {result:?}"
);
assert_eq!(
counter(&rows, "cano_stream_runs_total", &[("outcome", "failed")]),
1,
"a genuine error is recorded as failed"
);
assert_eq!(
counter(&rows, "cano_stream_items_total", &[("result", "ok")]),
1,
"item 1 processed ok"
);
assert_eq!(
counter(&rows, "cano_stream_items_total", &[("result", "err")]),
1,
"item 2 recorded as an err item"
);
}
#[test]
fn inmemory_completed_records_completed_outcome() {
let (result, rows) = run_with_recorder(|| async {
let res = Resources::new();
Task::run(&FiveItems, &res).await
});
assert!(result.is_ok(), "{result:?}");
assert_eq!(
counter(&rows, "cano_stream_runs_total", &[("outcome", "completed")]),
1,
"the in-memory companion records a completed run"
);
}
#[test]
fn inmemory_failed_records_failed_outcome() {
let (result, rows) = run_with_recorder(|| async {
let res = Resources::new();
Task::run(&FailSecond, &res).await
});
assert!(result.is_err(), "{result:?}");
assert_eq!(
counter(&rows, "cano_stream_runs_total", &[("outcome", "failed")]),
1,
"the in-memory companion records a failed run (never cancelled)"
);
}
}