use core::fmt;
use std::collections::HashMap;
use std::time::Duration;
use futures_util::StreamExt;
use futures_util::stream::FuturesUnordered;
use shep_client::{Lagged, RequestError};
use shep_core::protocol::BusEvent;
use tokio::sync::mpsc;
use tokio::time::MissedTickBehavior;
use super::app::{Msg, Sent};
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) => {
attempt += 1;
if attempt > RECONNECT_ATTEMPTS {
let _ = msgs
.send(Msg::Frozen {
at_local: local_now(),
why: err.to_string(),
})
.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);
let mut pulls = FuturesUnordered::new();
let mut pending: HashMap<u64, Sent> = HashMap::new();
let mut next_pull = 0u64;
loop {
tokio::select! {
_ = ticker.tick() => reconcile(&flock, &msgs).await?,
_ = channels.polls.recv() => reconcile(&flock, &msgs).await?,
Some(sent) = channels.requests.recv() => {
if matches!(sent, Sent::DogPane { .. } | Sent::DogPaneAction { .. }) {
next_pull += 1;
pending.insert(next_pull, sent.clone());
pulls.push(pull_reply(&flock, next_pull, sent));
} else {
msgs.send(reply_to(&flock, sent).await)
.await
.map_err(|_| UiGone)?;
}
}
Some((id, msg)) = pulls.next() => {
pending.remove(&id);
msgs.send(msg).await.map_err(|_| UiGone)?;
}
next = events.next_event() => match next {
None => {
for (_, sent) in pending.drain() {
let lost = Msg::Replied { sent, result: Err(RequestError::Closed) };
msgs.send(lost).await.map_err(|_| UiGone)?;
}
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 reply_to<F: FlockSource>(flock: &F, sent: Sent) -> Msg {
let result = flock.send(sent.request()).await;
Msg::Replied { sent, result }
}
async fn pull_reply<F: FlockSource>(flock: &F, id: u64, sent: Sent) -> (u64, Msg) {
let result = flock.pull(sent.request()).await;
(id, Msg::Replied { sent, result })
}
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"
);
}
struct ScriptedFlock {
gate: Arc<tokio::sync::Notify>,
hang_triggers: bool,
}
impl FlockSource for ScriptedFlock {
async fn flock(&self) -> Result<Vec<ProcessInfo>, RequestError> {
Ok(vec![sheep(1)])
}
async fn send(&self, request: Request) -> Result<Response, RequestError> {
assert!(
!matches!(request, Request::Trigger { .. }),
"a pull went over the connection the polls share"
);
self.gate.notified().await;
Ok(Response::Described(Vec::new()))
}
async fn pull(&self, request: Request) -> Result<Response, RequestError> {
assert!(matches!(request, Request::Trigger { .. }));
if self.hang_triggers {
std::future::pending().await
}
Ok(Response::Triggered(Vec::new()))
}
}
fn dog_pane() -> Sent {
Sent::DogPane {
name: "paddock".to_string(),
opening: 1,
}
}
#[tokio::test(start_paused = true)]
async fn a_pane_pull_that_never_answers_holds_up_neither_a_poll_nor_an_event() {
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 task = tokio::spawn(run_connected(
ScriptedFlock {
gate: Arc::new(tokio::sync::Notify::new()),
hang_triggers: true,
},
BroadcastEvents(events_rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
request_tx.send(dog_pane()).await.unwrap();
poll_tx.send(()).await.unwrap();
let _ = events_tx.send(BusEvent::Dropped { count: 3 });
let mut snapshots = 0;
let mut event = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while snapshots < 3 || !event {
let Ok(Some(msg)) = tokio::time::timeout_at(deadline, msg_rx.recv()).await else {
break;
};
match msg {
Msg::Snapshot { .. } => snapshots += 1,
Msg::Event(BusEvent::Dropped { count: 3 }) => event = true,
Msg::Replied { .. } => panic!("a request that never answers replied"),
_ => {}
}
}
task.abort();
assert!(event, "the event reached the UI while the pull was out");
assert_eq!(
snapshots, 3,
"the opening listing, the poll and the event's repair all landed"
);
}
#[tokio::test(start_paused = true)]
async fn a_pane_pull_reply_arrives_as_replied() {
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 task = tokio::spawn(run_connected(
ScriptedFlock {
gate: Arc::new(tokio::sync::Notify::new()),
hang_triggers: false,
},
BroadcastEvents(events_rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
request_tx.send(dog_pane()).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, dog_pane(), "tagged with what it answered");
assert!(matches!(result, Ok(Response::Triggered(_))));
}
#[tokio::test(start_paused = true)]
async fn an_ordinary_request_is_still_awaited_inline() {
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 gate = Arc::new(tokio::sync::Notify::new());
let task = tokio::spawn(run_connected(
ScriptedFlock {
gate: Arc::clone(&gate),
hang_triggers: false,
},
BroadcastEvents(events_rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
assert!(matches!(msg_rx.recv().await, Some(Msg::Snapshot { .. })));
request_tx.send(Sent::Lambs { id: 7 }).await.unwrap();
tokio::time::sleep(Duration::from_millis(10)).await;
poll_tx.send(()).await.unwrap();
let held = tokio::time::timeout(Duration::from_secs(5), msg_rx.recv()).await;
assert!(held.is_err(), "the poll got past a request in flight");
gate.notify_one();
let first = msg_rx.recv().await;
let second = msg_rx.recv().await;
task.abort();
assert!(
matches!(first, Some(Msg::Replied { .. })),
"the reply first: {first:?}"
);
assert!(
matches!(second, Some(Msg::Snapshot { .. })),
"then the poll that waited: {second:?}"
);
}
#[tokio::test(start_paused = true)]
async fn a_pull_still_out_when_the_subscription_ends_is_answered_closed() {
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 task = tokio::spawn(run_connected(
ScriptedFlock {
gate: Arc::new(tokio::sync::Notify::new()),
hang_triggers: true,
},
BroadcastEvents(events_rx),
msg_tx,
Channels {
polls: poll_rx,
requests: request_rx,
},
Duration::from_secs(3600),
));
assert!(matches!(msg_rx.recv().await, Some(Msg::Snapshot { .. })));
request_tx.send(dog_pane()).await.unwrap();
tokio::time::sleep(Duration::from_millis(10)).await;
drop(events_tx);
let lost = tokio::time::timeout(Duration::from_secs(5), msg_rx.recv()).await;
assert!(
matches!(
lost,
Ok(Some(Msg::Replied { ref sent, result: Err(RequestError::Closed) }))
if *sent == dog_pane()
),
"{lost:?}"
);
assert!(task.await.unwrap().is_ok(), "and the ladder gets its turn");
}
}