#![allow(clippy::print_stdout, clippy::print_stderr)]
use spate::prelude::*;
use spate::source::LaneId;
use spate_test::{
BytesPassthrough, CaptureSink, MemorySource, PipelineRun, ScriptedResult, TestEncoder,
WriteOutcome, capture_sink, decode_rows, memory_source, wait_until,
};
use std::time::Duration;
const ORDERS: [&str; 3] = [
"order_placed:1001",
"payment_captured:1001",
"refund_issued:1001",
];
fn prompt_pool() -> SinkPoolConfig {
let mut cfg = SinkPoolConfig::default();
cfg.batch.linger = Duration::from_millis(20);
cfg
}
fn assemble(
config: &str,
sink: CaptureSink,
source: MemorySource,
) -> Result<PipelineRuntime<MemorySource>, Box<dyn std::error::Error>> {
Ok(Pipeline::from_config(PipelineConfig::from_str(config)?)?
.sink(sink)?
.chains(|ctx| {
let chunk_cfg = ctx.chunk();
chain_owned::<Vec<u8>, _>(BytesPassthrough)
.with_metrics(ctx.pipeline, "main")
.sink(
TestEncoder,
KeyHashRouter,
chunk_cfg,
ctx.queues,
ctx.budget,
)
.build()
})
.runtime_options(RuntimeOptions {
handle_signals: false, ..RuntimeOptions::default()
})
.into_runtime(source)?)
}
fn retryable_write() -> Result<(), Box<dyn std::error::Error>> {
const CONFIG: &str = r#"
pipeline: { name: sink-failures-retryable, threads: 1, io_threads: 2 }
admin: { listen: none }
metrics: { exporter: none }
checkpoint: { interval: 100ms, drain_timeout: 5s }
source: { memory: {} }
sink: { capture: {} }
"#;
let (source, handle) = memory_source();
let (sink, script) = capture_sink(1, 2);
let sink = sink.with_pool_config({
let mut pool = prompt_pool();
pool.inflight.max_per_shard = 1;
pool
});
script.enqueue_global(WriteOutcome::retryable("payments replica restarting"));
let runtime = assemble(CONFIG, sink, source)?;
let shutdown = runtime.shutdown_handle();
let run = PipelineRun::spawn(move || runtime.run());
let p = PartitionId(0);
handle.assign_lanes(&[(LaneId(0), p)]);
let last = *handle.push_many(p, ORDERS).last().expect("three offsets");
assert!(
handle.wait_committed(p, last + 1, Duration::from_secs(30)),
"the watermark advances after the retry (last committed: {:?})",
handle.last_committed(p),
);
shutdown.trigger();
let report = run
.wait_exit(Duration::from_secs(30))
.expect("the pipeline drains after shutdown")?;
assert_eq!(report.exit_code(), 0, "a retried write still drains clean");
let writes = script.writes();
let failed = writes
.iter()
.position(|w| matches!(w.result, ScriptedResult::Retryable(_)))
.expect("the scripted failure was attempted");
let retry = writes[failed + 1..]
.iter()
.find(|w| w.dedup_token == writes[failed].dedup_token)
.expect("the same batch was re-sent");
assert_eq!(retry.result, ScriptedResult::Ok, "the retry landed");
assert_ne!(
retry.replica, writes[failed].replica,
"the retry rotated onto the other replica"
);
assert_eq!(retry.rows, writes[failed].rows, "the same sealed batch");
println!(
"\n1. retryable: batch {} refused by replica {}, accepted by replica {} ({} rows)",
retry.dedup_token, writes[failed].replica, retry.replica, retry.rows,
);
Ok(())
}
fn failed_probe() -> Result<(), Box<dyn std::error::Error>> {
const CONFIG: &str = r#"
pipeline: { name: sink-failures-probe, threads: 1, io_threads: 2 }
checkpoint: { interval: 100ms, drain_timeout: 5s }
admin: { listen: none }
metrics: { exporter: none }
source: { memory: {} }
sink: { capture: {} }
"#;
let (source, handle) = memory_source();
let (sink, script) = capture_sink(1, 2);
let sink = sink.with_pool_config(prompt_pool());
let probe = sink
.clone()
.into_parts()
.probe
.expect("the capturing sink ships a readiness probe");
let asker = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
asker
.block_on(probe())
.expect("every replica answers before the outage");
let runtime = assemble(CONFIG, sink, source)?;
let shutdown = runtime.shutdown_handle();
let run = PipelineRun::spawn(move || runtime.run());
let p = PartitionId(0);
handle.assign_lanes(&[(LaneId(0), p)]);
let before = *handle.push_many(p, ORDERS).last().expect("three offsets");
assert!(handle.wait_committed(p, before + 1, Duration::from_secs(30)));
script.fail_probe(0, 1, "orders replica 1 unreachable");
assert!(
asker.block_on(probe()).is_err(),
"one unreachable replica fails readiness for the sink"
);
let during = *handle.push_many(p, ORDERS).last().expect("three offsets");
assert!(
handle.wait_committed(p, during + 1, Duration::from_secs(30)),
"records still commit while readiness is red (last committed: {:?})",
handle.last_committed(p),
);
script.heal_probe(0, 1);
asker.block_on(probe()).expect("readiness recovers");
shutdown.trigger();
let report = run
.wait_exit(Duration::from_secs(30))
.expect("the pipeline drains after shutdown")?;
assert_eq!(report.exit_code(), 0, "a red probe does not fail the run");
println!("\n2. probe: readiness went red and back while every record committed");
Ok(())
}
fn fatal_write() -> Result<(), Box<dyn std::error::Error>> {
const CONFIG: &str = r#"
pipeline: { name: sink-failures-fatal, threads: 1, io_threads: 2 }
checkpoint: { interval: 100ms, drain_timeout: 5s, stalled_fail_after: 1s }
admin: { listen: none }
metrics: { exporter: none }
source: { memory: {} }
sink: { capture: {} }
"#;
let (source, handle) = memory_source();
let (sink, script) = capture_sink(1, 1);
let sink = sink.with_pool_config(prompt_pool());
let runtime = assemble(CONFIG, sink, source)?;
let run = PipelineRun::spawn(move || runtime.run());
let p = PartitionId(0);
handle.assign_lanes(&[(LaneId(0), p)]);
let healthy = *handle.push_many(p, ORDERS).last().expect("three offsets");
assert!(handle.wait_committed(p, healthy + 1, Duration::from_secs(30)));
let committed_before = handle.last_committed(p).expect("first wave committed");
script.enqueue_global(WriteOutcome::fatal("table orders_local does not exist"));
let doomed = handle.push_many(p, ORDERS);
let (first_doomed, last_doomed) = (doomed[0], *doomed.last().expect("three offsets"));
let report = run
.wait_exit(Duration::from_secs(30))
.expect("the stall watchdog fails the pipeline")?;
assert_eq!(
handle.last_committed(p),
Some(committed_before),
"the watermark stalls rather than committing past unacknowledged data",
);
assert!(
committed_before <= first_doomed,
"the whole abandoned wave is uncommitted ({committed_before} vs {first_doomed})",
);
let ExitState::Failed(failure) = &report.state else {
panic!("a permanent stall must fail the pipeline, not idle forever");
};
assert_eq!(failure.component, "checkpoint");
assert!(failure.reason.contains("stalled"), "{}", failure.reason);
assert_ne!(report.exit_code(), 0, "the process must exit non-zero");
println!(
"\n3. fatal: watermark held at {committed_before} (records {first_doomed}..={last_doomed} \
unacknowledged), exit code {} — {}",
report.exit_code(),
failure.reason,
);
Ok(())
}
fn slow_sink() -> Result<(), Box<dyn std::error::Error>> {
const CONFIG: &str = r#"
pipeline: { name: sink-failures-slow, threads: 1, io_threads: 2 }
checkpoint: { interval: 100ms, drain_timeout: 10s }
backpressure: { max_inflight_bytes: 1KiB, high_ratio: 0.5, low_ratio: 0.25, min_pause: 50ms }
admin: { listen: none }
metrics: { exporter: none }
source: { memory: {} }
sink: { capture: { chunk: { target_bytes: 256B } } }
"#;
let (source, handle) = memory_source();
let (sink, script) = capture_sink(1, 1);
let sink = sink.with_pool_config(prompt_pool());
for _ in 0..6 {
script.enqueue_global(WriteOutcome::ok().after(Duration::from_millis(300)));
}
let runtime = assemble(CONFIG, sink, source)?;
let shutdown = runtime.shutdown_handle();
let run = PipelineRun::spawn(move || runtime.run());
let p = PartitionId(0);
handle.assign_lanes(&[(LaneId(0), p)]);
let backlog: Vec<String> = (0..60).map(|i| format!("order_placed:{i:08}")).collect();
let last = *handle.push_many(p, &backlog).last().expect("offsets");
wait_until(Duration::from_secs(30), "lanes paused under load", || {
!handle.paused_lanes().is_empty()
});
assert!(
handle.wait_committed(p, last + 1, Duration::from_secs(60)),
"the backlog drains once the sink catches up (last committed: {:?})",
handle.last_committed(p),
);
shutdown.trigger();
let report = run
.wait_exit(Duration::from_secs(30))
.expect("the pipeline drains after shutdown")?;
assert_eq!(report.exit_code(), 0, "a slow sink still drains clean");
let rows: usize = script
.writes()
.iter()
.map(|w| decode_rows(&w.payload).len())
.sum();
assert!(rows >= backlog.len(), "every record reached the sink");
println!("\n4. slow sink: lanes paused, then {rows} rows written and committed");
Ok(())
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
spate::telemetry::init(spate::telemetry::LogFormat::Pretty, "warn");
retryable_write()?;
failed_probe()?;
fatal_write()?;
slow_sink()?;
Ok(())
}
#[cfg(test)]
mod tests {
#[test]
fn runs_to_completion() {
super::main().expect("the example must run clean");
}
}