#[cfg(test)]
mod checkpoint_end_order_tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use crate::runtime::chain_segment::{NamedStep, sink_chain_factory, source_chain_factory};
use crate::runtime::chained_stream_task::ChainedStreamTask;
use tinyflow_api::checkpoint::CheckpointContext;
use tinyflow_api::context::RuntimeContext;
use tinyflow_api::error::{StreamError, StreamResult};
use tinyflow_api::functions::{LifecycleAware, SinkFunction, SourceFunction};
use tinyflow_api::state::in_memory_job_state;
static SINK_CHECKPOINT_END_CALLS: AtomicUsize = AtomicUsize::new(0);
#[derive(Clone, Default)]
struct FailCheckpointEndSource;
#[async_trait]
impl LifecycleAware for FailCheckpointEndSource {
async fn on_checkpoint_end(
&mut self,
_runtime: &mut RuntimeContext,
_ckpt: &CheckpointContext,
) -> StreamResult<()> {
Err(StreamError::TaskFailed {
task_id: 0,
reason: "simulated kafka commit failure".into(),
})
}
}
#[async_trait]
impl SourceFunction<i32> for FailCheckpointEndSource {
async fn poll_next(&mut self) -> StreamResult<Option<i32>> {
Ok(None)
}
}
#[derive(Clone, Default)]
struct RecordingSink;
#[async_trait]
impl LifecycleAware for RecordingSink {
async fn on_checkpoint_end(
&mut self,
_: &mut RuntimeContext,
_: &CheckpointContext,
) -> StreamResult<()> {
SINK_CHECKPOINT_END_CALLS.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[async_trait]
impl SinkFunction<i32> for RecordingSink {
async fn accept(&mut self, _: i32, _: &mut RuntimeContext) -> StreamResult<()> {
Ok(())
}
}
#[tokio::test]
async fn checkpoint_end_stops_when_source_fails_before_sink() {
SINK_CHECKPOINT_END_CALLS.store(0, Ordering::SeqCst);
let source_factory = Arc::new(Mutex::new(|| FailCheckpointEndSource));
let factories = [
source_chain_factory(Arc::clone(&source_factory)),
sink_chain_factory(RecordingSink),
];
let steps: Vec<NamedStep> = factories.iter().map(|f| f()).collect();
let (ack_tx, _ack_rx) = tokio::sync::mpsc::channel(8);
let (_ctrl_tx, ctrl_rx) = tokio::sync::mpsc::channel(8);
let mut task = ChainedStreamTask::from_steps(
1,
"test-0",
0,
0,
1,
steps,
ctrl_rx,
ack_tx,
in_memory_job_state(),
None,
)
.unwrap();
let mut ctx = RuntimeContext::new(1, "test-0", 0, 1, 0, in_memory_job_state(), None);
task.open_all(&mut ctx).await.unwrap();
let ckpt = CheckpointContext::new(1, 0);
let err = task
.on_checkpoint_end_all(&mut ctx, &ckpt)
.await
.unwrap_err();
assert!(
format!("{err}").contains("simulated kafka commit failure"),
"unexpected error: {err}"
);
assert_eq!(
SINK_CHECKPOINT_END_CALLS.load(Ordering::SeqCst),
0,
"sink on_checkpoint_end must not run after source failure"
);
}
}
#[cfg(test)]
mod collect_metrics_tests {
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use crate::runtime::chain_segment::{NamedStep, sink_chain_factory, source_chain_factory};
use crate::runtime::chained_stream_task::ChainedStreamTask;
use tinyflow_api::checkpoint::CheckpointContext;
use tinyflow_api::context::RuntimeContext;
use tinyflow_api::error::StreamResult;
use tinyflow_api::functions::{LifecycleAware, SinkFunction, SourceFunction};
use tinyflow_api::state::in_memory_job_state;
#[derive(Clone, Default)]
struct CountingSource;
#[async_trait]
impl LifecycleAware for CountingSource {}
#[async_trait]
impl SourceFunction<i32> for CountingSource {
async fn poll_next(&mut self) -> StreamResult<Option<i32>> {
Ok(None)
}
}
#[derive(Clone, Default)]
struct DummySink;
#[async_trait]
impl LifecycleAware for DummySink {}
#[async_trait]
impl SinkFunction<i32> for DummySink {
async fn accept(&mut self, _: i32, _: &mut RuntimeContext) -> StreamResult<()> {
Ok(())
}
}
#[tokio::test]
async fn checkpoint_stores_metrics_each_time() {
let source_factory = Arc::new(Mutex::new(|| CountingSource::default()));
let factories = [
source_chain_factory(Arc::clone(&source_factory)),
sink_chain_factory(DummySink::default()),
];
let steps: Vec<NamedStep> = factories.iter().map(|f| f()).collect();
let job_state = in_memory_job_state();
let (ack_tx, _ack_rx) = tokio::sync::mpsc::channel(8);
let (_ctrl_tx, ctrl_rx) = tokio::sync::mpsc::channel(8);
let mut task = ChainedStreamTask::from_steps(
1,
"test-0",
0,
0,
1,
steps,
ctrl_rx,
ack_tx,
Arc::clone(&job_state),
None,
)
.unwrap();
let mut ctx = RuntimeContext::new(1, "test-0", 0, 1, 0, Arc::clone(&job_state), None);
task.open_all(&mut ctx).await.unwrap();
let ckpt1 = CheckpointContext::new(1, 1000);
task.on_checkpoint_end_all(&mut ctx, &ckpt1).await.unwrap();
let val = job_state
.get("job_status", b"metric/0")
.await
.unwrap()
.unwrap();
let metrics1: serde_json::Value = serde_json::from_slice(&val).unwrap();
assert!(
metrics1.get("Source/0").is_some(),
"metric on first checkpoint"
);
let ckpt2 = CheckpointContext::new(2, 2000);
task.on_checkpoint_end_all(&mut ctx, &ckpt2).await.unwrap();
let val = job_state
.get("job_status", b"metric/0")
.await
.unwrap()
.unwrap();
let metrics2: serde_json::Value = serde_json::from_slice(&val).unwrap();
assert!(
metrics2.get("Source/0").is_some(),
"metric on second checkpoint"
);
assert!(metrics2.get("Sink/0").is_some(), "sink metric should exist");
}
#[tokio::test]
async fn rate_is_calculated_correctly() {
#[derive(Clone)]
struct PollingSource {
poll_count: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}
impl Default for PollingSource {
fn default() -> Self {
Self {
poll_count: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LifecycleAware for PollingSource {}
#[async_trait]
impl SourceFunction<i32> for PollingSource {
async fn poll_next(&mut self) -> StreamResult<Option<i32>> {
if self
.poll_count
.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
< 5
{
Ok(Some(42))
} else {
Ok(None)
}
}
}
#[derive(Clone)]
struct CountingSink {
count: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}
impl Default for CountingSink {
fn default() -> Self {
Self {
count: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LifecycleAware for CountingSink {}
#[async_trait]
impl SinkFunction<i32> for CountingSink {
async fn accept(&mut self, _: i32, _: &mut RuntimeContext) -> StreamResult<()> {
self.count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(())
}
}
let source_factory = Arc::new(Mutex::new(move || PollingSource::default()));
let factories = [
source_chain_factory(Arc::clone(&source_factory)),
sink_chain_factory(CountingSink::default()),
];
let mut steps: Vec<NamedStep> = factories.iter().map(|f| f()).collect();
steps[0].name = "my_source".into();
steps[1].name = "my_sink".into();
let job_state = in_memory_job_state();
let (ack_tx, _ack_rx) = tokio::sync::mpsc::channel(8);
let (_ctrl_tx, ctrl_rx) = tokio::sync::mpsc::channel(8);
let mut task = ChainedStreamTask::from_steps(
1,
"test-0",
0,
0,
1,
steps,
ctrl_rx,
ack_tx,
Arc::clone(&job_state),
None,
)
.unwrap();
let mut ctx = RuntimeContext::new(1, "test-0", 0, 1, 0, Arc::clone(&job_state), None);
task.open_all(&mut ctx).await.unwrap();
let ckpt1 = CheckpointContext::new(1, 1000);
task.on_checkpoint_end_all(&mut ctx, &ckpt1).await.unwrap();
let ckpt2 = CheckpointContext::new(2, 2000);
task.on_checkpoint_end_all(&mut ctx, &ckpt2).await.unwrap();
let val = job_state
.get("job_status", b"metric/0")
.await
.unwrap()
.unwrap();
let metrics: serde_json::Value = serde_json::from_slice(&val).unwrap();
assert!(
metrics.get("my_source/0").is_some(),
"metric key should use resolved step name"
);
assert!(
metrics.get("my_sink/0").is_some(),
"metric key should use resolved step name"
);
}
}