use core::future::Future;
use std::path::Path;
use std::time::Duration;
use shep_client::dogs::{Stop, parse_section};
use shep_core::protocol::BusEvent;
use tokio::time::{Instant, MissedTickBehavior};
use super::clock::Clock;
use super::config::{BarkConfig, rules_for};
use super::delivery::{Delivery, spawn_firings};
use super::rules::Rules;
use super::source::{ConfigSource, EventSource, FlockSource, Resubscribe};
use crate::exit::ExitCode;
const LAG_POLL_MIN_GAP: Duration = Duration::from_secs(1);
#[derive(Debug, Default)]
struct LagPacing {
last_poll: Option<Instant>,
lag_pending: bool,
}
impl LagPacing {
fn polled(&mut self) {
self.last_poll = Some(Instant::now());
self.lag_pending = false;
}
fn lag_polls_now(&mut self) -> bool {
match self.last_poll {
Some(last) if last.elapsed() < LAG_POLL_MIN_GAP => {
self.lag_pending = true;
false
}
_ => true,
}
}
fn deferred_poll_at(&self) -> Option<Instant> {
self.last_poll
.filter(|_| self.lag_pending)
.map(|last| last + LAG_POLL_MIN_GAP)
}
}
pub fn run_loop<E: EventSource, F: FlockSource, C: ConfigSource>(
events: E,
flock: F,
rules: Rules,
config: &BarkConfig,
barks_path: &Path,
config_source: C,
stop: Stop,
) -> impl Future<Output = ExitCode> + Send + use<E, F, C> {
let mut delivery = Delivery::new(config, barks_path);
let mut poll_period = config.poll.as_duration();
async move {
let mut events = events;
let mut rules = rules;
let mut stop = stop;
let mut clock = Clock::system();
let mut poll_interval = poll_timer(poll_period);
let mut pacing = LagPacing::default();
if !rules.log_topics().is_empty() {
match flock.flock().await {
Ok(snapshot) => rules.learn(&snapshot),
Err(err) => eprintln!("shep dog bark: could not read the flock's names: {err}"),
}
}
loop {
tokio::select! {
() = stop.wait() => break ExitCode::Success,
next = events.next() => {
match next {
None => match events.resubscribe().await {
Ok(()) => {
reconcile(&flock, &mut rules, &mut clock, &delivery).await;
pacing.polled();
}
Err(failed) => break exit_for(&failed),
},
Some(Ok(BusEvent::DogConfigChanged { .. })) => {
if let Some((next, next_rules)) = reloaded_config(&config_source).await {
delivery.reconfigure(&next);
let mut next_rules = next_rules;
next_rules.take_names_from(&mut rules);
rules = next_rules;
if next.poll.as_duration() != poll_period {
poll_period = next.poll.as_duration();
poll_interval = poll_timer(poll_period);
}
eprintln!("shep dog bark: reloaded [bark] from dogs.toml");
let wanted = wanted_topics(&events, &rules);
if wanted != events.topics() {
if let Err(failed) = events.set_topics(wanted).await {
break exit_for(&failed);
}
reconcile(&flock, &mut rules, &mut clock, &delivery).await;
pacing.polled();
}
}
}
Some(Ok(event)) => {
let firings = rules.on_event(&event, clock.now_ms());
spawn_firings(firings, &delivery);
}
Some(Err(_dropped)) => {
if pacing.lag_polls_now() {
reconcile(&flock, &mut rules, &mut clock, &delivery).await;
pacing.polled();
}
}
}
}
() = deferred_poll(pacing.deferred_poll_at()) => {
reconcile(&flock, &mut rules, &mut clock, &delivery).await;
pacing.polled();
}
_ = poll_interval.tick() => {
reconcile(&flock, &mut rules, &mut clock, &delivery).await;
pacing.polled();
}
}
}
}
}
fn wanted_topics<E: EventSource>(events: &E, rules: &Rules) -> Vec<String> {
let mut topics: Vec<String> = events
.topics()
.iter()
.filter(|topic| !topic.starts_with("log."))
.cloned()
.collect();
topics.extend(rules.log_topics());
topics
}
fn exit_for(failed: &Resubscribe) -> ExitCode {
eprintln!("shep dog bark: {failed}");
match failed {
Resubscribe::Lost(lost) => ExitCode::from(lost),
Resubscribe::Request(err) => ExitCode::from(err),
}
}
pub(super) async fn reconcile<F: FlockSource>(
flock: &F,
rules: &mut Rules,
clock: &mut Clock,
delivery: &Delivery,
) {
match flock.flock().await {
Ok(snapshot) => {
let firings = rules.on_poll(&snapshot, clock.now_ms());
spawn_firings(firings, delivery);
}
Err(err) => eprintln!("shep dog bark: reconciliation poll failed: {err}"),
}
}
pub(super) async fn reloaded_config<C: ConfigSource>(source: &C) -> Option<(BarkConfig, Rules)> {
let section = match source.section().await {
Ok(section) => section,
Err(err) => {
eprintln!("shep dog bark: could not re-read [bark] from the shepherd: {err}");
return None;
}
};
let config = match parse_section::<BarkConfig>("bark", §ion) {
Ok(config) => config,
Err(err) => {
eprintln!("shep dog bark: {err}; see `shep dogs`");
return None;
}
};
match rules_for(&config) {
Ok(rules) => Some((config, rules)),
Err(err) => {
eprintln!("shep dog bark: keeping the running rules; the new ones are refused: {err}");
None
}
}
}
async fn deferred_poll(due: Option<Instant>) {
match due {
Some(due) => tokio::time::sleep_until(due).await,
None => core::future::pending().await,
}
}
fn poll_timer(period: Duration) -> tokio::time::Interval {
let mut interval = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
interval
}
#[cfg(test)]
mod log_match_tests;
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::time::Duration;
use shep_client::LinkLost;
use shep_core::barks;
use shep_core::protocol::{ProcessEventKind, RpcErrorCode};
use shep_core::values::UpDuration;
use tokio::sync::broadcast;
use super::*;
use crate::dog::bark::testing::*;
#[tokio::test]
async fn a_dropped_frame_makes_bark_poll_and_catch_up() {
let (tx, mut rx) = tokio::sync::broadcast::channel(4);
for i in 0..64 {
tx.send(log_event(i)).unwrap();
}
tx.send(errored_event("web")).unwrap();
assert!(
matches!(rx.recv().await, Err(broadcast::error::RecvError::Lagged(n)) if n > 0),
"the fixture must actually overflow the channel"
);
let (tx2, rx2) = tokio::sync::broadcast::channel(4);
for i in 0..64 {
tx2.send(log_event(i)).unwrap();
}
tx2.send(errored_event("web")).unwrap();
let (addr, captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let flock = ScriptedFlock::answering(vec![errored_info("web", 16)]);
let loop_handle = tokio::spawn(run_loop(
rx2,
flock.clone(),
gave_up_rules(),
&config_with_sink(addr),
&barks_path,
ScriptedConfig::answering(String::new()),
Stop::never(),
));
let req = tokio::time::timeout(Duration::from_secs(5), captured)
.await
.expect("a dropped frame must produce a delivered bark")
.unwrap();
assert!(String::from_utf8_lossy(&req.body).contains("web"));
let recorded = tokio::time::timeout(Duration::from_secs(5), async {
loop {
let records = shep_core::barks::read(&barks_path).unwrap();
if !records.is_empty() {
break records;
}
tokio::task::yield_now().await;
}
})
.await
.expect("the delivered bark must be recorded promptly after delivery");
assert_eq!(recorded.len(), 1);
assert_eq!(recorded[0].subject, "web");
assert_eq!(recorded[0].sinks[0].error, None);
assert_eq!(
flock.calls(),
1,
"the poll ran because of the lag, not because an interval elapsed: \
the interval is 60s and this test is milliseconds old"
);
loop_handle.abort();
}
#[tokio::test]
async fn a_handover_re_subscribes_instead_of_ending_the_dog() {
let (first_tx, first_rx) = tokio::sync::broadcast::channel(8);
let (second_tx, second_rx) = tokio::sync::broadcast::channel(8);
let (source, resubscribes) = HandoverSource::across(
vec![first_rx, second_rx],
LinkLost::Budget {
waited: Duration::ZERO,
},
);
let (addr, captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let flock = ScriptedFlock::answering(Vec::new());
let loop_handle = tokio::spawn(run_loop(
source,
flock.clone(),
gave_up_rules(),
&config_with_sink(addr),
&barks_path,
ScriptedConfig::answering(String::new()),
Stop::never(),
));
drop(first_tx);
let delivered = tokio::time::timeout(Duration::from_secs(5), async {
loop {
let _ = second_tx.send(errored_event("web"));
tokio::time::sleep(Duration::from_millis(20)).await;
if loop_handle.is_finished() {
panic!("the loop ended on a handover instead of re-subscribing");
}
}
});
let captured = tokio::time::timeout(Duration::from_secs(5), captured);
let request = tokio::select! {
() = async { delivered.await.ok(); } => panic!(
"no bark was delivered within 5s of the handover, so the \
second generation never reached the loop"
),
request = captured => request.expect("a bark must be delivered after the handover"),
};
assert!(String::from_utf8_lossy(&request.unwrap().body).contains("web"));
assert_eq!(
resubscribes.load(std::sync::atomic::Ordering::SeqCst),
1,
"the handover must have been answered by exactly one re-subscribe"
);
assert!(
flock.calls() >= 1,
"a re-subscribe must reconcile: frames sent while there was no \
subscription are gone, and only the shepherd knows the flock now"
);
assert!(
!loop_handle.is_finished(),
"the dog must still be running after a handover"
);
loop_handle.abort();
}
#[tokio::test]
async fn a_slow_sink_never_stalls_the_loop() {
let (slow_addr, slow_connected) = slow_sink().await;
let (fast_addr, fast_captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let mut sinks = BTreeMap::new();
sinks.insert(
"slow".to_owned(),
json_sink(format!("http://{slow_addr}/hook")),
);
sinks.insert(
"fast".to_owned(),
json_sink(format!("http://{fast_addr}/hook")),
);
let rules = Rules::new(
vec![
super::super::rules::Rule {
when: super::super::rules::Trigger::GaveUp {},
sinks: vec!["slow".to_owned()],
debounce: UpDuration::from_millis(0),
},
super::super::rules::Rule {
when: super::super::rules::Trigger::Event {
kinds: vec!["online".to_owned()],
},
sinks: vec!["fast".to_owned()],
debounce: UpDuration::from_millis(0),
},
],
&sinks,
)
.unwrap();
let config = BarkConfig {
sinks,
rules: Vec::new(),
poll: UpDuration::from_millis(60_000),
history_bytes: barks::DEFAULT_MAX_BYTES,
sink_timeout: UpDuration::from_millis(600_000),
};
let (tx, rx) = tokio::sync::broadcast::channel(8);
tx.send(errored_event("web")).unwrap();
tx.send(process_event("api", ProcessEventKind::Online))
.unwrap();
let flock = ScriptedFlock::answering(Vec::new());
let loop_handle = tokio::spawn(run_loop(
rx,
flock,
rules,
&config,
&barks_path,
ScriptedConfig::answering(String::new()),
Stop::never(),
));
tokio::time::timeout(Duration::from_secs(30), slow_connected)
.await
.expect("the slow sink must be reached at all, or this test proves nothing")
.unwrap();
let req = tokio::time::timeout(Duration::from_secs(30), fast_captured)
.await
.expect(
"the fast sink must be reached while a slow sink is still in \
flight; a slow sink must not stall the loop",
)
.unwrap();
assert_eq!(req.method, "POST");
loop_handle.abort();
}
#[tokio::test]
async fn a_dog_rides_out_three_shepherds_in_a_row() {
let (first_tx, first_rx) = tokio::sync::broadcast::channel(8);
let (second_tx, second_rx) = tokio::sync::broadcast::channel(8);
let (third_tx, third_rx) = tokio::sync::broadcast::channel(8);
let (source, resubscribes) = HandoverSource::across(
vec![first_rx, second_rx, third_rx],
LinkLost::Budget {
waited: Duration::ZERO,
},
);
let (addr, captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let flock = ScriptedFlock::answering(Vec::new());
let loop_handle = tokio::spawn(run_loop(
source,
flock.clone(),
gave_up_rules(),
&config_with_sink(addr),
&barks_path,
ScriptedConfig::answering(String::new()),
Stop::never(),
));
drop(first_tx);
let two = tokio::time::timeout(Duration::from_secs(5), async {
while resubscribes.load(std::sync::atomic::Ordering::SeqCst) < 1 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
drop(second_tx);
while resubscribes.load(std::sync::atomic::Ordering::SeqCst) < 2 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await;
assert!(two.is_ok(), "both handovers must be answered within 5s");
let delivered = tokio::time::timeout(Duration::from_secs(5), async {
loop {
let _ = third_tx.send(errored_event("web"));
tokio::time::sleep(Duration::from_millis(20)).await;
assert!(
!loop_handle.is_finished(),
"the dog ended on one of the two handovers"
);
}
});
let captured = tokio::time::timeout(Duration::from_secs(5), captured);
let request = tokio::select! {
() = async { delivered.await.ok(); } => panic!(
"no bark was delivered on the third generation within 5s"
),
request = captured => request.expect("a bark must be delivered after two handovers"),
};
assert!(String::from_utf8_lossy(&request.unwrap().body).contains("web"));
assert_eq!(
resubscribes.load(std::sync::atomic::Ordering::SeqCst),
2,
"two handovers, two re-subscribes, no more"
);
loop_handle.abort();
}
struct LagSource(tokio::sync::mpsc::UnboundedReceiver<u64>);
impl EventSource for LagSource {
async fn next(&mut self) -> Option<Result<BusEvent, u64>> {
self.0.recv().await.map(Err)
}
fn topics(&self) -> &[String] {
&[]
}
async fn set_topics(&mut self, _topics: Vec<String>) -> Result<(), Resubscribe> {
Ok(())
}
async fn resubscribe(&mut self) -> Result<(), Resubscribe> {
Err(Resubscribe::Lost(LinkLost::Budget {
waited: Duration::ZERO,
}))
}
}
fn lagging_loop(
poll: Duration,
) -> (
tokio::sync::mpsc::UnboundedSender<u64>,
ScriptedFlock,
tokio::task::JoinHandle<ExitCode>,
tempfile::TempDir,
) {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let flock = ScriptedFlock::answering(Vec::new());
let dir = tempfile::tempdir().unwrap();
let mut config = config_with_sink("127.0.0.1:1".parse().unwrap());
config.poll = UpDuration::from_millis(poll.as_millis() as u64);
let handle = tokio::spawn(run_loop(
LagSource(rx),
flock.clone(),
gave_up_rules(),
&config,
&dir.path().join("barks.jsonl"),
ScriptedConfig::answering(String::new()),
Stop::never(),
));
(tx, flock, handle, dir)
}
async fn advance_and_settle(by: Duration) {
tokio::time::advance(by).await;
for _ in 0..16 {
tokio::task::yield_now().await;
}
}
#[tokio::test(start_paused = true)]
async fn a_burst_of_lags_costs_one_immediate_poll_and_one_deferred() {
let (tx, flock, handle, _dir) = lagging_loop(Duration::from_secs(60));
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
assert_eq!(flock.calls(), 1, "the first lag after quiet polls at once");
for _ in 0..4 {
advance_and_settle(Duration::from_millis(200)).await;
tx.send(1).unwrap();
}
advance_and_settle(Duration::from_millis(199)).await;
assert_eq!(
flock.calls(),
1,
"four lags inside the gap wait for its end"
);
advance_and_settle(Duration::from_millis(1)).await;
assert_eq!(flock.calls(), 2, "the deferred poll runs when the gap ends");
advance_and_settle(LAG_POLL_MIN_GAP * 5).await;
assert_eq!(flock.calls(), 2, "the deferred poll is spent, not repeated");
handle.abort();
}
#[tokio::test(start_paused = true)]
async fn a_lag_after_the_gap_polls_immediately() {
let (tx, flock, handle, _dir) = lagging_loop(Duration::from_secs(60));
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
assert_eq!(flock.calls(), 1);
advance_and_settle(LAG_POLL_MIN_GAP).await;
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
assert_eq!(flock.calls(), 2, "a full gap since the last poll is enough");
advance_and_settle(LAG_POLL_MIN_GAP * 5).await;
assert_eq!(flock.calls(), 2, "nothing was deferred");
handle.abort();
}
#[tokio::test(start_paused = true)]
async fn a_lag_does_not_postpone_the_interval_poll() {
let (tx, flock, handle, _dir) = lagging_loop(Duration::from_secs(2));
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
advance_and_settle(Duration::from_millis(1_500)).await;
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
assert_eq!(flock.calls(), 2, "the second lag is past the gap");
advance_and_settle(Duration::from_millis(499)).await;
assert_eq!(flock.calls(), 2);
advance_and_settle(Duration::from_millis(1)).await;
assert_eq!(flock.calls(), 3, "the tick polls at its own time");
handle.abort();
}
#[tokio::test(start_paused = true)]
async fn an_interval_poll_answers_a_pending_lag() {
let (tx, flock, handle, _dir) = lagging_loop(Duration::from_secs(2));
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
advance_and_settle(Duration::from_millis(1_500)).await;
tx.send(1).unwrap();
advance_and_settle(Duration::from_millis(250)).await;
tx.send(1).unwrap();
advance_and_settle(Duration::ZERO).await;
assert_eq!(flock.calls(), 2, "the third lag is inside the gap");
advance_and_settle(Duration::from_millis(250)).await;
assert_eq!(flock.calls(), 3, "the tick polls");
advance_and_settle(LAG_POLL_MIN_GAP).await;
assert_eq!(flock.calls(), 3, "the tick already covered the lag");
handle.abort();
}
async fn exit_code_once_the_stream_ends(
source: HandoverSource,
tx: broadcast::Sender<BusEvent>,
) -> ExitCode {
let (addr, _captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let loop_handle = tokio::spawn(run_loop(
source,
ScriptedFlock::answering(Vec::new()),
gave_up_rules(),
&config_with_sink(addr),
&barks_path,
ScriptedConfig::answering(String::new()),
Stop::never(),
));
drop(tx);
tokio::time::timeout(Duration::from_secs(5), loop_handle)
.await
.expect("a dog whose stream ended for good must exit, not linger")
.unwrap()
}
#[tokio::test]
async fn a_shepherd_that_never_comes_back_ends_the_dog() {
let (tx, rx) = tokio::sync::broadcast::channel(8);
let (source, resubscribes) = HandoverSource::across(
vec![rx],
LinkLost::Budget {
waited: Duration::from_secs(5),
},
);
let code = exit_code_once_the_stream_ends(source, tx).await;
assert_eq!(
code,
ExitCode::DaemonUnreachable,
"exiting 0 would read as a dog that finished its work"
);
assert_eq!(
resubscribes.load(std::sync::atomic::Ordering::SeqCst),
1,
"it must have tried once before giving up"
);
}
#[tokio::test]
async fn a_refused_re_subscribe_exits_on_the_refusals_own_code() {
let (tx, rx) = tokio::sync::broadcast::channel(8);
let (source, _resubscribes) = HandoverSource::refusing(vec![rx], RpcErrorCode::Unsupported);
let code = exit_code_once_the_stream_ends(source, tx).await;
assert_eq!(
code,
ExitCode::Unsupported,
"exiting DaemonUnreachable would name a shepherd that answered"
);
}
#[tokio::test]
async fn a_refused_dog_exits_saying_it_was_refused() {
let (tx, rx) = tokio::sync::broadcast::channel(8);
let (source, _resubscribes) = HandoverSource::across(
vec![rx],
LinkLost::Refused {
daemon_version: Some("0.9.9".into()),
message: "daemon speaks protocol 9, client speaks 8".into(),
},
);
let code = exit_code_once_the_stream_ends(source, tx).await;
assert_eq!(code, ExitCode::ProtocolMismatch);
}
#[tokio::test]
async fn a_config_change_swaps_barks_sinks_in_place() {
let (old_addr, mut old_captured) = one_shot_sink(200, "").await;
let (new_addr, new_captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let source = ScriptedConfig::answering(format!(
"[sinks.ops]\nkind = \"json\"\nurl = \"http://{new_addr}/hook\"\n"
));
let (tx, rx) = broadcast::channel(16);
let loop_handle = tokio::spawn(run_loop(
rx,
ScriptedFlock::answering(Vec::new()),
gave_up_rules(),
&config_with_sink(old_addr),
&barks_path,
source.clone(),
Stop::never(),
));
tx.send(BusEvent::DogConfigChanged {
dog: "bark".to_owned(),
})
.unwrap();
tx.send(errored_event("web")).unwrap();
let req = tokio::time::timeout(Duration::from_secs(5), new_captured)
.await
.expect("the bark must reach the sink the new section names")
.unwrap();
assert!(String::from_utf8_lossy(&req.body).contains("web"));
assert_eq!(source.calls(), 1, "one frame, one re-ask");
assert!(
!loop_handle.is_finished(),
"bark swaps its config in place; it does not exit to pick one up"
);
assert!(
old_captured.try_recv().is_err(),
"the sink the old section named must be left alone"
);
loop_handle.abort();
}
#[tokio::test]
async fn a_config_change_that_empties_the_section_keeps_the_running_sinks() {
let (addr, captured) = one_shot_sink(200, "").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let source = ScriptedConfig::answering(String::new());
let (tx, rx) = broadcast::channel(16);
let loop_handle = tokio::spawn(run_loop(
rx,
ScriptedFlock::answering(Vec::new()),
gave_up_rules(),
&config_with_sink(addr),
&barks_path,
source.clone(),
Stop::never(),
));
tx.send(BusEvent::DogConfigChanged {
dog: "bark".to_owned(),
})
.unwrap();
tx.send(errored_event("web")).unwrap();
let req = tokio::time::timeout(Duration::from_secs(5), captured)
.await
.expect("the running sink must still get the bark")
.unwrap();
assert!(String::from_utf8_lossy(&req.body).contains("web"));
assert_eq!(source.calls(), 1, "the frame must still be re-asked");
loop_handle.abort();
}
}