use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use flyby::api::{
CountingCollector, Decoder, DefaultSchemaId, Error, ErrorKind, Lifecycle, Message, Metadata,
MetricKey, Pipeline, Result, Sink, SinkId, Timestamp,
};
use flyby::pipeline::{
CallbackPlacement, DropAllPlacement, FixedPlacement, HashPlacement, IdentityPreProcessor,
RawBatchSource, RoundRobinPlacement, SimplePipeline, schema_hash_placement,
};
use flyby::runtime::{
AffinityRequest, BackpressureStrategy, CpuAffinityPolicy, Runtime, RuntimeConfig, RuntimeEvent,
RuntimeMetricKey, RuntimePhase, Scheduler, SchedulerKind, SingleThreadScheduler,
WorkerPoolScheduler, apply_affinity, build_scheduler, emit_lifecycle, run_pipeline,
run_with_worker_factory,
};
#[derive(Debug, Clone, PartialEq, Eq)]
struct Tick {
seq: u64,
}
impl Message for Tick {
type Schema = DefaultSchemaId;
fn schema_id(&self) -> Self::Schema {
DefaultSchemaId(1)
}
fn timestamp(&self) -> Timestamp {
Timestamp::from_nanos(self.seq)
}
fn metadata(&self) -> Metadata {
Metadata {
sequence: self.seq,
..Metadata::default()
}
}
}
struct TickDecoder;
impl Decoder for TickDecoder {
type Output = Tick;
fn decode(&mut self, bytes: &[u8]) -> Result<Option<Tick>> {
if bytes.len() < 8 {
return Ok(None);
}
let mut arr = [0u8; 8];
arr.copy_from_slice(&bytes[..8]);
Ok(Some(Tick {
seq: u64::from_le_bytes(arr),
}))
}
}
struct VecSource {
frames: Vec<Vec<u8>>,
idx: usize,
inited: bool,
}
impl VecSource {
fn new(seqs: &[u64]) -> Self {
Self {
frames: seqs.iter().map(|s| s.to_le_bytes().to_vec()).collect(),
idx: 0,
inited: false,
}
}
}
impl Lifecycle for VecSource {
fn init(&mut self) -> Result<()> {
self.inited = true;
Ok(())
}
fn shutdown(&mut self) -> Result<()> {
self.inited = false;
Ok(())
}
}
impl RawBatchSource for VecSource {
fn poll_frames(&mut self, out: &mut Vec<Vec<u8>>) -> Result<usize> {
if self.idx >= self.frames.len() {
return Ok(0);
}
let n = (self.frames.len() - self.idx).min(4);
for _ in 0..n {
out.push(self.frames[self.idx].clone());
self.idx += 1;
}
Ok(n)
}
fn is_exhausted(&self) -> bool {
self.idx >= self.frames.len()
}
}
struct CountingSink {
written: Arc<Mutex<Vec<u64>>>,
fail_after: Option<usize>,
}
impl CountingSink {
fn new(written: Arc<Mutex<Vec<u64>>>) -> Self {
Self {
written,
fail_after: None,
}
}
fn with_bp_after(written: Arc<Mutex<Vec<u64>>>, n: usize) -> Self {
Self {
written,
fail_after: Some(n),
}
}
}
impl Lifecycle for CountingSink {}
impl Sink for CountingSink {
type Message = Tick;
fn write(&mut self, message: &Tick) -> Result<()> {
let mut g = self.written.lock().unwrap();
if let Some(limit) = self.fail_after
&& g.len() >= limit
{
return Err(Error::new(ErrorKind::BackPressure, "full"));
}
g.push(message.seq);
Ok(())
}
}
struct RecoveringSink {
written: Arc<Mutex<Vec<u64>>>,
fails_left: usize,
}
impl Lifecycle for RecoveringSink {}
impl Sink for RecoveringSink {
type Message = Tick;
fn write(&mut self, message: &Tick) -> Result<()> {
if self.fails_left > 0 {
self.fails_left -= 1;
return Err(Error::back_pressure("retry"));
}
self.written.lock().unwrap().push(message.seq);
Ok(())
}
}
#[test]
fn runtime_lifecycle_phases() {
let written = Arc::new(Mutex::new(Vec::new()));
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written))),
)
.unwrap();
let mut rt = Runtime::build(pipe, RuntimeConfig::default()).unwrap();
assert_eq!(rt.phase(), RuntimePhase::Built);
rt.validate().unwrap();
assert_eq!(rt.phase(), RuntimePhase::Validated);
let stats = rt.run().unwrap();
assert_eq!(rt.phase(), RuntimePhase::CleanedUp);
assert!(stats.steps > 0);
assert_eq!(written.lock().unwrap().len(), 3);
}
#[test]
fn drop_newest_backpressure() {
let written = Arc::new(Mutex::new(Vec::new()));
let cfg = RuntimeConfig::default()
.with_backpressure(BackpressureStrategy::DropNewest)
.with_batch_size(8);
let mut pipe = SimplePipeline::new(
VecSource::new(&[10, 20, 30]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_runtime(cfg.clone());
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::with_bp_after(Arc::clone(&written), 1)),
)
.unwrap();
run_pipeline(&mut pipe, &cfg).unwrap();
assert_eq!(written.lock().unwrap().len(), 1);
assert!(pipe.messages_dropped() >= 2);
assert!(pipe.backpressure_events() >= 2);
}
#[test]
fn round_robin_placement() {
let a = Arc::new(Mutex::new(Vec::new()));
let b = Arc::new(Mutex::new(Vec::new()));
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3, 4]),
TickDecoder,
IdentityPreProcessor::default(),
RoundRobinPlacement::new(vec![SinkId::new(1), SinkId::new(2)]).unwrap(),
);
pipe.register_sink(SinkId::new(1), Box::new(CountingSink::new(Arc::clone(&a))))
.unwrap();
pipe.register_sink(SinkId::new(2), Box::new(CountingSink::new(Arc::clone(&b))))
.unwrap();
run_pipeline(&mut pipe, &RuntimeConfig::default()).unwrap();
assert_eq!(a.lock().unwrap().len(), 2);
assert_eq!(b.lock().unwrap().len(), 2);
}
#[test]
fn callback_and_schema_hash_placement() {
let written = Arc::new(Mutex::new(Vec::new()));
let place = CallbackPlacement::new(|m: &Tick| {
if m.seq.is_multiple_of(2) {
Ok(SinkId::new(1))
} else {
Ok(SinkId::NONE)
}
});
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3, 4]),
TickDecoder,
IdentityPreProcessor::default(),
place,
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written))),
)
.unwrap();
run_pipeline(&mut pipe, &RuntimeConfig::default()).unwrap();
assert_eq!(*written.lock().unwrap(), vec![2, 4]);
let _ = schema_hash_placement::<Tick>(vec![SinkId::new(1)]).unwrap();
}
#[test]
fn single_thread_scheduler_runs() {
let written = Arc::new(Mutex::new(Vec::new()));
let mut pipe = SimplePipeline::new(
VecSource::new(&[7, 8]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written))),
)
.unwrap();
pipe.init().unwrap();
let mut sched = SingleThreadScheduler::new(RuntimeConfig::default());
sched.run(&mut pipe).unwrap();
pipe.shutdown().unwrap();
assert_eq!(written.lock().unwrap().len(), 2);
}
#[test]
fn worker_pool_factory() {
let written = Arc::new(Mutex::new(Vec::new()));
let written2 = Arc::clone(&written);
let cfg = RuntimeConfig::default()
.with_scheduler(SchedulerKind::WorkerPool)
.with_workers(2);
let pool = WorkerPoolScheduler::new(cfg).unwrap();
pool.run_factory(move |worker| {
let seqs = if worker == 0 {
vec![1u64, 2]
} else {
vec![3, 4]
};
let mut pipe = SimplePipeline::new(
VecSource::new(&seqs),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written2))),
)?;
Ok(pipe)
})
.unwrap();
assert_eq!(written.lock().unwrap().len(), 4);
}
#[test]
fn affinity_noop_and_strict() {
let req = AffinityRequest::for_worker(CpuAffinityPolicy::None, 0);
apply_affinity(&req, true).unwrap();
let req = AffinityRequest::for_worker(CpuAffinityPolicy::PinWorkers, 0);
assert!(apply_affinity(&req, true).is_err());
apply_affinity(&req, false).unwrap();
}
#[test]
fn toml_scheduler_worker_pool() {
let cfg = RuntimeConfig::from_toml_str(
r#"
[runtime]
workers = 2
batch_size = 64
backpressure = "drop_newest"
scheduler = "worker_pool"
metrics = false
"#,
)
.unwrap();
assert_eq!(cfg.scheduler, SchedulerKind::WorkerPool);
assert_eq!(cfg.backpressure, BackpressureStrategy::DropNewest);
}
#[test]
fn overflow_routes_to_overflow_sink() {
let primary = Arc::new(Mutex::new(Vec::new()));
let overflow = Arc::new(Mutex::new(Vec::new()));
let cfg = RuntimeConfig::default()
.with_backpressure(BackpressureStrategy::Overflow)
.with_batch_size(8);
let mut cfg = cfg;
cfg.overflow_sink = Some(2);
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_runtime(cfg.clone());
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::with_bp_after(Arc::clone(&primary), 0)),
)
.unwrap();
pipe.register_sink(
SinkId::new(2),
Box::new(CountingSink::new(Arc::clone(&overflow))),
)
.unwrap();
run_pipeline(&mut pipe, &cfg).unwrap();
assert!(primary.lock().unwrap().is_empty());
assert_eq!(overflow.lock().unwrap().len(), 3);
}
#[test]
fn overflow_without_sink_drops() {
let written = Arc::new(Mutex::new(Vec::new()));
let mut cfg = RuntimeConfig::default().with_backpressure(BackpressureStrategy::Overflow);
cfg.overflow_sink = None;
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_runtime(cfg.clone());
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::with_bp_after(Arc::clone(&written), 0)),
)
.unwrap();
run_pipeline(&mut pipe, &cfg).unwrap();
assert!(written.lock().unwrap().is_empty());
assert!(pipe.messages_dropped() >= 2);
}
#[test]
fn drop_oldest_matches_drop_newest_for_single_message() {
let written = Arc::new(Mutex::new(Vec::new()));
let cfg = RuntimeConfig::default().with_backpressure(BackpressureStrategy::DropOldest);
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_runtime(cfg.clone());
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::with_bp_after(Arc::clone(&written), 1)),
)
.unwrap();
run_pipeline(&mut pipe, &cfg).unwrap();
assert_eq!(written.lock().unwrap().len(), 1);
assert!(pipe.messages_dropped() >= 2);
}
#[test]
fn block_retries_then_succeeds() {
let written = Arc::new(Mutex::new(Vec::new()));
let cfg = RuntimeConfig::default()
.with_backpressure(BackpressureStrategy::Block)
.with_batch_size(4);
let mut cfg = cfg;
cfg.backpressure_retries = Some(8);
cfg.backpressure_yield_ms = 0;
let mut pipe = SimplePipeline::new(
VecSource::new(&[42]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_runtime(cfg.clone());
pipe.register_sink(
SinkId::new(1),
Box::new(RecoveringSink {
written: Arc::clone(&written),
fails_left: 3,
}),
)
.unwrap();
run_pipeline(&mut pipe, &cfg).unwrap();
assert_eq!(*written.lock().unwrap(), vec![42]);
assert!(pipe.backpressure_events() >= 3);
}
#[test]
fn spin_exhausts_retry_budget() {
let written = Arc::new(Mutex::new(Vec::new()));
let mut cfg = RuntimeConfig::default().with_backpressure(BackpressureStrategy::Spin);
cfg.backpressure_retries = Some(2);
let mut pipe = SimplePipeline::new(
VecSource::new(&[1]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_runtime(cfg);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::with_bp_after(Arc::clone(&written), 0)),
)
.unwrap();
pipe.init().unwrap();
let mut saw_bp = false;
for _ in 0..8 {
match pipe.step_outcome().unwrap() {
flyby::api::StepOutcome::BackPressured => {
saw_bp = true;
break;
}
flyby::api::StepOutcome::Exhausted => break,
_ => {}
}
}
pipe.shutdown().unwrap();
assert!(saw_bp);
assert!(written.lock().unwrap().is_empty());
assert!(pipe.backpressure_events() >= 2);
}
#[test]
fn hash_and_drop_all_placement() {
let a = Arc::new(Mutex::new(Vec::new()));
let b = Arc::new(Mutex::new(Vec::new()));
let place = HashPlacement::new(vec![SinkId::new(1), SinkId::new(2)], |m: &Tick| m.seq).unwrap();
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3, 4]),
TickDecoder,
IdentityPreProcessor::default(),
place,
);
pipe.register_sink(SinkId::new(1), Box::new(CountingSink::new(Arc::clone(&a))))
.unwrap();
pipe.register_sink(SinkId::new(2), Box::new(CountingSink::new(Arc::clone(&b))))
.unwrap();
run_pipeline(&mut pipe, &RuntimeConfig::default()).unwrap();
assert_eq!(a.lock().unwrap().len() + b.lock().unwrap().len(), 4);
let mut drop_pipe = SimplePipeline::new(
VecSource::new(&[9, 10]),
TickDecoder,
IdentityPreProcessor::default(),
DropAllPlacement::default(),
);
drop_pipe
.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::new(Mutex::new(Vec::new())))),
)
.unwrap();
run_pipeline(&mut drop_pipe, &RuntimeConfig::default()).unwrap();
assert_eq!(drop_pipe.messages_out(), 0);
assert!(FixedPlacement::<Tick>::new(SinkId::NONE).is_err());
assert!(HashPlacement::<Tick, _>::new(vec![], |_| 0).is_err());
assert!(RoundRobinPlacement::<Tick>::new(vec![SinkId::NONE]).is_err());
}
#[test]
fn runtime_cooperative_shutdown() {
let written = Arc::new(Mutex::new(Vec::new()));
let seqs: Vec<u64> = (1..=200).collect();
let mut pipe = SimplePipeline::new(
VecSource::new(&seqs),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written))),
)
.unwrap();
let mut rt = Runtime::build(pipe, RuntimeConfig::default().with_metrics(true)).unwrap();
let handle = rt.shutdown_handle();
handle.store(true, Ordering::SeqCst);
rt.request_shutdown();
let stats = rt.run().unwrap();
assert_eq!(rt.phase(), RuntimePhase::CleanedUp);
assert!(stats.steps < 200 || written.lock().unwrap().len() < 200);
}
#[test]
fn runtime_config_builders_and_validation() {
let cfg = RuntimeConfig::default()
.with_workers(0) .with_batch_size(0)
.with_scheduler(SchedulerKind::SingleThread)
.with_metrics(false)
.with_backpressure(BackpressureStrategy::AdaptiveBatching);
assert_eq!(cfg.workers, 1);
assert_eq!(cfg.batch_size, 1);
assert_eq!(cfg.idle_sleep(), None);
assert_eq!(cfg.backpressure_yield(), Duration::ZERO);
let bad = RuntimeConfig {
workers: 0,
..RuntimeConfig::default()
};
assert!(bad.validate().is_err());
let dir = std::env::temp_dir().join(format!("flyby-runtime-cfg-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("runtime.toml");
std::fs::write(
&path,
r#"
[runtime]
workers = 3
batch_size = 32
idle_sleep_ms = 1
backpressure_yield_ms = 1
"#,
)
.unwrap();
let loaded = RuntimeConfig::from_toml_path(&path).unwrap();
assert_eq!(loaded.workers, 3);
assert_eq!(loaded.idle_sleep(), Some(Duration::from_millis(1)));
assert!(RuntimeConfig::from_toml_path(dir.join("missing.toml")).is_err());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn backpressure_and_metric_key_labels() {
for (bp, name) in [
(BackpressureStrategy::Block, "block"),
(BackpressureStrategy::Spin, "spin"),
(BackpressureStrategy::DropNewest, "drop_newest"),
(BackpressureStrategy::DropOldest, "drop_oldest"),
(BackpressureStrategy::Overflow, "overflow"),
(BackpressureStrategy::AdaptiveBatching, "adaptive_batching"),
] {
assert_eq!(bp.as_str(), name);
}
let keys = [
(RuntimeMetricKey::Steps, "runtime.steps"),
(RuntimeMetricKey::MessagesOut, "runtime.messages_out"),
(
RuntimeMetricKey::MessagesDropped,
"runtime.messages_dropped",
),
(
RuntimeMetricKey::BackpressureEvents,
"runtime.backpressure_events",
),
(RuntimeMetricKey::IdleSkips, "runtime.idle_skips"),
(RuntimeMetricKey::StepDurationNs, "runtime.step_duration_ns"),
(RuntimeMetricKey::Phase, "runtime.phase"),
(RuntimeMetricKey::Workers, "runtime.workers"),
];
for (key, name) in keys {
assert_eq!(key.name(), name);
}
for (phase, n) in [
(RuntimePhase::Built, 0),
(RuntimePhase::Validated, 1),
(RuntimePhase::Initialized, 2),
(RuntimePhase::Running, 3),
(RuntimePhase::Draining, 4),
(RuntimePhase::Shutdown, 5),
(RuntimePhase::CleanedUp, 6),
] {
assert_eq!(phase.as_u8(), n);
}
let metrics = CountingCollector::new();
emit_lifecycle(&metrics, false, &RuntimeEvent::Built);
assert_eq!(metrics.calls.load(Ordering::Relaxed), 0);
for event in [
RuntimeEvent::Built,
RuntimeEvent::Validated,
RuntimeEvent::Initialized,
RuntimeEvent::Started,
RuntimeEvent::Draining,
RuntimeEvent::Shutdown,
RuntimeEvent::Cleanup,
RuntimeEvent::Step { n: 1 },
] {
emit_lifecycle(&metrics, true, &event);
}
assert!(metrics.calls.load(Ordering::Relaxed) >= 8);
let sched = build_scheduler(RuntimeConfig::default());
assert!(!sched.is_shutdown_requested());
sched.request_shutdown();
assert!(sched.is_shutdown_requested());
}
#[test]
fn run_with_worker_factory_helper() {
let written = Arc::new(Mutex::new(Vec::new()));
let written2 = Arc::clone(&written);
let cfg = RuntimeConfig::default()
.with_scheduler(SchedulerKind::WorkerPool)
.with_workers(1);
run_with_worker_factory(cfg, move |_| {
let mut pipe = SimplePipeline::new(
VecSource::new(&[5, 6]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written2))),
)?;
Ok(pipe)
})
.unwrap();
assert_eq!(written.lock().unwrap().len(), 2);
}
#[test]
fn pipeline_run_and_accessors() {
let written = Arc::new(Mutex::new(Vec::new()));
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2]),
TickDecoder,
IdentityPreProcessor::default(),
FixedPlacement::new(SinkId::new(1)).unwrap(),
)
.with_metrics(CountingCollector::new());
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written))),
)
.unwrap();
assert!(!pipe.source().is_exhausted());
let _ = pipe.source_mut();
pipe.init().unwrap();
pipe.run().unwrap();
assert_eq!(written.lock().unwrap().len(), 2);
assert_eq!(pipe.messages_out(), 2);
assert_eq!(pipe.runtime_config().batch_size, 512);
}
#[test]
fn schema_hash_routes_consistently() {
let written = Arc::new(Mutex::new(Vec::new()));
let place = schema_hash_placement::<Tick>(vec![SinkId::new(1)]).unwrap();
let mut pipe = SimplePipeline::new(
VecSource::new(&[1, 2, 3]),
TickDecoder,
IdentityPreProcessor::default(),
place,
);
pipe.register_sink(
SinkId::new(1),
Box::new(CountingSink::new(Arc::clone(&written))),
)
.unwrap();
run_pipeline(&mut pipe, &RuntimeConfig::default()).unwrap();
assert_eq!(written.lock().unwrap().len(), 3);
}