use core::fmt;
use core::future::Future;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use shep_client::{Client, ConnectError, EventStream, Lagged, RequestError, SshOptions};
use shep_core::paths::ShepPaths;
use shep_core::protocol::{BusEvent, ProcessInfo, Request, Response};
use sysinfo::{MemoryRefreshKind, RefreshKind, System};
use crate::commands::rpc::deadline_for;
use crate::dog_settings::{ask_dog, dog_write};
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: Client,
home: Option<ShepPaths>,
}
impl ClientFlock {
pub(crate) fn client(&self) -> &Client {
&self.client
}
}
impl FlockSource for ClientFlock {
async fn flock(&self) -> Result<Vec<ProcessInfo>, RequestError> {
match self.client.list_flock().await {
Err(RequestError::UnexpectedReply { .. }) => Ok(Vec::new()),
answer => answer,
}
}
async fn send(&self, request: Request) -> Result<Response, RequestError> {
if let Some(write) = dog_write(&request)
&& let Some(paths) = &self.home
{
ask_dog(&self.client, paths, write.dog, &write.table, write.sheep)
.await
.map_err(|failure| failure.request_error())?;
}
let deadline = deadline_for(&request);
self.client.request_with_deadline(request, deadline).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,
home: ShepPaths,
}
impl UnixShepherd {
#[must_use]
pub fn new(socket: &Path, paths: &ShepPaths) -> Self {
Self {
socket: socket.to_path_buf(),
home: paths.clone(),
}
}
}
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?;
subscribed(client, Some(self.home.clone())).await
}
}
#[derive(Debug)]
pub struct HostShepherd {
dest: String,
options: SshOptions,
dialled: bool,
}
impl HostShepherd {
#[must_use]
pub fn new(dest: &str, options: SshOptions) -> Self {
Self {
dest: dest.to_owned(),
options,
dialled: false,
}
}
}
impl Shepherd for HostShepherd {
type Flock = ClientFlock;
type Events = EventStream;
async fn link(&mut self) -> Result<(Self::Flock, Self::Events), LinkError> {
let options = if core::mem::replace(&mut self.dialled, true) {
self.options.clone().batch_mode(true)
} else {
self.options.clone()
};
let client = Client::connect_ssh_with(&self.dest, &options).await?;
subscribed(client, None).await
}
}
async fn subscribed(
client: Client,
home: Option<ShepPaths>,
) -> Result<(ClientFlock, EventStream), LinkError> {
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, home }, 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;
fn log_sizes(&mut self, out: Option<&Path>, err: Option<&Path>) -> Option<u64>;
}
#[derive(Debug)]
pub struct LocalReader {
cores: Option<usize>,
system: System,
seen: std::collections::BTreeMap<PathBuf, u64>,
}
impl LocalReader {
#[must_use]
pub fn new() -> Self {
Self {
cores: std::thread::available_parallelism()
.ok()
.map(NonZeroUsize::get),
system: System::new_with_specifics(
RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()),
),
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;
}
self.system
.refresh_memory_specifics(MemoryRefreshKind::everything());
let load = System::load_average();
Some(HostSample {
load: (load.one, load.five, load.fifteen),
cores: self.cores,
memory_total_bytes: self.system.total_memory(),
memory_used_bytes: self.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)
}
fn log_sizes(&mut self, out: Option<&Path>, err: Option<&Path>) -> Option<u64> {
let out = std::fs::metadata(out?).ok()?.len();
let err = std::fs::metadata(err?).ok()?.len();
Some(out.saturating_add(err))
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct NoLocal;
impl Local for NoLocal {
fn host(&mut self) -> Option<HostSample> {
None
}
fn tail(&mut self, _out: Option<&Path>, _err: Option<&Path>) -> super::tail::Tail {
super::tail::Tail::default()
}
fn log_sizes(&mut self, _out: Option<&Path>, _err: Option<&Path>) -> Option<u64> {
None
}
}
#[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;
#[cfg(unix)]
let budget = std::time::Duration::from_millis(2);
#[cfg(windows)]
let budget = std::time::Duration::from_millis(15);
assert!(
each < budget,
"one host sample took {each:?} (budget {budget:?}); \
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),
);
}
#[tokio::test]
async fn a_version_skewed_shepherd_is_refused_through_client_flock() {
let dir = tempfile::tempdir().unwrap();
let addr = shep_client::testing::control_address(dir.path());
let ack = shep_core::protocol::HelloAck {
daemon_version: "0.1.8".to_string(),
protocol: shep_core::protocol::PROTOCOL_VERSION,
pid: 4242,
min_supported: None,
};
let (client, _fake) = shep_client::testing::fake_client_with_ack(&addr, ack).await;
let flock = ClientFlock { client, home: None };
let mut out = Vec::new();
let mut err = Vec::new();
let mut streams =
crate::output::Streams::bare(&mut out, &mut err, crate::cli::Format::Table);
let code = crate::version_guard::refuse_version_skew(
&mut streams,
flock.client(),
crate::version_guard::VersionGuard::Enforce,
)
.expect_err("a differing crate version must be refused");
assert_eq!(code, crate::exit::ExitCode::VersionSkew);
}
#[tokio::test]
async fn a_matching_version_proceeds_through_client_flock() {
let dir = tempfile::tempdir().unwrap();
let addr = shep_client::testing::control_address(dir.path());
let ack = shep_core::protocol::HelloAck {
daemon_version: env!("CARGO_PKG_VERSION").to_string(),
protocol: shep_core::protocol::PROTOCOL_VERSION,
pid: 4242,
min_supported: None,
};
let (client, _fake) = shep_client::testing::fake_client_with_ack(&addr, ack).await;
let flock = ClientFlock { client, home: None };
let mut out = Vec::new();
let mut err = Vec::new();
let mut streams =
crate::output::Streams::bare(&mut out, &mut err, crate::cli::Format::Table);
crate::version_guard::refuse_version_skew(
&mut streams,
flock.client(),
crate::version_guard::VersionGuard::Enforce,
)
.expect("a matching version is not a skew");
}
#[tokio::test]
async fn a_stop_sent_through_client_flock_asks_for_the_stop_budget() {
let dir = tempfile::tempdir().unwrap();
let addr = shep_client::testing::control_address(dir.path());
let (client, _fake) = shep_client::testing::fake_client_needing_budget(
&addr,
std::time::Duration::from_secs(7),
Response::Stopped(Vec::new()),
)
.await;
let flock = ClientFlock { client, home: None };
let answer = flock
.send(Request::Stop {
selector: shep_core::protocol::SelectorSpec::All,
})
.await;
assert_eq!(answer, Ok(Response::Stopped(Vec::new())));
}
#[test]
fn two_readable_logs_size_to_their_sum() {
let dir = tempfile::Builder::new()
.prefix("shep-logsize-")
.tempdir()
.expect("a tempdir for the logs");
let out = dir.path().join("web-out.log");
let err = dir.path().join("web-err.log");
std::fs::write(&out, b"listening on :8080\n").expect("write the out log");
std::fs::write(&err, b"warn\n").expect("write the err log");
let mut local = LocalReader::new();
assert_eq!(local.log_sizes(Some(&out), Some(&err)), Some(19 + 5));
}
#[test]
fn a_missing_log_sizes_to_nothing_rather_than_to_one_file() {
let dir = tempfile::Builder::new()
.prefix("shep-logsize-")
.tempdir()
.expect("a tempdir for the logs");
let out = dir.path().join("only-out.log");
std::fs::write(&out, b"listening on :8080\n").expect("write the out log");
let rotated_away = dir.path().join("rotated-away-err.log");
let mut local = LocalReader::new();
assert_eq!(local.log_sizes(Some(&out), Some(&rotated_away)), None);
assert_eq!(
local.log_sizes(Some(&rotated_away), Some(&out)),
None,
"either side missing is the same answer"
);
}
#[test]
fn an_unreported_log_path_sizes_to_nothing() {
let mut local = LocalReader::new();
assert_eq!(local.log_sizes(None, None), None);
}
#[cfg(unix)]
mod stand_in {
use std::os::unix::fs::PermissionsExt as _;
use std::path::{Path, PathBuf};
use shep_core::protocol::{
HelloAck, HelloReply, PROTOCOL_VERSION, Reply, Response, ServerFrame, encode_frame,
};
fn framed(payload: &impl serde::Serialize) -> Vec<u8> {
let payload = encode_frame(payload).unwrap();
let mut bytes = u32::try_from(payload.len()).unwrap().to_be_bytes().to_vec();
bytes.extend_from_slice(&payload);
bytes
}
pub(super) fn write(dir: &Path) -> PathBuf {
let ack: HelloReply = Ok(HelloAck {
daemon_version: env!("CARGO_PKG_VERSION").to_owned(),
protocol: PROTOCOL_VERSION,
pid: 4242,
min_supported: None,
});
let subscribed = ServerFrame::Reply(Reply {
id: 1,
result: Ok(Response::Subscribed),
});
std::fs::write(dir.join("hello"), framed(&ack)).unwrap();
std::fs::write(dir.join("subscribed"), framed(&subscribed)).unwrap();
let program = dir.join("ssh");
std::fs::write(
&program,
r#"#!/bin/sh
dir=$(dirname "$0")
echo "$@" >> "$dir/argv"
skip() {
n=0
for b in $(dd bs=1 count=4 2>/dev/null | od -An -v -tu1); do n=$((n * 256 + b)); done
[ "$n" -eq 0 ] || dd bs=1 count="$n" of=/dev/null 2>/dev/null
}
skip; cat "$dir/hello"
skip; cat "$dir/subscribed"
cat >/dev/null
"#,
)
.unwrap();
std::fs::set_permissions(&program, std::fs::Permissions::from_mode(0o755)).unwrap();
program
}
pub(super) fn dials(dir: &Path) -> Vec<String> {
std::fs::read_to_string(dir.join("argv"))
.unwrap()
.lines()
.map(str::to_owned)
.collect()
}
}
#[cfg(unix)]
#[tokio::test]
async fn host_shepherd_dials_a_stand_in_and_subscribes() {
let dir = tempfile::tempdir().unwrap();
let program = stand_in::write(dir.path());
let mut shepherd = HostShepherd::new("build-01", SshOptions::default().program(program));
let (flock, _events) = shepherd
.link()
.await
.expect("the stand-in answers the handshake and the subscribe");
assert_eq!(flock.client().host(), Some("build-01"));
assert_eq!(
stand_in::dials(dir.path()),
["-T -- build-01 shep dial-stdio"]
);
}
#[cfg(unix)]
#[tokio::test]
async fn the_second_dial_is_batch_mode_and_the_first_is_not() {
let dir = tempfile::tempdir().unwrap();
let program = stand_in::write(dir.path());
let mut shepherd = HostShepherd::new("build-01", SshOptions::default().program(program));
let _first = shepherd.link().await.expect("the first dial");
let _second = shepherd.link().await.expect("the redial");
assert_eq!(
stand_in::dials(dir.path()),
[
"-T -- build-01 shep dial-stdio",
"-T -o BatchMode=yes -- build-01 shep dial-stdio",
]
);
}
#[tokio::test]
async fn a_host_that_cannot_be_dialled_is_unreachable() {
let dir = tempfile::tempdir().unwrap();
let program = dir.path().join("no-such-ssh");
let mut shepherd = HostShepherd::new("build-01", SshOptions::default().program(program));
let err = shepherd.link().await.expect_err("there is no ssh to run");
assert!(matches!(err, LinkError::Unreachable(_)), "{err:?}");
}
#[test]
fn no_local_reads_nothing_of_this_machine() {
let log = Path::new(file!());
let mut local = NoLocal;
assert_eq!(local.host(), None);
assert_eq!(
local.tail(Some(log), Some(log)),
crate::lookout::tail::Tail::default()
);
assert_eq!(local.log_sizes(Some(log), Some(log)), None);
}
}
#[cfg(test)]
mod validate_tests;