use crate::{id_or_span::IdOrSpan, store::Store}; use actix_rt::task::JoinHandle; use actix_web::web::Bytes; use std::{ future::Future, pin::Pin, process::Stdio, task::{Context, Poll}, }; use tokio::{ io::{AsyncRead, AsyncWriteExt, ReadBuf}, process::{Child, Command}, sync::oneshot::{channel, Receiver}, }; use tracing::Instrument; use tracing::{Id, Span}; #[derive(Debug)] struct StatusError; pub(crate) struct Process { child: Child, id: Option, } struct DropHandle { inner: JoinHandle<()>, } pin_project_lite::pin_project! { struct ProcessRead { #[pin] inner: I, span: IdOrSpan, err_recv: Receiver, err_closed: bool, handle: DropHandle, } } impl Process { pub(crate) fn run(command: &str, args: &[&str]) -> std::io::Result { Self::spawn(Command::new(command).args(args)) } fn spawn_span(&self) -> Span { let span = tracing::info_span!(parent: None, "Spawned command writer",); span.follows_from(self.id.clone()); span } #[tracing::instrument(name = "Spawning Command")] pub(crate) fn spawn(cmd: &mut Command) -> std::io::Result { let cmd = cmd.stdin(Stdio::piped()).stdout(Stdio::piped()); cmd.spawn().map(|child| Process { child, id: Span::current().id(), }) } pub(crate) async fn wait(mut self) -> std::io::Result<()> { let status = self.child.wait().await?; if !status.success() { return Err(std::io::Error::new(std::io::ErrorKind::Other, &StatusError)); } Ok(()) } pub(crate) fn bytes_read(mut self, mut input: Bytes) -> Option { let mut stdin = self.child.stdin.take()?; let stdout = self.child.stdout.take()?; let (tx, rx) = channel::(); let span = self.spawn_span(); let mut child = self.child; let handle = actix_rt::spawn( async move { if let Err(e) = stdin.write_all_buf(&mut input).await { let _ = tx.send(e); return; } drop(stdin); match child.wait().await { Ok(status) => { if !status.success() { let _ = tx .send(std::io::Error::new(std::io::ErrorKind::Other, &StatusError)); } } Err(e) => { let _ = tx.send(e); } } } .instrument(span), ); Some(ProcessRead { inner: stdout, span: IdOrSpan::from_id(self.id), err_recv: rx, err_closed: false, handle: DropHandle { inner: handle }, }) } pub(crate) fn read(mut self) -> Option { let stdout = self.child.stdout.take()?; let (tx, rx) = channel(); let span = self.spawn_span(); let mut child = self.child; let handle = actix_rt::spawn( async move { match child.wait().await { Ok(status) => { if !status.success() { let _ = tx .send(std::io::Error::new(std::io::ErrorKind::Other, &StatusError)); } } Err(e) => { let _ = tx.send(e); } } } .instrument(span), ); Some(ProcessRead { inner: stdout, span: IdOrSpan::from_id(self.id), err_recv: rx, err_closed: false, handle: DropHandle { inner: handle }, }) } pub(crate) fn store_read( mut self, store: S, identifier: S::Identifier, ) -> Option { let mut stdin = self.child.stdin.take()?; let stdout = self.child.stdout.take()?; let (tx, rx) = channel(); let span = self.spawn_span(); let mut child = self.child; let handle = actix_rt::spawn( async move { if let Err(e) = store.read_into(&identifier, &mut stdin).await { let _ = tx.send(e); return; } drop(stdin); match child.wait().await { Ok(status) => { if !status.success() { let _ = tx .send(std::io::Error::new(std::io::ErrorKind::Other, &StatusError)); } } Err(e) => { let _ = tx.send(e); } } } .instrument(span), ); Some(ProcessRead { inner: stdout, span: IdOrSpan::from_id(self.id), err_recv: rx, err_closed: false, handle: DropHandle { inner: handle }, }) } } impl AsyncRead for ProcessRead where I: AsyncRead, { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let this = self.as_mut().project(); let err_recv = this.err_recv; let err_closed = this.err_closed; let inner = this.inner; let span = this.span.as_span(|id| { let span = tracing::info_span!("Processing Command"); span.follows_from(id); span }); span.in_scope(|| { if !*err_closed { if let Poll::Ready(res) = Pin::new(err_recv).poll(cx) { *err_closed = true; if let Ok(err) = res { let display = format!("{}", err); let debug = format!("{:?}", err); span.record("exception.message", &display.as_str()); span.record("exception.details", &debug.as_str()); return Poll::Ready(Err(err)); } } } if let Poll::Ready(res) = inner.poll_read(cx, buf) { if let Err(err) = &res { let display = format!("{}", err); let debug = format!("{:?}", err); span.record("exception.message", &display.as_str()); span.record("exception.details", &debug.as_str()); } return Poll::Ready(res); } Poll::Pending }) } } impl Drop for DropHandle { fn drop(&mut self) { self.inner.abort(); } } impl std::fmt::Display for StatusError { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "Command failed with bad status") } } impl std::error::Error for StatusError {}