pub mod rules;
pub mod sinks;
use core::future::Future;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use serde::Deserialize;
use shep_client::RequestError;
use shep_core::barks::{self, SinkOutcome};
use shep_core::protocol::{BusEvent, ProcessInfo};
use shep_core::values::UpDuration;
use tokio::sync::Mutex;
use tokio::time::MissedTickBehavior;
use self::rules::{Firing, Rule, Rules};
use self::sinks::Sink;
use crate::exit::ExitCode;
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(deny_unknown_fields, default)]
pub struct BarkConfig {
pub sinks: BTreeMap<String, Sink>,
pub rules: Vec<Rule>,
pub poll: UpDuration,
pub history_bytes: u64,
pub sink_timeout: UpDuration,
}
impl Default for BarkConfig {
fn default() -> Self {
Self {
sinks: BTreeMap::new(),
rules: Vec::new(),
poll: UpDuration::from_millis(30_000),
history_bytes: barks::DEFAULT_MAX_BYTES,
sink_timeout: UpDuration::from_millis(10_000),
}
}
}
pub trait EventSource: Send {
fn next(&mut self) -> impl Future<Output = Option<Result<BusEvent, u64>>> + Send;
}
pub trait FlockSource: Send + Sync {
fn flock(&self) -> impl Future<Output = Result<Vec<ProcessInfo>, RequestError>> + Send;
}
pub fn run_loop<E: EventSource, F: FlockSource>(
events: E,
flock: F,
rules: Rules,
config: &BarkConfig,
barks_path: &Path,
) -> impl Future<Output = ExitCode> + Send + use<E, F> {
let sinks = Arc::new(config.sinks.clone());
let sink_timeout = config.sink_timeout.as_duration();
let max_bytes = config.history_bytes;
let poll_period = config.poll.as_duration();
let barks_path = Arc::new(barks_path.to_path_buf());
async move {
let mut events = events;
let mut rules = rules;
let append_lock = Arc::new(Mutex::new(()));
let mut sigterm =
match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
Ok(sigterm) => sigterm,
Err(err) => {
eprintln!("shep dog bark: could not install a SIGTERM handler: {err}");
return ExitCode::Failure;
}
};
let mut poll_interval =
tokio::time::interval_at(tokio::time::Instant::now() + poll_period, poll_period);
poll_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
_ = tokio::signal::ctrl_c() => break,
_ = sigterm.recv() => break,
next = events.next() => {
match next {
None => break,
Some(Ok(event)) => {
let firings = rules.on_event(&event, now_ms());
spawn_firings(firings, &sinks, &append_lock, &barks_path, sink_timeout, max_bytes);
}
Some(Err(_dropped)) => {
reconcile(&flock, &mut rules, &sinks, &append_lock, &barks_path, sink_timeout, max_bytes).await;
}
}
}
_ = poll_interval.tick() => {
reconcile(&flock, &mut rules, &sinks, &append_lock, &barks_path, sink_timeout, max_bytes).await;
}
}
}
ExitCode::Success
}
}
async fn reconcile<F: FlockSource>(
flock: &F,
rules: &mut Rules,
sinks: &Arc<BTreeMap<String, Sink>>,
append_lock: &Arc<Mutex<()>>,
barks_path: &Arc<PathBuf>,
sink_timeout: Duration,
max_bytes: u64,
) {
match flock.flock().await {
Ok(snapshot) => {
let firings = rules.on_poll(&snapshot, now_ms());
spawn_firings(
firings,
sinks,
append_lock,
barks_path,
sink_timeout,
max_bytes,
);
}
Err(err) => eprintln!("shep dog bark: reconciliation poll failed: {err}"),
}
}
fn spawn_firings(
firings: Vec<Firing>,
sinks: &Arc<BTreeMap<String, Sink>>,
append_lock: &Arc<Mutex<()>>,
barks_path: &Arc<PathBuf>,
sink_timeout: Duration,
max_bytes: u64,
) {
for firing in firings {
let sinks = Arc::clone(sinks);
let append_lock = Arc::clone(append_lock);
let barks_path = Arc::clone(barks_path);
tokio::spawn(async move {
deliver_and_record(
firing,
&sinks,
&append_lock,
&barks_path,
sink_timeout,
max_bytes,
)
.await;
});
}
}
async fn deliver_and_record(
firing: Firing,
sinks: &BTreeMap<String, Sink>,
append_lock: &Mutex<()>,
barks_path: &Path,
sink_timeout: Duration,
max_bytes: u64,
) {
let mut bark = firing.bark;
let mut outcomes = Vec::with_capacity(firing.sinks.len());
for name in &firing.sinks {
let outcome = match sinks.get(name) {
Some(sink) => match sinks::deliver(sink, &bark, sink_timeout).await {
Ok(()) => SinkOutcome {
sink: name.clone(),
error: None,
},
Err(err) => SinkOutcome {
sink: name.clone(),
error: Some(err.to_string()),
},
},
None => SinkOutcome {
sink: name.clone(),
error: Some("sink not configured".to_owned()),
},
};
outcomes.push(outcome);
}
bark.sinks = outcomes;
let _guard = append_lock.lock().await;
if let Err(err) = barks::append(barks_path, &bark, max_bytes) {
eprintln!("shep dog bark: could not record a fired bark: {err}");
}
}
fn now_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::SystemTime::UNIX_EPOCH)
.map(|elapsed| elapsed.as_millis() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use std::net::SocketAddr;
use shep_core::barks::Bark;
use shep_core::protocol::ProcessEventKind;
use shep_core::status::ProcStatus;
use tokio::sync::{broadcast, oneshot};
use super::*;
use crate::http::{HttpRequest, read_request, write_response};
impl EventSource for broadcast::Receiver<BusEvent> {
async fn next(&mut self) -> Option<Result<BusEvent, u64>> {
match self.recv().await {
Ok(event) => Some(Ok(event)),
Err(broadcast::error::RecvError::Lagged(count)) => Some(Err(count)),
Err(broadcast::error::RecvError::Closed) => None,
}
}
}
#[derive(Clone)]
struct ScriptedFlock {
answer: Arc<Vec<ProcessInfo>>,
calls: Arc<std::sync::atomic::AtomicU32>,
}
impl ScriptedFlock {
fn answering(answer: Vec<ProcessInfo>) -> Self {
Self {
answer: Arc::new(answer),
calls: Arc::new(std::sync::atomic::AtomicU32::new(0)),
}
}
fn calls(&self) -> u32 {
self.calls.load(std::sync::atomic::Ordering::SeqCst)
}
}
impl FlockSource for ScriptedFlock {
async fn flock(&self) -> Result<Vec<ProcessInfo>, RequestError> {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok((*self.answer).clone())
}
}
async fn one_shot_sink(
status: u16,
body: &str,
) -> (SocketAddr, oneshot::Receiver<HttpRequest>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (tx, rx) = oneshot::channel();
let body = body.to_string();
tokio::spawn(async move {
let (mut stream, _peer) = listener.accept().await.unwrap();
let req = read_request(&mut stream, Duration::from_secs(5))
.await
.unwrap();
write_response(&mut stream, status, "application/json", body.as_bytes())
.await
.unwrap();
let _ = tx.send(req);
});
(addr, rx)
}
async fn slow_sink() -> SocketAddr {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (_stream, _peer) = listener.accept().await.unwrap();
core::future::pending::<()>().await;
});
addr
}
async fn await_real_io<T: Send + 'static>(
timeout: Duration,
fut: impl Future<Output = T> + Send + 'static,
) -> Result<T, tokio::time::error::Elapsed> {
let handle = tokio::runtime::Handle::current();
tokio::task::spawn_blocking(move || handle.block_on(tokio::time::timeout(timeout, fut)))
.await
.expect("the spawn_blocking bridge task must not itself panic")
}
fn base_info(name: &str, status: ProcStatus, restarts: u32) -> ProcessInfo {
ProcessInfo::builder(1, name, status)
.pid(Some(4242))
.restarts(restarts)
.uptime_ms(1_000)
.build()
}
fn errored_info(name: &str, restarts: u32) -> ProcessInfo {
base_info(name, ProcStatus::Errored, restarts)
}
fn process_event(name: &str, kind: ProcessEventKind) -> BusEvent {
BusEvent::Process {
event: kind,
info: base_info(name, ProcStatus::Online, 0),
manually: false,
at_ms: 0,
}
}
fn errored_event(name: &str) -> BusEvent {
process_event(name, ProcessEventKind::Errored)
}
fn log_event(i: u32) -> BusEvent {
BusEvent::LogOut {
id: i,
line: format!("log line {i}"),
}
}
fn gave_up_rules() -> Rules {
let mut sinks = BTreeMap::new();
sinks.insert(
"ops".to_owned(),
Sink::Json {
url: "http://127.0.0.1:1/hook".to_owned(),
body: None,
},
);
Rules::new(
vec![rules::Rule {
when: rules::Trigger::GaveUp,
sinks: vec!["ops".to_owned()],
debounce: UpDuration::from_millis(5 * 60_000),
}],
&sinks,
)
.unwrap()
}
fn config_with_sink(addr: SocketAddr, _barks_path: &Path) -> BarkConfig {
let mut sinks = BTreeMap::new();
sinks.insert(
"ops".to_owned(),
Sink::Json {
url: format!("http://{addr}/hook"),
body: None,
},
);
BarkConfig {
sinks,
rules: Vec::new(),
poll: UpDuration::from_millis(60_000),
history_bytes: barks::DEFAULT_MAX_BYTES,
sink_timeout: UpDuration::from_millis(5_000),
}
}
#[tokio::test(start_paused = true)]
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),
&barks_path,
));
let req = await_real_io(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 = await_real_io(Duration::from_secs(5), {
let barks_path = barks_path.clone();
async move {
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 clock is paused, so no interval has"
);
loop_handle.abort();
}
#[tokio::test]
async fn a_bark_is_recorded_even_when_every_sink_refuses_it() {
let (addr, _captured) = one_shot_sink(500, "refused").await;
let dir = tempfile::tempdir().unwrap();
let barks_path = dir.path().join("barks.jsonl");
let mut sinks = BTreeMap::new();
sinks.insert(
"ops".to_owned(),
Sink::Json {
url: format!("http://{addr}/hook"),
body: None,
},
);
let append_lock = Mutex::new(());
let firing = Firing {
bark: Bark {
at_ms: 1_000,
rule: "gave_up".to_owned(),
subject: "web".to_owned(),
message: "web gave up: restart budget exhausted".to_owned(),
sinks: Vec::new(),
},
sinks: vec!["ops".to_owned()],
};
deliver_and_record(
firing,
&sinks,
&append_lock,
&barks_path,
Duration::from_secs(5),
barks::DEFAULT_MAX_BYTES,
)
.await;
let recorded = shep_core::barks::read(&barks_path).unwrap();
assert_eq!(
recorded.len(),
1,
"a refused delivery must still be recorded"
);
assert_eq!(recorded[0].subject, "web");
assert!(
recorded[0].sinks[0].error.is_some(),
"the 500 must be recorded as a failed delivery, not silently dropped"
);
}
#[tokio::test(start_paused = true)]
async fn a_slow_sink_never_stalls_the_loop() {
let slow_addr = 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(),
Sink::Json {
url: format!("http://{slow_addr}/hook"),
body: None,
},
);
sinks.insert(
"fast".to_owned(),
Sink::Json {
url: format!("http://{fast_addr}/hook"),
body: None,
},
);
let rules = Rules::new(
vec![
rules::Rule {
when: rules::Trigger::GaveUp,
sinks: vec!["slow".to_owned()],
debounce: UpDuration::from_millis(0),
},
rules::Rule {
when: 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(10_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));
let req = await_real_io(Duration::from_millis(50), fast_captured)
.await
.expect(
"the fast sink must be reached promptly; a slow sink in flight \
must not stall the loop",
)
.unwrap();
assert_eq!(req.method, "POST");
loop_handle.abort();
}
#[test]
fn an_empty_section_gets_sane_defaults_not_zeros() {
let parsed: BarkConfig = toml::from_str("").unwrap();
assert_eq!(parsed, BarkConfig::default());
assert_eq!(BarkConfig::default().poll.as_millis(), 30_000);
assert_eq!(
BarkConfig::default().history_bytes,
barks::DEFAULT_MAX_BYTES
);
assert_eq!(BarkConfig::default().sink_timeout.as_millis(), 10_000);
}
}