use crate::lifecycle::{listener, supervisor};
use std::io;
use std::net::TcpListener;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
pub const RELEASE_POINTER_PATH: &str = "storage/releases/current";
pub fn resolve_binary_path() -> io::Result<PathBuf> {
resolve_binary_path_from(RELEASE_POINTER_PATH)
}
fn resolve_binary_path_from(pointer_path: &str) -> io::Result<PathBuf> {
if let Ok(contents) = std::fs::read_to_string(pointer_path) {
let path = contents.trim();
if !path.is_empty() {
return Ok(PathBuf::from(path));
}
}
std::env::current_exe()
}
pub const READY_MARKER: &str = "__LARUST_HANDOFF_READY__";
pub async fn spawn_replacement_and_wait_for_ready(
listener: &TcpListener,
binary_path: &Path,
ready_timeout: Duration,
register_with_supervisor: bool,
) -> io::Result<Option<Child>> {
let mut command = Command::new(binary_path);
command
.env(listener::INHERIT_LISTENER_ENV, "1")
.stdin(Stdio::piped())
.stdout(Stdio::inherit())
.stderr(Stdio::piped());
#[cfg(unix)]
let pre_spawn_encoded = Some(listener::prepare_for_handoff(listener, 0)?);
#[cfg(not(unix))]
let pre_spawn_encoded: Option<String> = None;
let spawn_result = command.spawn();
if let Some(encoded) = &pre_spawn_encoded {
listener::close_duplicated_fd(encoded);
}
let mut child = spawn_result?;
if register_with_supervisor {
supervisor::register(&child);
}
let child_pid = child
.id()
.ok_or_else(|| io::Error::other("spawned replacement has no pid"))?;
let encoded = match pre_spawn_encoded {
Some(encoded) => encoded,
None => listener::prepare_for_handoff(listener, child_pid)?,
};
{
let mut stdin = child.stdin.take().expect("stdin was piped above");
let write_result = async {
stdin.write_all(encoded.as_bytes()).await?;
stdin.write_all(b"\n").await
}
.await;
if let Err(source) = write_result {
tracing::warn!(
error = %source,
"restart handoff replacement's stdin closed before this process could \
hand off the listener - it exited or crashed immediately"
);
let _ = child.kill().await;
let _ = child.wait().await;
return Ok(None);
}
}
let stderr = child.stderr.take().expect("stderr was piped above");
let mut lines = BufReader::new(stderr).lines();
let mut early_output: Vec<String> = Vec::new();
let ready_result = tokio::time::timeout(ready_timeout, async {
while let Ok(Some(line)) = lines.next_line().await {
if line.trim() == READY_MARKER {
return true;
}
if early_output.len() < 10 {
early_output.push(line);
}
}
false
})
.await;
let became_ready = match ready_result {
Ok(true) => true,
Ok(false) => {
tracing::warn!(
?early_output,
"restart handoff replacement's stderr closed before it announced \
readiness - it exited or crashed rather than merely being slow"
);
false
}
Err(_) => {
tracing::warn!(
?ready_timeout,
?early_output,
"restart handoff replacement did not announce readiness within the timeout"
);
false
}
};
if became_ready {
tokio::spawn(forward_drained_lines(lines, |line| eprintln!("{line}")));
Ok(Some(child))
} else {
let _ = child.kill().await;
let _ = child.wait().await;
Ok(None)
}
}
async fn forward_drained_lines<R: AsyncBufRead + Unpin>(
mut lines: tokio::io::Lines<R>,
mut sink: impl FnMut(&str),
) {
while let Ok(Some(line)) = lines.next_line().await {
sink(&line);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolves_to_the_pointer_files_contents_when_present_and_non_empty() {
let dir = tempfile::tempdir().unwrap();
let pointer = dir.path().join("current");
std::fs::write(&pointer, " /releases/abc123/my-app \n").unwrap();
let resolved = resolve_binary_path_from(pointer.to_str().unwrap()).unwrap();
assert_eq!(resolved, PathBuf::from("/releases/abc123/my-app"));
}
#[test]
fn falls_back_to_current_exe_when_the_pointer_file_is_missing() {
let dir = tempfile::tempdir().unwrap();
let pointer = dir.path().join("does-not-exist");
let resolved = resolve_binary_path_from(pointer.to_str().unwrap()).unwrap();
assert_eq!(resolved, std::env::current_exe().unwrap());
}
#[test]
fn falls_back_to_current_exe_when_the_pointer_file_is_empty_or_blank() {
let dir = tempfile::tempdir().unwrap();
let pointer = dir.path().join("current");
std::fs::write(&pointer, " \n").unwrap();
let resolved = resolve_binary_path_from(pointer.to_str().unwrap()).unwrap();
assert_eq!(resolved, std::env::current_exe().unwrap());
}
#[tokio::test]
async fn forward_drained_lines_hands_every_line_to_the_sink_instead_of_discarding_it() {
let (mut writer, reader) = tokio::io::duplex(1024);
let lines = tokio::io::AsyncBufReadExt::lines(BufReader::new(reader));
let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let seen_in_task = seen.clone();
let task = tokio::spawn(forward_drained_lines(lines, move |line| {
seen_in_task.lock().unwrap().push(line.to_string());
}));
writer
.write_all(b"couldn't open the log file - falling back to stdout only\n")
.await
.unwrap();
writer
.write_all(b"a second unrelated line\n")
.await
.unwrap();
drop(writer);
task.await.unwrap();
assert_eq!(
*seen.lock().unwrap(),
vec![
"couldn't open the log file - falling back to stdout only".to_string(),
"a second unrelated line".to_string(),
],
"every line written after the handshake must reach the sink, not be discarded"
);
}
}