use std::path::{Path, PathBuf};
use std::time::Duration;
use alba_executors::protocol::PROTOCOL_VERSION;
use alba_executors::{
BeamContext, CommandSpec, ExecContext, ExecResult, ExecSession, Executor, OutputLine,
PluginExecutor,
};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
use crate::args::PluginCommand;
use crate::exit::EXIT_ALBA_ERROR;
use crate::render::LineSink;
const CANCEL_AFTER: Duration = Duration::from_millis(100);
const CANCEL_TOLERANCE: Duration = Duration::from_secs(1);
const CANCEL_PROMPT_FRACTION: f32 = 0.8;
const EXECUTE_TIMEOUT: Duration = Duration::from_secs(30);
pub fn run(command: &PluginCommand) -> i32 {
let PluginCommand::Check {
binary,
command,
cancel_command,
} = command;
let runtime = match tokio::runtime::Runtime::new() {
Ok(runtime) => runtime,
Err(error) => {
LineSink::stderr().line(&format!("cannot start the async runtime: {error}"));
return EXIT_ALBA_ERROR;
}
};
runtime.block_on(check(binary, command, cancel_command))
}
async fn check(binary: &Path, command: &str, cancel_command: &str) -> i32 {
let mut out = LineSink::stdout();
if !binary.is_file() {
out.line(&format!(
"cannot check `{}`: no such file",
binary.display()
));
return EXIT_ALBA_ERROR;
}
let (session, tx, mut rx) = match open_session(PluginExecutor::new(binary.to_path_buf())).await
{
Ok(opened) => {
out.line("\u{2713} handshake");
opened
}
Err(message) => {
out.line(&format!("\u{2717} handshake: {message}"));
out.line("not conformant");
return 1;
}
};
let mut conformant = execute_check(session, command, tx, &mut rx, &mut out).await;
conformant &= cancel_and_close_check(binary, cancel_command, &mut out).await;
if conformant {
out.line(&format!("conformant: protocol v{PROTOCOL_VERSION}"));
0
} else {
out.line("not conformant");
1
}
}
async fn open_session(
executor: PluginExecutor,
) -> Result<
(
Box<dyn ExecSession>,
UnboundedSender<OutputLine>,
UnboundedReceiver<OutputLine>,
),
String,
> {
let (tx, rx) = unbounded_channel();
let context = BeamContext {
beam: "plugin-check".to_string(),
dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
options: serde_json::json!({}),
output: tx.clone(),
cancel: CancellationToken::new(),
};
let session = executor
.open(context)
.await
.map_err(|error| error.to_string())?;
Ok((session, tx, rx))
}
async fn execute_check(
mut session: Box<dyn ExecSession>,
command: &str,
tx: UnboundedSender<OutputLine>,
rx: &mut UnboundedReceiver<OutputLine>,
out: &mut LineSink<std::io::Stdout>,
) -> bool {
let cmd = CommandSpec {
command: command.to_string(),
env: Vec::new(),
cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
};
let ctx = ExecContext {
output: tx,
cancel: CancellationToken::new(),
};
let outcome = tokio::time::timeout(EXECUTE_TIMEOUT, session.execute(cmd, ctx)).await;
session.kill().await;
match outcome {
Ok(Ok(ExecResult { exit_code })) => {
let mut lines = 0usize;
while rx.try_recv().is_ok() {
lines += 1;
}
let plural = if lines == 1 { "" } else { "s" };
out.line(&format!(
"\u{2713} execute: exit code {exit_code}, {lines} output line{plural}"
));
true
}
Ok(Err(error)) => {
out.line(&format!("\u{2717} execute: {error}"));
false
}
Err(_elapsed) => {
out.line(&format!(
"\u{2717} execute: no `exit` message within {EXECUTE_TIMEOUT:?}"
));
false
}
}
}
async fn cancel_and_close_check(
binary: &Path,
cancel_command: &str,
out: &mut LineSink<std::io::Stdout>,
) -> bool {
let executor = PluginExecutor::with_timeouts(
binary.to_path_buf(),
PluginExecutor::DEFAULT_HANDSHAKE_TIMEOUT,
PluginExecutor::DEFAULT_GRACE,
);
let (mut session, tx, _rx) = match open_session(executor).await {
Ok(opened) => opened,
Err(message) => {
out.line(&format!(
"\u{2717} cancel: cannot open a fresh session — {message}"
));
out.line("\u{2717} close: skipped — no session to close");
return false;
}
};
let cancel = CancellationToken::new();
let fire = cancel.clone();
tokio::spawn(async move {
tokio::time::sleep(CANCEL_AFTER).await;
fire.cancel();
});
let cmd = CommandSpec {
command: cancel_command.to_string(),
env: Vec::new(),
cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
};
let ctx = ExecContext { output: tx, cancel };
let grace = PluginExecutor::DEFAULT_GRACE;
let bound = grace + CANCEL_TOLERANCE;
let started = Instant::now();
let outcome = tokio::time::timeout(bound, session.execute(cmd, ctx)).await;
let total_elapsed = started.elapsed();
let answered_in = total_elapsed.saturating_sub(CANCEL_AFTER);
let cancel_conforms = match outcome {
Ok(Ok(_result)) if total_elapsed < CANCEL_AFTER => {
out.line(&format!(
"\u{2717} cancel: `--cancel-command` finished in {total_elapsed:?}, before \
`cancel` was even sent at {CANCEL_AFTER:?} — pick a command that runs longer"
));
false
}
Ok(Ok(_result)) if answered_in < grace.mul_f32(CANCEL_PROMPT_FRACTION) => {
out.line(&format!("\u{2713} cancel (answered in {answered_in:?})"));
true
}
Ok(Ok(_result)) => {
out.line(&format!(
"\u{2717} cancel: no answer within the {grace:?} grace ({answered_in:?} after \
cancel); the host had to kill the plugin"
));
false
}
Ok(Err(error)) => {
out.line(&format!("\u{2717} cancel: {error}"));
false
}
Err(_elapsed) => {
out.line(&format!(
"\u{2717} cancel: no response within {bound:?} of asking it to stop"
));
out.line("\u{2717} close: skipped — the cancel check did not return");
session.kill().await;
return false;
}
};
match session.close().await {
Ok(()) => {
out.line("\u{2713} close");
cancel_conforms
}
Err(error) => {
out.line(&format!("\u{2717} close: {error}"));
false
}
}
}