use std::ffi::{OsStr, OsString};
use std::fmt;
use std::io;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;
use std::time::Duration;
use lgwks_deps::tokio::process::Command;
use super::supervise::{CleanupReceipt, Containment};
pub use lgwks_std::process::{
OrphanReap, ProcessIdentity, ProcessIdentityError, reap_orphaned_group,
};
use crate::journal::frame::{LENGTH_BYTES, declared_length, is_possible_length};
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[non_exhaustive]
pub enum StdioPolicy {
#[default]
Inherit,
Null,
Capture(NonZeroUsize),
}
#[derive(Clone, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum EnvDelta {
Set {
key: OsString,
value: OsString,
},
Remove {
key: OsString,
},
Clear,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ProcessSpec {
program: OsString,
args: Vec<OsString>,
env: Vec<EnvDelta>,
cwd: Option<PathBuf>,
stdin: StdioPolicy,
stdout: StdioPolicy,
stderr: StdioPolicy,
stdout_file: Option<PathBuf>,
stderr_file: Option<PathBuf>,
deadline: Option<Duration>,
}
impl ProcessSpec {
pub fn new(program: impl AsRef<OsStr>) -> Self {
Self {
program: program.as_ref().to_os_string(),
args: Vec::new(),
env: Vec::new(),
cwd: None,
stdin: StdioPolicy::default(),
stdout: StdioPolicy::default(),
stderr: StdioPolicy::default(),
stdout_file: None,
stderr_file: None,
deadline: None,
}
}
pub fn arg(&mut self, arg: impl AsRef<OsStr>) -> &mut Self {
self.args.push(arg.as_ref().to_os_string());
self
}
pub fn env(&mut self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) -> &mut Self {
self.env.push(EnvDelta::Set {
key: key.as_ref().to_os_string(),
value: value.as_ref().to_os_string(),
});
self
}
pub fn env_remove(&mut self, key: impl AsRef<OsStr>) -> &mut Self {
self.env.push(EnvDelta::Remove {
key: key.as_ref().to_os_string(),
});
self
}
pub fn env_clear(&mut self) -> &mut Self {
self.env.push(EnvDelta::Clear);
self
}
pub fn current_dir(&mut self, cwd: impl AsRef<Path>) -> &mut Self {
self.cwd = Some(cwd.as_ref().to_path_buf());
self
}
pub fn stdin(&mut self, policy: StdioPolicy) -> &mut Self {
self.stdin = policy;
self
}
pub fn stdout(&mut self, policy: StdioPolicy) -> &mut Self {
self.stdout = policy;
self.stdout_file = None;
self
}
pub fn stderr(&mut self, policy: StdioPolicy) -> &mut Self {
self.stderr = policy;
self.stderr_file = None;
self
}
pub fn capture_stdout(&mut self, limit: NonZeroUsize) -> &mut Self {
self.stdout = StdioPolicy::Capture(limit);
self.stdout_file = None;
self
}
pub fn capture_stderr(&mut self, limit: NonZeroUsize) -> &mut Self {
self.stderr = StdioPolicy::Capture(limit);
self.stderr_file = None;
self
}
pub fn stdout_to_file(&mut self, path: impl AsRef<Path>) -> &mut Self {
self.stdout = StdioPolicy::Null;
self.stdout_file = Some(path.as_ref().to_path_buf());
self
}
pub fn stderr_to_file(&mut self, path: impl AsRef<Path>) -> &mut Self {
self.stderr = StdioPolicy::Null;
self.stderr_file = Some(path.as_ref().to_path_buf());
self
}
pub fn deadline(&mut self, deadline: Duration) -> &mut Self {
self.deadline = Some(deadline);
self
}
#[must_use]
pub fn program(&self) -> &OsStr {
&self.program
}
#[must_use]
pub fn args(&self) -> &[OsString] {
&self.args
}
#[must_use]
pub fn env_deltas(&self) -> &[EnvDelta] {
&self.env
}
#[must_use]
pub fn cwd(&self) -> Option<&Path> {
self.cwd.as_deref()
}
#[must_use]
pub const fn stdin_policy(&self) -> StdioPolicy {
self.stdin
}
#[must_use]
pub const fn stdout_policy(&self) -> StdioPolicy {
self.stdout
}
#[must_use]
pub const fn stderr_policy(&self) -> StdioPolicy {
self.stderr
}
#[must_use]
pub const fn deadline_duration(&self) -> Option<Duration> {
self.deadline
}
#[must_use]
pub(crate) const fn stdout_capture(&self) -> Option<NonZeroUsize> {
match self.stdout {
StdioPolicy::Capture(limit) => Some(limit),
StdioPolicy::Inherit | StdioPolicy::Null => None,
}
}
#[must_use]
pub(crate) const fn stderr_capture(&self) -> Option<NonZeroUsize> {
match self.stderr {
StdioPolicy::Capture(limit) => Some(limit),
StdioPolicy::Inherit | StdioPolicy::Null => None,
}
}
pub(crate) fn configure(&self, command: &mut Command) -> io::Result<()> {
command.args(&self.args);
for delta in &self.env {
match *delta {
EnvDelta::Set { ref key, ref value } => {
command.env(key, value);
}
EnvDelta::Remove { ref key } => {
command.env_remove(key);
}
EnvDelta::Clear => {
command.env_clear();
}
}
}
if let Some(cwd) = self.cwd.as_ref() {
command.current_dir(cwd);
}
command.stdin(self.stdin.into_stdio());
command.stdout(match self.stdout_file.as_deref() {
Some(path) => file_stdio(path)?,
None => self.stdout.into_stdio(),
});
command.stderr(match self.stderr_file.as_deref() {
Some(path) => file_stdio(path)?,
None => self.stderr.into_stdio(),
});
Ok(())
}
}
fn file_stdio(path: &Path) -> io::Result<std::process::Stdio> {
let file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path)?;
Ok(std::process::Stdio::from(file))
}
impl StdioPolicy {
fn into_stdio(self) -> std::process::Stdio {
match self {
Self::Inherit => std::process::Stdio::inherit(),
Self::Null => std::process::Stdio::null(),
Self::Capture(_) => std::process::Stdio::piped(),
}
}
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
#[non_exhaustive]
pub struct CapturedStream {
bytes: Vec<u8>,
total_bytes: u64,
truncated: bool,
}
impl CapturedStream {
#[must_use]
pub fn bytes(&self) -> &[u8] {
&self.bytes
}
#[must_use]
pub const fn total_bytes(&self) -> u64 {
self.total_bytes
}
#[must_use]
pub const fn truncated(&self) -> bool {
self.truncated
}
#[must_use]
pub fn retained_capacity(&self) -> usize {
self.bytes.capacity()
}
#[must_use]
pub fn frames(&self, ceiling: usize) -> Frames {
let mut pass = Frames::of_slice(self.bytes.as_slice(), ceiling);
if self.truncated
&& matches!(
pass.ended,
FrameRead::EndOfStream
| FrameRead::TruncatedPrefix { .. }
| FrameRead::TruncatedPayload { .. }
)
{
pass.ended = FrameRead::CeilingReached {
ceiling: self.retained_capacity(),
};
}
pass
}
#[cfg(unix)]
pub(crate) fn from_parts(bytes: Vec<u8>, total_bytes: u64, truncated: bool) -> Self {
Self {
bytes,
total_bytes,
truncated,
}
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct ProcessRun {
pub(crate) status: Option<ExitStatus>,
pub(crate) deadline_fired: bool,
pub(crate) stdout: CapturedStream,
pub(crate) stderr: CapturedStream,
pub(crate) cleanup: CleanupReceipt,
pub(crate) containment: Containment,
}
impl ProcessRun {
#[must_use]
pub const fn status(&self) -> Option<ExitStatus> {
self.status
}
#[must_use]
pub fn exit_code(&self) -> Option<i32> {
self.status.and_then(|status| status.code())
}
#[cfg(unix)]
#[must_use]
pub fn signal(&self) -> Option<i32> {
use std::os::unix::process::ExitStatusExt;
self.status.and_then(|status| status.signal())
}
#[must_use]
pub const fn deadline_fired(&self) -> bool {
self.deadline_fired
}
#[must_use]
pub const fn stdout(&self) -> &CapturedStream {
&self.stdout
}
#[must_use]
pub const fn stderr(&self) -> &CapturedStream {
&self.stderr
}
#[must_use]
pub const fn cleanup(&self) -> &CleanupReceipt {
&self.cleanup
}
#[must_use]
pub const fn containment(&self) -> &Containment {
&self.containment
}
}
pub const DEFAULT_FRAME_CEILING: usize = 64 * 1024;
#[derive(Clone, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum FrameRead {
Frame {
declared: usize,
payload: Vec<u8>,
},
EndOfStream,
TruncatedPrefix {
partial: Vec<u8>,
},
TruncatedPayload {
declared: usize,
partial: Vec<u8>,
},
MalformedPrefix {
declared: usize,
ceiling: usize,
},
CeilingReached {
ceiling: usize,
},
}
impl FrameRead {
#[must_use]
pub const fn is_frame(&self) -> bool {
matches!(*self, Self::Frame { .. })
}
#[must_use]
pub fn payload(&self) -> Option<&[u8]> {
match *self {
Self::Frame { ref payload, .. } => Some(payload.as_slice()),
Self::EndOfStream
| Self::TruncatedPrefix { .. }
| Self::TruncatedPayload { .. }
| Self::MalformedPrefix { .. }
| Self::CeilingReached { .. } => None,
}
}
#[must_use]
pub const fn is_refusal(&self) -> bool {
matches!(
*self,
Self::TruncatedPrefix { .. }
| Self::TruncatedPayload { .. }
| Self::MalformedPrefix { .. }
| Self::CeilingReached { .. }
)
}
#[must_use]
pub const fn is_truncated(&self) -> bool {
matches!(
*self,
Self::TruncatedPrefix { .. } | Self::TruncatedPayload { .. }
)
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct FrameReadError {
source: io::Error,
frames: usize,
payload_bytes: usize,
}
impl FrameReadError {
#[must_use]
pub const fn frames(&self) -> usize {
self.frames
}
#[must_use]
pub const fn payload_bytes(&self) -> usize {
self.payload_bytes
}
#[must_use]
pub fn kind(&self) -> io::ErrorKind {
self.source.kind()
}
}
impl fmt::Display for FrameReadError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
formatter,
"reading framed output failed after {} records and {} payload bytes: {}",
self.frames, self.payload_bytes, self.source
)
}
}
impl std::error::Error for FrameReadError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(&self.source)
}
}
#[must_use = "the records and the reason the stream stopped are the whole result"]
pub fn read_frames<R: std::io::Read>(
stream: &mut R,
ceiling: usize,
) -> Result<Frames, FrameReadError> {
Frames::read(stream, ceiling)
}
#[derive(Clone, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub struct Frames {
records: Vec<FrameRead>,
ended: FrameRead,
retained_bytes: usize,
}
impl Frames {
pub fn read<R: std::io::Read>(stream: &mut R, ceiling: usize) -> Result<Self, FrameReadError> {
let pass = read_pass(stream, ceiling);
match pass {
Ok((records, ended, retained_bytes)) => Ok(Self {
records,
ended,
retained_bytes,
}),
Err(Failed {
source,
records,
retained_bytes,
}) => Err(FrameReadError {
source,
frames: records.len(),
payload_bytes: retained_bytes,
}),
}
}
fn of_slice(bytes: &[u8], ceiling: usize) -> Self {
let mut stream = bytes;
match read_pass(&mut stream, ceiling) {
Ok((records, ended, retained_bytes)) => Self {
records,
ended,
retained_bytes,
},
Err(Failed {
records,
retained_bytes,
..
}) => Self {
records,
ended: FrameRead::CeilingReached {
ceiling: retained_bytes,
},
retained_bytes,
},
}
}
#[must_use]
pub fn records(&self) -> &[FrameRead] {
&self.records
}
#[must_use]
pub const fn ended(&self) -> &FrameRead {
&self.ended
}
#[must_use]
pub const fn retained_bytes(&self) -> usize {
self.retained_bytes
}
#[must_use]
pub fn is_complete(&self) -> bool {
matches!(self.ended, FrameRead::EndOfStream)
}
}
#[derive(Debug)]
struct Failed {
source: io::Error,
records: Vec<FrameRead>,
retained_bytes: usize,
}
fn read_pass<R: std::io::Read>(
stream: &mut R,
ceiling: usize,
) -> Result<(Vec<FrameRead>, FrameRead, usize), Failed> {
let mut records: Vec<FrameRead> = Vec::new();
let mut retained = 0_usize;
let ended = loop {
match read_one(stream, ceiling, &mut retained, &mut records) {
Ok(Some(ended)) => break ended,
Ok(None) => {}
Err(source) => {
let refusal = Err(Failed {
source,
records,
retained_bytes: retained,
});
lgwks_std::trace::debug!(error = ?refusal.as_ref().err(), "read_pass: returning an error to the caller");
return refusal;
}
}
};
Ok((records, ended, retained))
}
fn read_one<R: std::io::Read>(
stream: &mut R,
ceiling: usize,
retained: &mut usize,
records: &mut Vec<FrameRead>,
) -> Result<Option<FrameRead>, io::Error> {
if *retained >= ceiling {
return Ok(Some(FrameRead::CeilingReached { ceiling }));
}
let mut prefix = [0_u8; LENGTH_BYTES];
match read_counted(stream, &mut prefix)? {
0 => return Ok(Some(FrameRead::EndOfStream)),
filled if filled < LENGTH_BYTES => {
return Ok(Some(FrameRead::TruncatedPrefix {
partial: prefix[..filled].to_vec(),
}));
}
_ => {}
}
let declared = declared_length(&prefix);
if !is_possible_length(declared, ceiling) {
return Ok(Some(FrameRead::MalformedPrefix { declared, ceiling }));
}
let remaining = ceiling.saturating_sub(*retained);
if declared > remaining {
return Ok(Some(FrameRead::CeilingReached { ceiling }));
}
let mut payload: Vec<u8> = vec![0; declared];
let read = read_counted(stream, &mut payload)?;
*retained = retained.saturating_add(declared);
if read < declared {
payload.truncate(read);
return Ok(Some(FrameRead::TruncatedPayload {
declared,
partial: payload,
}));
}
records.push(FrameRead::Frame { declared, payload });
Ok(None)
}
fn read_counted<R: std::io::Read>(stream: &mut R, buf: &mut [u8]) -> Result<usize, io::Error> {
match crate::journal::frame::read_exact_or_eof(stream, buf) {
Ok(None) => Ok(0),
Ok(Some(filled)) => Ok(filled),
Err(crate::journal::JournalError::Storage(source)) => Err(source),
Err(other) => Err(io::Error::other(other.to_string())),
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum ProcessRunError {
Refused,
NotStarted {
source: io::Error,
},
AfterStart {
source: io::Error,
},
}
impl fmt::Display for ProcessRunError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
Self::Refused => formatter.write_str("the supervisor refused to start a process"),
Self::NotStarted { ref source } => {
write!(formatter, "the program did not start: {source}")
}
Self::AfterStart { ref source } => {
write!(
formatter,
"the process started but did not settle: {source}"
)
}
}
}
}
impl std::error::Error for ProcessRunError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match *self {
Self::Refused => None,
Self::NotStarted { ref source } | Self::AfterStart { ref source } => Some(source),
}
}
}