use std::{
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
mpsc,
},
thread,
time::{Duration, Instant},
};
use datum::{Channel, Keep, RestartSettings, Sink, Source, StreamError};
use datum_agent::{
Agent, AgentConfig, DesiredJobState, JobEventKind, JobMat, JobRestartPolicy, JobSpec, JobState,
JobStatus,
};
fn start_agent() -> datum_agent::AgentHandle {
Agent::start_with_config(AgentConfig {
poll_interval: Duration::from_millis(5),
event_buffer: 64,
..AgentConfig::default()
})
.expect("agent starts")
}
fn graceful_job(name: &str, source: Source<u64>) -> JobSpec {
JobSpec::new(name, move |context| {
let control = context.control();
Ok(source
.clone()
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}))
})
}
fn wait_for_status<F>(
registry: &datum_agent::JobRegistryHandle,
name: &str,
predicate: F,
) -> JobStatus
where
F: Fn(&JobStatus) -> bool,
{
let deadline = Instant::now() + Duration::from_secs(2);
loop {
let status = registry.status(name).expect("job status");
if predicate(&status) {
return status;
}
assert!(
Instant::now() < deadline,
"timed out waiting for status; latest={status:?}"
);
thread::sleep(Duration::from_millis(5));
}
}
#[test]
fn submit_start_status_and_events() {
let agent = start_agent();
let registry = agent.registry();
let events = registry
.events()
.run_with(Sink::queue())
.expect("event feed materializes");
let submitted = registry
.submit(graceful_job("finite", Source::from_iterable([1_u64, 2, 3])))
.expect("job submits");
assert_eq!(submitted.state, JobState::Submitted);
let first_event = events.pull().expect("event pull").expect("event");
assert_eq!(first_event.kind, JobEventKind::Submitted);
assert_eq!(first_event.name, "finite");
let started = registry.start("finite").expect("job starts");
assert_eq!(started.state, JobState::Running);
assert_eq!(started.starts_total, 1);
let mut saw_started = false;
for _ in 0..4 {
let event = events.pull().expect("event pull").expect("event");
if event.kind == JobEventKind::Started {
saw_started = true;
break;
}
}
assert!(saw_started, "started event was not emitted");
let completed = wait_for_status(registry, "finite", |status| {
status.state == JobState::Completed
});
assert_eq!(completed.desired_state, DesiredJobState::Running);
assert_eq!(completed.starts_total, 1);
registry.shutdown().expect("registry shuts down");
}
#[test]
fn drain_completes_in_flight_then_stops() {
let agent = start_agent();
let registry = agent.registry();
let channel = Channel::bounded(4);
let producer = channel.clone();
let (started_tx, started_rx) = mpsc::channel();
let (done_tx, done_rx) = mpsc::channel();
let spec = JobSpec::new("drain", move |context| {
let control = context.control();
let started_tx = started_tx.clone();
let done_tx = done_tx.clone();
Ok(channel
.source()
.via_mat(context.drain_flow(), Keep::right)
.to_mat(
Sink::foreach(move |item| {
let _ = started_tx.send(item);
thread::sleep(Duration::from_millis(40));
let _ = done_tx.send(item);
}),
move |_switch, completion| JobMat::new(completion, control.clone()),
))
});
registry.submit(spec).expect("job submits");
registry.start("drain").expect("job starts");
producer.try_send(7).expect("send item");
assert_eq!(
started_rx
.recv_timeout(Duration::from_secs(1))
.expect("item enters sink"),
7
);
let draining = registry.drain("drain").expect("drain starts");
assert_eq!(draining.state, JobState::Draining);
assert_eq!(draining.desired_state, DesiredJobState::Draining);
assert_eq!(
done_rx
.recv_timeout(Duration::from_secs(1))
.expect("in-flight item drains"),
7
);
let drained = wait_for_status(registry, "drain", |status| {
status.state == JobState::Drained
});
assert_eq!(drained.desired_state, DesiredJobState::Stopped);
registry.shutdown().expect("registry shuts down");
}
#[test]
fn crash_restarts_with_backoff_policy() {
let agent = start_agent();
let registry = agent.registry();
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_for_factory = Arc::clone(&attempts);
let settings = RestartSettings::new(Duration::from_millis(20), Duration::from_millis(20), 0.0)
.with_max_restarts(1, Duration::from_secs(1));
let spec = JobSpec::new("restart", move |context| {
let control = context.control();
let attempt = attempts_for_factory.fetch_add(1, Ordering::SeqCst);
let source = if attempt == 0 {
Source::failed(StreamError::Failed("boom".to_owned()))
} else {
Source::<u64>::empty()
};
Ok(source
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}))
})
.with_restart_policy(JobRestartPolicy::on_failure(settings));
registry.submit(spec).expect("job submits");
registry.start("restart").expect("job starts");
let restarted = wait_for_status(registry, "restart", |status| {
status.generation >= 2 && status.restarts_total == 1
});
assert_eq!(restarted.starts_total, 2);
assert!(attempts.load(Ordering::SeqCst) >= 2);
registry.shutdown().expect("registry shuts down");
}
#[test]
fn stop_during_drain_stops_job() {
let agent = start_agent();
let registry = agent.registry();
let channel = Channel::bounded(4);
let producer = channel.clone();
let (started_tx, started_rx) = mpsc::channel();
let spec = JobSpec::new("stop-drain", move |context| {
let control = context.control();
let started_tx = started_tx.clone();
Ok(channel
.source()
.via_mat(context.drain_flow(), Keep::right)
.to_mat(
Sink::foreach(move |item| {
let _ = started_tx.send(item);
thread::sleep(Duration::from_millis(100));
}),
move |_switch, completion| JobMat::new(completion, control.clone()),
))
});
registry.submit(spec).expect("job submits");
registry.start("stop-drain").expect("job starts");
producer.try_send(1).expect("send item");
started_rx
.recv_timeout(Duration::from_secs(1))
.expect("item enters sink");
registry.drain("stop-drain").expect("drain starts");
let stopped = registry.stop("stop-drain").expect("stop succeeds");
assert_eq!(stopped.state, JobState::Stopped);
assert_eq!(stopped.desired_state, DesiredJobState::Stopped);
registry.shutdown().expect("registry shuts down");
}
#[test]
fn registry_survives_job_factory_panics() {
let agent = start_agent();
let registry = agent.registry();
let bad = JobSpec::new("panic", |_context| -> datum::StreamResult<_> {
panic!("factory panic")
});
registry.submit(bad).expect("job submits");
assert!(registry.start("panic").is_err());
let failed = registry.status("panic").expect("job status");
assert_eq!(failed.state, JobState::Failed);
registry
.submit(graceful_job("healthy", Source::<u64>::empty()))
.expect("healthy job submits");
let healthy = registry.start("healthy").expect("healthy job starts");
assert_eq!(healthy.state, JobState::Running);
registry.shutdown().expect("registry shuts down");
}