use std::future::Future;
use std::io;
use std::time::Duration;
use futures_util::future::join_all;
use shep_client::{Client, ConnectError, RequestError, SshOptions};
use shep_core::destination::Destination;
use shep_core::paths::ShepPaths;
use shep_core::protocol::{HostUsage, ProcessInfo};
use tokio::time::error::Elapsed;
use super::flock_display::{flock_from_roll, host_usage, print_listing};
use crate::cli::Format;
use crate::client::unreachable_message;
use crate::commands::hosts::{HOST_CONNECT_TIMEOUT, unattended};
use crate::commands::rpc::UNRECOGNISED;
use crate::commands::shep_toml::ShepToml;
use crate::exit::ExitCode;
use crate::output::{HostListing, Streams, emit_fleet_flock, write_outcome};
use crate::version_guard::{
VersionGuard, flock_connect_refusal_message, host_skew, refuse_version_skew,
};
const LOCAL_HEADING: &str = "local";
struct Answer {
procs: Vec<ProcessInfo>,
usage: Option<Option<HostUsage>>,
}
struct HostAnswer {
answer: Answer,
skew: Option<String>,
}
struct Failure {
code: ExitCode,
message: String,
}
impl Failure {
fn of_connect(err: &ConnectError) -> Self {
Self {
code: ExitCode::from(err),
message: unreachable_message(err),
}
}
fn abandoned(timeout: Duration) -> Self {
Self {
code: ExitCode::DaemonUnreachable,
message: format!("did not answer within {timeout:?}"),
}
}
fn of_request(err: &RequestError) -> Self {
match err {
RequestError::UnexpectedReply { .. } => Self {
code: ExitCode::Internal,
message: UNRECOGNISED.to_owned(),
},
other => Self {
code: ExitCode::from(other),
message: other.to_string(),
},
}
}
}
pub(crate) fn listed_hosts(streams: &mut Streams<'_>, paths: &ShepPaths) -> Vec<String> {
let names = ShepToml::read_only(&paths.daemon_config).and_then(|doc| doc.host_names());
names.unwrap_or_else(|_| {
let message = format!(
"the hosts list in {} was not read, so only this machine's flock is listed; \
`shep host list` says why",
paths.daemon_config.display()
);
streams.aside("hosts_unread", &message);
Vec::new()
})
}
pub(crate) async fn flock_all(
streams: &mut Streams<'_>,
paths: &ShepPaths,
guard: VersionGuard,
hosts: &[String],
) -> ExitCode {
let ssh = unattended(HOST_CONNECT_TIMEOUT);
fleet(streams, paths, guard, hosts, &ssh, HOST_CONNECT_TIMEOUT).await
}
pub(super) async fn fleet(
streams: &mut Streams<'_>,
paths: &ShepPaths,
guard: VersionGuard,
hosts: &[String],
ssh: &SshOptions,
timeout: Duration,
) -> ExitCode {
let (local, remote) = tokio::join!(
tokio::time::timeout(timeout, Client::connect(&paths.socket)),
join_all(hosts.iter().map(|host| ask_host(host, ssh, timeout, guard))),
);
let mut tables = Tables::new(streams.fmt);
let mut first_failure = None;
let mut note = |code: ExitCode| {
first_failure.get_or_insert(code);
};
match ask_local(streams, guard, local, timeout).await {
Local::Listed(answer) => {
if let Err(done) = tables.print(streams, None, answer, None) {
return done;
}
}
Local::Roll if streams.fmt == Format::Json => {
let message = format!("{LOCAL_HEADING}: no shepherd running");
note(streams.fail(ExitCode::DaemonUnreachable, &message));
}
Local::Roll => match print_roll(streams, &mut tables, paths) {
Ok(code) => note(code),
Err(done) => return done,
},
Local::Failed(code) => note(code),
}
for (host, result) in hosts.iter().zip(remote) {
match result {
Ok(HostAnswer { answer, skew }) => {
if let Err(done) = tables.print(streams, Some(host), answer, skew) {
return done;
}
}
Err(Failure { code, message }) => {
note(streams.fail(code, &format!("{host}: {message}")));
}
}
}
if let Err(done) = tables.finish(streams) {
return done;
}
first_failure.unwrap_or(ExitCode::Success)
}
async fn within<T>(
timeout: Duration,
ask: impl Future<Output = Result<T, Failure>>,
) -> Result<T, Failure> {
tokio::time::timeout(timeout, ask)
.await
.unwrap_or_else(|_elapsed| Err(Failure::abandoned(timeout)))
}
struct Tables {
printed: bool,
held: Option<Vec<HostListing>>,
}
impl Tables {
fn new(fmt: Format) -> Self {
Self {
printed: false,
held: (fmt == Format::Json).then(Vec::new),
}
}
fn open(&mut self, streams: &mut Streams<'_>, heading: &str) -> io::Result<()> {
if self.printed {
writeln!(streams.out)?;
}
self.printed = true;
writeln!(streams.out, "{heading}")
}
fn print(
&mut self,
streams: &mut Streams<'_>,
host: Option<&str>,
answer: Answer,
skew: Option<String>,
) -> Result<(), ExitCode> {
if let Some(warning) = skew {
streams.aside("version_skew", &warning);
}
if let Some(held) = &mut self.held {
held.push(HostListing {
host: host.map(str::to_owned),
procs: answer.procs,
usage: answer.usage,
});
return Ok(());
}
self.open(streams, host.unwrap_or(LOCAL_HEADING))
.and_then(|()| print_listing(streams, answer.procs, answer.usage, None))
.map_err(|err| write_outcome(Err(err)))
}
fn finish(self, streams: &mut Streams<'_>) -> Result<(), ExitCode> {
match self.held {
Some(held) => emit_fleet_flock(&mut *streams.out, "flock", held)
.map_err(|err| write_outcome(Err(err))),
None => Ok(()),
}
}
}
enum Local {
Listed(Answer),
Roll,
Failed(ExitCode),
}
async fn ask_local(
streams: &mut Streams<'_>,
guard: VersionGuard,
connected: Result<Result<Client, ConnectError>, Elapsed>,
timeout: Duration,
) -> Local {
let client = match connected {
Ok(Ok(client)) => client,
Ok(Err(ConnectError::Connect { .. })) => return Local::Roll,
Err(_elapsed) => {
let message = format!("{LOCAL_HEADING}: {}", Failure::abandoned(timeout).message);
return Local::Failed(streams.fail(ExitCode::DaemonUnreachable, &message));
}
Ok(Err(err)) => {
let message = format!("{LOCAL_HEADING}: {}", flock_connect_refusal_message(&err));
return Local::Failed(streams.fail(ExitCode::from(&err), &message));
}
};
if let Err(code) = refuse_version_skew(streams, &client, guard) {
return Local::Failed(code);
}
match within(timeout, fetch(&client)).await {
Ok(answer) => Local::Listed(answer),
Err(Failure { code, message }) => {
Local::Failed(streams.fail(code, &format!("{LOCAL_HEADING}: {message}")))
}
}
}
fn print_roll(
streams: &mut Streams<'_>,
tables: &mut Tables,
paths: &ShepPaths,
) -> Result<ExitCode, ExitCode> {
let mut rolled = Vec::new();
let code = flock_from_roll(
&mut Streams {
out: &mut rolled,
err: &mut *streams.err,
style: streams.style,
fmt: streams.fmt,
},
paths,
);
if !rolled.is_empty() {
tables
.open(streams, LOCAL_HEADING)
.and_then(|()| streams.out.write_all(&rolled))
.map_err(|err| write_outcome(Err(err)))?;
}
Ok(code)
}
async fn fetch(client: &Client) -> Result<Answer, Failure> {
let procs = client
.list_flock()
.await
.map_err(|e| Failure::of_request(&e))?;
let usage = host_usage(client).await;
Ok(Answer { procs, usage })
}
async fn ask_host(
host: &str,
ssh: &SshOptions,
timeout: Duration,
guard: VersionGuard,
) -> Result<HostAnswer, Failure> {
let destination = Destination::parse(host).map_err(|err| Failure {
code: ExitCode::InvalidConfig,
message: format!("not a destination: {err}"),
})?;
within(timeout, async {
let client = Client::connect_ssh_with(destination.as_str(), ssh)
.await
.map_err(|e| Failure::of_connect(&e))?;
let answer = match host_skew(&client, guard) {
Ok(skew) => fetch(&client)
.await
.map(|answer| HostAnswer { answer, skew }),
Err(message) => Err(Failure {
code: ExitCode::VersionSkew,
message,
}),
};
let _ = client.close().await;
answer
})
.await
}
#[cfg(all(test, unix))]
mod tests;