use super::config::BarkConfig;
use super::rules::Rules;
use super::sinks::Sink;
use super::source::{ConfigSource, EventSource, FlockSource, Resubscribe};
use super::*;
use crate::http::{HttpRequest, read_request, write_response};
use shep_client::{LinkLost, RequestError};
use shep_core::barks::{self};
use shep_core::protocol::{BusEvent, ProcessInfo};
use shep_core::protocol::{ProcessEventKind, RpcError, RpcErrorCode};
use shep_core::status::ProcStatus;
use shep_core::values::UpDuration;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{broadcast, oneshot};
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,
}
}
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,
}))
}
}
pub(super) struct HandoverSource {
current: Option<broadcast::Receiver<BusEvent>>,
later: std::collections::VecDeque<broadcast::Receiver<BusEvent>>,
when_gone: LinkLost,
refuses: Option<RpcErrorCode>,
resubscribes: Arc<std::sync::atomic::AtomicU32>,
topics: Vec<String>,
}
impl HandoverSource {
pub(super) fn across(
generations: Vec<broadcast::Receiver<BusEvent>>,
when_gone: LinkLost,
) -> (Self, Arc<std::sync::atomic::AtomicU32>) {
Self::across_inner(generations, when_gone, None)
}
pub(super) fn refusing(
generations: Vec<broadcast::Receiver<BusEvent>>,
code: RpcErrorCode,
) -> (Self, Arc<std::sync::atomic::AtomicU32>) {
Self::across_inner(
generations,
LinkLost::Budget {
waited: Duration::ZERO,
},
Some(code),
)
}
fn across_inner(
generations: Vec<broadcast::Receiver<BusEvent>>,
when_gone: LinkLost,
refuses: Option<RpcErrorCode>,
) -> (Self, Arc<std::sync::atomic::AtomicU32>) {
let mut later: std::collections::VecDeque<_> = generations.into();
let current = later.pop_front();
let resubscribes = Arc::new(std::sync::atomic::AtomicU32::new(0));
(
Self {
current,
later,
when_gone,
refuses,
resubscribes: Arc::clone(&resubscribes),
topics: Vec::new(),
},
resubscribes,
)
}
}
impl EventSource for HandoverSource {
async fn next(&mut self) -> Option<Result<BusEvent, u64>> {
let next = EventSource::next(self.current.as_mut()?).await;
if next.is_none() {
self.current = None;
}
next
}
fn topics(&self) -> &[String] {
&self.topics
}
async fn set_topics(&mut self, topics: Vec<String>) -> Result<(), Resubscribe> {
self.topics = topics;
Ok(())
}
async fn resubscribe(&mut self) -> Result<(), Resubscribe> {
self.resubscribes
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if let Some(code) = self.refuses {
return Err(Resubscribe::Request(RequestError::Rpc(RpcError {
code,
message: "this shepherd does not serve that topic".into(),
daemon_version: None,
})));
}
match self.later.pop_front() {
Some(stream) => {
self.current = Some(stream);
Ok(())
}
None => Err(Resubscribe::Lost(self.when_gone.clone())),
}
}
}
#[derive(Clone)]
pub(super) struct ScriptedFlock {
answer: Arc<Vec<ProcessInfo>>,
calls: Arc<std::sync::atomic::AtomicU32>,
}
impl ScriptedFlock {
pub(super) fn answering(answer: Vec<ProcessInfo>) -> Self {
Self {
answer: Arc::new(answer),
calls: Arc::new(std::sync::atomic::AtomicU32::new(0)),
}
}
pub(super) 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())
}
}
#[derive(Clone)]
pub(super) struct ScriptedConfig {
section: Arc<String>,
calls: Arc<std::sync::atomic::AtomicU32>,
}
impl ScriptedConfig {
pub(super) fn answering(section: String) -> Self {
Self {
section: Arc::new(section),
calls: Arc::new(std::sync::atomic::AtomicU32::new(0)),
}
}
pub(super) fn calls(&self) -> u32 {
self.calls.load(std::sync::atomic::Ordering::SeqCst)
}
}
impl ConfigSource for ScriptedConfig {
async fn section(&self) -> Result<String, RequestError> {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok((*self.section).clone())
}
}
pub(super) 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)
}
pub(super) async fn slow_sink() -> (SocketAddr, tokio::sync::oneshot::Receiver<()>) {
let (addr, connected, release) = slow_sink_with_release().await;
drop(release);
(addr, connected)
}
pub(super) async fn slow_sink_with_release() -> (
SocketAddr,
tokio::sync::oneshot::Receiver<()>,
tokio::sync::oneshot::Sender<()>,
) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (connected, rx) = tokio::sync::oneshot::channel();
let (release, wait) = tokio::sync::oneshot::channel();
tokio::spawn(async move {
let (_stream, _peer) = listener.accept().await.unwrap();
let _ = connected.send(());
match wait.await {
Ok(()) => {}
Err(_dropped_without_sending) => core::future::pending::<()>().await,
}
});
(addr, rx, release)
}
pub(super) 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()
}
pub(super) fn errored_info(name: &str, restarts: u32) -> ProcessInfo {
base_info(name, ProcStatus::Errored, restarts)
}
pub(super) fn process_event(name: &str, kind: ProcessEventKind) -> BusEvent {
BusEvent::Process {
event: kind,
info: base_info(name, ProcStatus::Online, 0),
manually: false,
at_ms: 0,
}
}
pub(super) fn errored_event(name: &str) -> BusEvent {
process_event(name, ProcessEventKind::Errored)
}
pub(super) fn json_sink(url: String) -> Sink {
Sink::Json { url, body: None }
}
pub(super) fn log_event(i: u32) -> BusEvent {
BusEvent::LogOut {
id: i,
line: format!("log line {i}"),
}
}
pub(super) fn gave_up_rules() -> Rules {
let mut sinks = BTreeMap::new();
sinks.insert(
"ops".to_owned(),
json_sink("http://127.0.0.1:1/hook".to_owned()),
);
Rules::new(
vec![rules::Rule {
when: rules::Trigger::GaveUp {},
sinks: vec!["ops".to_owned()],
debounce: UpDuration::from_millis(5 * 60_000),
}],
&sinks,
)
.unwrap()
}
pub(super) fn config_with_sink(addr: SocketAddr) -> BarkConfig {
let mut sinks = BTreeMap::new();
sinks.insert("ops".to_owned(), json_sink(format!("http://{addr}/hook")));
BarkConfig {
sinks,
rules: Vec::new(),
poll: UpDuration::from_millis(60_000),
history_bytes: barks::DEFAULT_MAX_BYTES,
sink_timeout: UpDuration::from_millis(5_000),
}
}