use core::fmt;
use core::future::Future;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use shep_client::{Client, ConnectError, EventStream, Lagged, RequestError};
use shep_core::protocol::{BusEvent, ProcessInfo, Request, Response};
use sysinfo::{MemoryRefreshKind, RefreshKind, System};
use crate::exit::ExitCode;
pub const TOPICS: &[&str] = &["process.*", "daemon.*"];
pub trait FlockSource: Send + Sync {
fn flock(&self) -> impl Future<Output = Result<Vec<ProcessInfo>, RequestError>> + Send;
fn send(&self, request: Request)
-> impl Future<Output = Result<Response, RequestError>> + Send;
}
pub trait EventSource: Send {
fn next_event(&mut self) -> impl Future<Output = Option<Result<BusEvent, Lagged>>> + Send;
}
pub trait Shepherd: Send {
type Flock: FlockSource;
type Events: EventSource;
fn link(
&mut self,
) -> impl Future<Output = Result<(Self::Flock, Self::Events), LinkError>> + Send;
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LinkError {
Unreachable(String),
Protocol(String),
Refused(String),
}
impl LinkError {
#[must_use]
pub fn exit_code(&self) -> ExitCode {
match self {
Self::Protocol(_) => ExitCode::ProtocolMismatch,
Self::Unreachable(_) | Self::Refused(_) => ExitCode::DaemonUnreachable,
}
}
}
impl fmt::Display for LinkError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Unreachable(why) => write!(f, "the shepherd did not answer: {why}"),
Self::Protocol(why) => write!(f, "{why}"),
Self::Refused(why) => write!(f, "the shepherd refused the subscription: {why}"),
}
}
}
impl core::error::Error for LinkError {}
impl From<ConnectError> for LinkError {
fn from(err: ConnectError) -> Self {
if ExitCode::from(&err) == ExitCode::ProtocolMismatch {
return Self::Protocol(err.to_string());
}
Self::Unreachable(err.to_string())
}
}
#[derive(Debug)]
pub struct ClientFlock(Client);
impl FlockSource for ClientFlock {
async fn flock(&self) -> Result<Vec<ProcessInfo>, RequestError> {
match self.0.request(Request::ListFlock).await? {
Response::Flock(flock) => Ok(flock),
_unrecognised => Ok(Vec::new()),
}
}
async fn send(&self, request: Request) -> Result<Response, RequestError> {
self.0.request(request).await
}
}
impl EventSource for EventStream {
async fn next_event(&mut self) -> Option<Result<BusEvent, Lagged>> {
self.next().await
}
}
#[derive(Debug)]
pub struct UnixShepherd {
socket: PathBuf,
}
impl UnixShepherd {
#[must_use]
pub fn new(socket: &Path) -> Self {
Self {
socket: socket.to_path_buf(),
}
}
}
impl Shepherd for UnixShepherd {
type Flock = ClientFlock;
type Events = EventStream;
async fn link(&mut self) -> Result<(Self::Flock, Self::Events), LinkError> {
let client = Client::connect(&self.socket).await?;
let topics = TOPICS.iter().map(|topic| (*topic).to_string()).collect();
let stream = client
.subscribe(topics)
.await
.map_err(|err| LinkError::Refused(err.to_string()))?;
Ok((ClientFlock(client), stream))
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct HostSample {
pub load: (f64, f64, f64),
pub cores: Option<usize>,
pub memory_total_bytes: u64,
pub memory_used_bytes: u64,
pub uptime_seconds: u64,
}
pub trait Local {
fn host(&mut self) -> Option<HostSample>;
fn tail(&mut self, out: Option<&Path>, err: Option<&Path>) -> super::tail::Tail;
}
#[derive(Debug)]
pub struct LocalReader {
cores: Option<usize>,
seen: std::collections::BTreeMap<PathBuf, u64>,
}
impl LocalReader {
#[must_use]
pub fn new() -> Self {
Self {
cores: std::thread::available_parallelism()
.ok()
.map(NonZeroUsize::get),
seen: std::collections::BTreeMap::new(),
}
}
}
impl Default for LocalReader {
fn default() -> Self {
Self::new()
}
}
impl Local for LocalReader {
fn host(&mut self) -> Option<HostSample> {
if !sysinfo::IS_SUPPORTED_SYSTEM {
return None;
}
let system = System::new_with_specifics(
RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()),
);
let load = System::load_average();
Some(HostSample {
load: (load.one, load.five, load.fifteen),
cores: self.cores,
memory_total_bytes: system.total_memory(),
memory_used_bytes: system.used_memory(),
uptime_seconds: System::uptime(),
})
}
fn tail(&mut self, out: Option<&Path>, err: Option<&Path>) -> super::tail::Tail {
super::tail::read(&mut self.seen, out, err)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_host_sample_is_cheap_enough_for_a_one_second_heartbeat() {
let mut local = LocalReader::new();
let _ = local.host();
let started = std::time::Instant::now();
for _ in 0..10 {
let _ = local.host();
}
let each = started.elapsed() / 10;
assert!(
each < std::time::Duration::from_millis(2),
"one host sample took {each:?}; the heartbeat fires every second"
);
}
#[test]
fn an_unsupported_platform_reports_nothing_rather_than_zero() {
let mut local = LocalReader::new();
if sysinfo::IS_SUPPORTED_SYSTEM {
let sample = local.host().expect("a supported platform samples");
assert!(sample.memory_total_bytes > 0, "a supported host has memory");
assert!(sample.memory_used_bytes < sample.memory_total_bytes);
assert!(sample.cores.is_some_and(|cores| cores >= 1));
} else {
assert!(
local.host().is_none(),
"no numbers where there is nothing to read"
);
}
}
#[test]
fn the_core_count_comes_from_std_and_not_from_sysinfo() {
let mut local = LocalReader::new();
assert_eq!(
local.host().and_then(|sample| sample.cores),
std::thread::available_parallelism()
.ok()
.map(NonZeroUsize::get),
);
}
}