use core::fmt;
use std::time::Duration;
use shep_client::{Lagged, RequestError};
use shep_core::protocol::BusEvent;
use tokio::sync::mpsc;
use tokio::time::MissedTickBehavior;
use super::app::Msg;
use super::source::{EventSource, FlockSource, Shepherd};
pub const FLOCK_POLL: Duration = Duration::from_secs(2);
pub const RECONNECT_ATTEMPTS: u32 = 5;
pub const RECONNECT_FIRST_WAIT: Duration = Duration::from_millis(250);
pub const RECONNECT_MAX_WAIT: Duration = Duration::from_secs(4);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UiGone;
impl fmt::Display for UiGone {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("the dashboard stopped listening")
}
}
impl core::error::Error for UiGone {}
#[derive(Debug)]
pub struct Channels {
pub polls: mpsc::Receiver<()>,
pub requests: mpsc::Receiver<super::app::Sent>,
}
pub async fn run_link<S: Shepherd>(
mut shepherd: S,
opened: (S::Flock, S::Events),
msgs: mpsc::Sender<Msg>,
mut channels: Channels,
period: Duration,
) {
let mut attempt = 0u32;
let mut wait = RECONNECT_FIRST_WAIT;
let mut connection = Some(opened);
loop {
let (flock, events) = match connection.take() {
Some(pair) => pair,
None => match shepherd.link().await {
Ok(pair) => pair,
Err(err) => {
let _ = err;
attempt += 1;
if attempt > RECONNECT_ATTEMPTS {
let _ = msgs
.send(Msg::Frozen {
at_local: local_now(),
})
.await;
return;
}
let _ = msgs.send(Msg::Retrying { attempt }).await;
tokio::time::sleep(wait).await;
wait = (wait * 2).min(RECONNECT_MAX_WAIT);
continue;
}
},
};
if attempt > 0 && msgs.send(Msg::Relinked).await.is_err() {
return;
}
attempt = 0;
wait = RECONNECT_FIRST_WAIT;
match run_connected(flock, events, msgs.clone(), channels, period).await {
Ok(returned) => channels = returned,
Err(UiGone) => return,
}
}
}
pub async fn run_connected<F: FlockSource, E: EventSource>(
flock: F,
mut events: E,
msgs: mpsc::Sender<Msg>,
mut channels: Channels,
period: Duration,
) -> Result<Channels, UiGone> {
reconcile(&flock, &msgs).await?;
let mut ticker = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
_ = ticker.tick() => reconcile(&flock, &msgs).await?,
_ = channels.polls.recv() => reconcile(&flock, &msgs).await?,
Some(sent) = channels.requests.recv() => {
let result = flock.send(sent.request()).await;
msgs.send(Msg::Replied { sent, result })
.await
.map_err(|_| UiGone)?;
}
next = events.next_event() => match next {
None => return Ok(channels),
Some(Ok(event)) => {
let repair = matches!(event, BusEvent::Dropped { .. });
msgs.send(Msg::Event(event)).await.map_err(|_| UiGone)?;
if repair {
reconcile(&flock, &msgs).await?;
}
}
Some(Err(Lagged { count })) => {
msgs.send(Msg::BusLagged { count })
.await
.map_err(|_| UiGone)?;
reconcile(&flock, &msgs).await?;
}
},
}
}
}
async fn reconcile<F: FlockSource>(flock: &F, msgs: &mpsc::Sender<Msg>) -> Result<(), UiGone> {
match flock.flock().await {
Ok(rows) => msgs
.send(Msg::Snapshot {
rows,
at: std::time::Instant::now(),
})
.await
.map_err(|_| UiGone),
Err(RequestError::Closed) => Ok(()),
Err(_other) => Ok(()),
}
}
fn local_now() -> String {
chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use shep_core::protocol::{ProcessEventKind, ProcessInfo, Request, Response, SelectorSpec};
use shep_core::status::ProcStatus;
use tokio::sync::broadcast;
use crate::lookout::app::Sent;
use crate::lookout::source::LinkError;
fn sheep(id: u32) -> ProcessInfo {
ProcessInfo::builder(id, format!("sheep-{id}"), ProcStatus::Online).build()
}
struct CountingFlock {
polls: Arc<AtomicU64>,
}
impl FlockSource for CountingFlock {
async fn flock(&self) -> Result<Vec<ProcessInfo>, RequestError> {
self.polls.fetch_add(1, Ordering::SeqCst);
Ok(vec![sheep(1)])
}
async fn send(&self, _request: Request) -> Result<Response, RequestError> {
Ok(Response::Described(Vec::new()))
}
}
struct BroadcastEvents(broadcast::Receiver<BusEvent>);
impl EventSource for BroadcastEvents {
async fn next_event(&mut self) -> Option<Result<BusEvent, Lagged>> {
match self.0.recv().await {
Ok(event) => Some(Ok(event)),
Err(broadcast::error::RecvError::Lagged(count)) => Some(Err(Lagged { count })),
Err(broadcast::error::RecvError::Closed) => None,
}
}
}
#[tokio::test(start_paused = true)]
async fn a_lagging_subscriber_polls_immediately_instead_of_waiting() {
let (tx, rx) = broadcast::channel(2);
let polls = Arc::new(AtomicU64::new(0));
let (msg_tx, mut msg_rx) = mpsc::channel(64);
let (_poll_tx, poll_rx) = mpsc::channel(1);
let (_request_tx, request_rx) = mpsc::channel(2);
for id in 0..8 {
let _ = tx.send(BusEvent::Process {
event: ProcessEventKind::Online,
info: sheep(id),
manually: false,
at_ms: 0,
});
}
let flock = CountingFlock {
polls: Arc::clone(&polls),
};
let task = tokio::spawn(run_connected(
flock,
BroadcastEvents(rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
let mut saw_lagged = false;
let mut repaired = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while !(saw_lagged && repaired) {
let Ok(Some(msg)) = tokio::time::timeout_at(deadline, msg_rx.recv()).await else {
break;
};
match msg {
Msg::BusLagged { .. } => saw_lagged = true,
Msg::Snapshot { .. } if saw_lagged => repaired = true,
_ => {}
}
}
task.abort();
assert!(saw_lagged, "the lag reached the reducer");
assert!(
repaired,
"the lag was repaired by a listing rather than left to the interval"
);
assert_eq!(
polls.load(Ordering::SeqCst),
2,
"the opening listing, plus exactly one repair for the lag; the one-hour interval caused none"
);
}
#[tokio::test(start_paused = true)]
async fn a_shepherd_side_drop_polls_and_is_forwarded() {
let (tx, rx) = broadcast::channel(16);
let polls = Arc::new(AtomicU64::new(0));
let (msg_tx, mut msg_rx) = mpsc::channel(64);
let (_poll_tx, poll_rx) = mpsc::channel(1);
let (_request_tx, request_rx) = mpsc::channel(2);
let _ = tx.send(BusEvent::Dropped { count: 9 });
let task = tokio::spawn(run_connected(
CountingFlock {
polls: Arc::clone(&polls),
},
BroadcastEvents(rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
let mut forwarded = false;
let mut repaired = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while !(forwarded && repaired) {
let Ok(Some(msg)) = tokio::time::timeout_at(deadline, msg_rx.recv()).await else {
break;
};
match msg {
Msg::Event(BusEvent::Dropped { count: 9 }) => forwarded = true,
Msg::Snapshot { .. } if forwarded => repaired = true,
_ => {}
}
}
task.abort();
assert!(forwarded, "the drop reached the reducer");
assert!(repaired, "and it triggered a repair listing");
assert_eq!(
polls.load(Ordering::SeqCst),
2,
"the opening listing, plus one repair for the drop"
);
}
#[tokio::test(start_paused = true)]
async fn the_ladder_is_bounded_and_ends_frozen() {
struct NeverConnects {
attempts: Arc<AtomicU64>,
}
impl Shepherd for NeverConnects {
type Flock = CountingFlock;
type Events = BroadcastEvents;
async fn link(&mut self) -> Result<(Self::Flock, Self::Events), LinkError> {
self.attempts.fetch_add(1, Ordering::SeqCst);
Err(LinkError::Unreachable("nothing is listening".to_string()))
}
}
let attempts = Arc::new(AtomicU64::new(0));
let (msg_tx, mut msg_rx) = mpsc::channel(64);
let (_poll_tx, poll_rx) = mpsc::channel(1);
let (_request_tx, request_rx) = mpsc::channel(2);
let (opening_tx, opening_rx) = broadcast::channel(1);
drop(opening_tx);
let task = tokio::spawn(run_link(
NeverConnects {
attempts: Arc::clone(&attempts),
},
(
CountingFlock {
polls: Arc::new(AtomicU64::new(0)),
},
BroadcastEvents(opening_rx),
),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(2),
));
let mut seen = Vec::new();
let done = tokio::time::timeout(Duration::from_secs(120), async {
while let Some(msg) = msg_rx.recv().await {
let frozen = matches!(msg, Msg::Frozen { .. });
seen.push(msg);
if frozen {
break;
}
}
})
.await;
assert!(
done.is_ok(),
"the ladder gave up rather than retrying forever"
);
assert_eq!(
attempts.load(Ordering::SeqCst),
u64::from(RECONNECT_ATTEMPTS) + 1
);
let retries = seen
.iter()
.filter(|msg| matches!(msg, Msg::Retrying { .. }))
.count();
assert_eq!(retries, usize::try_from(RECONNECT_ATTEMPTS).unwrap());
assert!(matches!(seen.last(), Some(Msg::Frozen { .. })));
let ended = tokio::time::timeout(Duration::from_secs(5), task).await;
assert!(ended.is_ok(), "the link task ended after freezing");
}
#[tokio::test(start_paused = true)]
async fn a_successful_relink_reports_live_on_the_first_success() {
struct FailsOnce {
done: bool,
polls: Arc<AtomicU64>,
keepalive: Option<broadcast::Sender<BusEvent>>,
}
impl Shepherd for FailsOnce {
type Flock = CountingFlock;
type Events = BroadcastEvents;
async fn link(&mut self) -> Result<(Self::Flock, Self::Events), LinkError> {
if self.done {
let (tx, rx) = broadcast::channel(16);
self.keepalive = Some(tx);
return Ok((
CountingFlock {
polls: Arc::clone(&self.polls),
},
BroadcastEvents(rx),
));
}
self.done = true;
Err(LinkError::Unreachable("not yet".to_string()))
}
}
let polls = Arc::new(AtomicU64::new(0));
let (msg_tx, mut msg_rx) = mpsc::channel(64);
let (_poll_tx, poll_rx) = mpsc::channel(1);
let (_request_tx, request_rx) = mpsc::channel(2);
let (opening_tx, opening_rx) = broadcast::channel(1);
drop(opening_tx);
let task = tokio::spawn(run_link(
FailsOnce {
done: false,
polls: Arc::clone(&polls),
keepalive: None,
},
(
CountingFlock {
polls: Arc::new(AtomicU64::new(0)),
},
BroadcastEvents(opening_rx),
),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(2),
));
let mut retried = false;
let mut relinked = false;
let mut listed_after_relink = false;
let _ = tokio::time::timeout(Duration::from_secs(30), async {
while let Some(msg) = msg_rx.recv().await {
match msg {
Msg::Retrying { attempt: 1 } => retried = true,
Msg::Relinked => relinked = true,
Msg::Snapshot { .. } if relinked => {
listed_after_relink = true;
break;
}
_ => {}
}
}
})
.await;
task.abort();
assert!(retried, "the failed dial put the banner up");
assert!(
relinked,
"and the FIRST successful re-dial took it down again"
);
assert!(
listed_after_relink,
"the fresh connection re-listed the flock"
);
assert_eq!(
polls.load(Ordering::SeqCst),
1,
"exactly the reconnected connection's opening listing"
);
}
#[tokio::test(start_paused = true)]
async fn the_scheduled_poll_lands_on_the_interval_and_not_at_zero() {
let (tx, rx) = broadcast::channel(16);
let polls = Arc::new(AtomicU64::new(0));
let (msg_tx, _msg_rx) = mpsc::channel(256);
let (_poll_tx, poll_rx) = mpsc::channel(1);
let (_request_tx, request_rx) = mpsc::channel(2);
let task = tokio::spawn(run_connected(
CountingFlock {
polls: Arc::clone(&polls),
},
BroadcastEvents(rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(2),
));
tokio::time::sleep(Duration::from_millis(1900)).await;
assert_eq!(
polls.load(Ordering::SeqCst),
1,
"the opening listing, and NOTHING from the timer before its period elapsed"
);
tokio::time::sleep(Duration::from_millis(2200)).await;
assert_eq!(
polls.load(Ordering::SeqCst),
3,
"the opening listing, plus t=2s and t=4s"
);
drop(tx);
task.abort();
}
struct RecordingFlock {
seen: Arc<std::sync::Mutex<Vec<Request>>>,
}
impl FlockSource for RecordingFlock {
async fn flock(&self) -> Result<Vec<ProcessInfo>, RequestError> {
Ok(vec![sheep(1)])
}
async fn send(&self, request: Request) -> Result<Response, RequestError> {
self.seen.lock().unwrap().push(request.clone());
Ok(Response::Described(vec![sheep(1)]))
}
}
#[tokio::test(start_paused = true)]
async fn a_request_reaches_the_shepherd_and_its_reply_comes_back() {
let (msg_tx, mut msg_rx) = mpsc::channel(64);
let (_poll_tx, poll_rx) = mpsc::channel(1);
let (request_tx, request_rx) = mpsc::channel(2);
let (_events_tx, events_rx) = broadcast::channel(4);
let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
let flock = RecordingFlock {
seen: Arc::clone(&seen),
};
let task = tokio::spawn(run_connected(
flock,
BroadcastEvents(events_rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
request_tx.send(Sent::Lambs { id: 7 }).await.unwrap();
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
let mut answered = None;
while answered.is_none() {
let Ok(Some(msg)) = tokio::time::timeout_at(deadline, msg_rx.recv()).await else {
break;
};
if let Msg::Replied { sent, result } = msg {
answered = Some((sent, result));
}
}
task.abort();
let (sent, result) = answered.expect("the reply came back");
assert_eq!(sent, Sent::Lambs { id: 7 }, "tagged with what it answered");
assert!(matches!(result, Ok(Response::Described(_))));
assert_eq!(
*seen.lock().unwrap(),
vec![Request::Describe {
selector: SelectorSpec::Id(7)
}],
"the id it was asked about, as a selector, and nothing else"
);
}
}