Inline process background future, clean tracing a bit

src/generate.rs

@@ -47,7 +47,7 @@ impl Drop for MetricsGuard {
 }
 
 #[allow(clippy::too_many_arguments)]
-#[tracing::instrument(skip(repo, store, hash, process_map, config))]
+#[tracing::instrument(skip(tmp_dir, repo, store, hash, process_map, config))]
 pub(crate) async fn generate<S: Store + 'static>(
     tmp_dir: &TmpDir,
     repo: &ArcRepo,

src/ingest.rs

@@ -144,7 +144,7 @@ where
     Ok((input_type, identifier, details, state))
 }
 
-#[tracing::instrument(skip(repo, store, client, stream, config))]
+#[tracing::instrument(skip(tmp_dir, repo, store, client, stream, config))]
 pub(crate) async fn ingest<S>(
     tmp_dir: &TmpDir,
     repo: &ArcRepo,

src/lib.rs

@@ -866,7 +866,7 @@ async fn not_found_hash(repo: &ArcRepo) -> Result<Option<(Alias, Hash)>, Error>
 #[allow(clippy::too_many_arguments)]
 #[tracing::instrument(
     name = "Serving processed image",
-    skip(repo, store, client, config, process_map)
+    skip(tmp_dir, repo, store, client, config, process_map)
 )]
 async fn process<S: Store + 'static>(
     range: Option<web::Header<Range>>,

src/process.rs

@@ -14,21 +14,25 @@ use std::{
 use tokio::{
     io::{AsyncRead, AsyncWriteExt, ReadBuf},
     process::{Child, ChildStdin, ChildStdout, Command},
-    task::JoinHandle,
 };
 use tracing::{Instrument, Span};
+use uuid::Uuid;
 
-use crate::{error_code::ErrorCode, future::WithTimeout};
+use crate::{
+    error_code::ErrorCode,
+    future::{LocalBoxFuture, WithTimeout},
+    sync::DropHandle,
+};
 
 struct MetricsGuard {
     start: Instant,
     armed: bool,
-    command: String,
+    command: Arc<str>,
 }
 
 impl MetricsGuard {
-    fn guard(command: String) -> Self {
+    fn guard(command: Arc<str>) -> Self {
-        metrics::increment_counter!("pict-rs.process.start", "command" => command.clone());
+        metrics::increment_counter!("pict-rs.process.start", "command" => command.to_string());
 
         Self {
             start: Instant::now(),
@@ -47,11 +51,11 @@ impl Drop for MetricsGuard {
         metrics::histogram!(
             "pict-rs.process.duration",
             self.start.elapsed().as_secs_f64(),
-            "command" => self.command.clone(),
+            "command" => self.command.to_string(),
             "completed" => (!self.armed).to_string(),
         );
 
-        metrics::increment_counter!("pict-rs.process.end", "completed" => (!self.armed).to_string() , "command" => self.command.clone());
+        metrics::increment_counter!("pict-rs.process.end", "completed" => (!self.armed).to_string() , "command" => self.command.to_string());
     }
 }
 
@@ -59,10 +63,11 @@ impl Drop for MetricsGuard {
 struct StatusError(ExitStatus);
 
 pub(crate) struct Process {
-    command: String,
+    command: Arc<str>,
     child: Child,
     guard: MetricsGuard,
     timeout: Duration,
+    id: Uuid,
 }
 
 impl std::fmt::Debug for Process {
@@ -71,10 +76,6 @@ impl std::fmt::Debug for Process {
     }
 }
 
-struct DropHandle {
-    inner: JoinHandle<std::io::Result<()>>,
-}
-
 struct ProcessReadState {
     flags: AtomicU8,
     parent: Mutex<Option<Waker>>,
@@ -86,29 +87,31 @@ struct ProcessReadWaker {
 }
 
 pub(crate) struct ProcessRead {
-    inner: ChildStdout,
+    stdout: ChildStdout,
-    handle: DropHandle,
+    handle: LocalBoxFuture<'static, std::io::Result<()>>,
     closed: bool,
     state: Arc<ProcessReadState>,
-    span: Span,
+    span: Option<Span>,
+    command: Arc<str>,
+    id: Uuid,
 }
 
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum ProcessError {
     #[error("Required command {0} not found, make sure it exists in pict-rs' $PATH")]
-    NotFound(String),
+    NotFound(Arc<str>),
 
     #[error("Cannot run command {0} due to invalid permissions on binary, make sure the pict-rs user has permission to run it")]
-    PermissionDenied(String),
+    PermissionDenied(Arc<str>),
 
     #[error("Reached process spawn limit")]
     LimitReached,
 
     #[error("{0} timed out")]
-    Timeout(String),
+    Timeout(Arc<str>),
 
     #[error("{0} Failed with {1}")]
-    Status(String, ExitStatus),
+    Status(Arc<str>, ExitStatus),
 
     #[error("Unknown process error")]
     Other(#[source] std::io::Error),
@@ -136,10 +139,14 @@ impl Process {
     where
         T: AsRef<OsStr>,
     {
+        let command: Arc<str> = Arc::from(String::from(command));
+
         let res = tracing::trace_span!(parent: None, "Create command", %command).in_scope(|| {
             Self::spawn(
-                command,
+                command.clone(),
-                Command::new(command).args(args).envs(envs.iter().copied()),
+                Command::new(command.as_ref())
+                    .args(args)
+                    .envs(envs.iter().copied()),
                 timeout,
             )
         });
@@ -147,9 +154,9 @@ impl Process {
         match res {
             Ok(this) => Ok(this),
             Err(e) => match e.kind() {
-                std::io::ErrorKind::NotFound => Err(ProcessError::NotFound(command.to_string())),
+                std::io::ErrorKind::NotFound => Err(ProcessError::NotFound(command)),
                 std::io::ErrorKind::PermissionDenied => {
-                    Err(ProcessError::PermissionDenied(command.to_string()))
+                    Err(ProcessError::PermissionDenied(command))
                 }
                 std::io::ErrorKind::WouldBlock => Err(ProcessError::LimitReached),
                 _ => Err(ProcessError::Other(e)),
@@ -157,9 +164,9 @@ impl Process {
         }
     }
 
-    fn spawn(command: &str, cmd: &mut Command, timeout: u64) -> std::io::Result<Self> {
+    fn spawn(command: Arc<str>, cmd: &mut Command, timeout: u64) -> std::io::Result<Self> {
         tracing::trace_span!(parent: None, "Spawn command", %command).in_scope(|| {
-            let guard = MetricsGuard::guard(command.into());
+            let guard = MetricsGuard::guard(command.clone());
 
             let cmd = cmd
                 .stdin(Stdio::piped())
@@ -168,20 +175,22 @@ impl Process {
 
             cmd.spawn().map(|child| Process {
                 child,
-                command: String::from(command),
+                command,
                 guard,
                 timeout: Duration::from_secs(timeout),
+                id: Uuid::now_v7(),
             })
         })
     }
 
-    #[tracing::instrument(skip(self))]
+    #[tracing::instrument(skip(self), fields(command = %self.command, id = %self.id))]
     pub(crate) async fn wait(self) -> Result<(), ProcessError> {
         let Process {
             command,
             mut child,
             guard,
             timeout,
+            id: _,
         } = self;
 
         let res = child.wait().with_timeout(timeout).await;
@@ -226,21 +235,13 @@ impl Process {
             mut child,
             guard,
             timeout,
+            id,
         } = self;
 
         let stdin = child.stdin.take().expect("stdin exists");
         let stdout = child.stdout.take().expect("stdout exists");
 
-        let background_span =
+        let handle = Box::pin(async move {
-            tracing::info_span!(parent: None, "Background process task", %command);
-        background_span.follows_from(Span::current());
-
-        let span = tracing::info_span!(parent: None, "Foreground process task", %command);
-        span.follows_from(Span::current());
-
-        let handle = crate::sync::spawn(
-            "await-process",
-            async move {
             let child_fut = async {
                 (f)(stdin).await?;
 
@@ -262,16 +263,16 @@ impl Process {
             child.kill().await?;
 
             Err(error)
-            }
+        });
-            .instrument(background_span),
-        );
 
         ProcessRead {
-            inner: stdout,
+            stdout,
-            handle: DropHandle { inner: handle },
+            handle,
             closed: false,
             state: ProcessReadState::new_woken(),
-            span,
+            span: None,
+            command,
+            id,
         }
     }
 }
@@ -345,14 +346,19 @@ impl AsyncRead for ProcessRead {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<std::io::Result<()>> {
-        let span = self.span.clone();
+        let command = self.command.clone();
+        let id = self.id;
+        let span = self
+            .span
+            .get_or_insert_with(|| tracing::info_span!("process task", %command, %id))
+            .clone();
         let guard = span.enter();
 
         let value = loop {
             // always poll for bytes when poll_read is called
             let before_size = buf.filled().len();
 
-            if let Poll::Ready(res) = Pin::new(&mut self.inner).poll_read(cx, buf) {
+            if let Poll::Ready(res) = Pin::new(&mut self.stdout).poll_read(cx, buf) {
                 if let Err(e) = res {
                     self.closed = true;
 
@@ -368,11 +374,10 @@ impl AsyncRead for ProcessRead {
                 // only poll handle if we've been explicitly woken
                 let mut handle_cx = Context::from_waker(&waker);
 
-                if let Poll::Ready(res) = Pin::new(&mut self.handle.inner).poll(&mut handle_cx) {
+                if let Poll::Ready(res) = Pin::new(&mut self.handle).poll(&mut handle_cx) {
                     let error = match res {
-                        Ok(Ok(())) => continue,
+                        Ok(()) => continue,
-                        Ok(Err(e)) => e,
+                        Err(e) => e,
-                        Err(e) => std::io::Error::new(std::io::ErrorKind::Other, e),
                     };
 
                     self.closed = true;
@@ -398,12 +403,6 @@ impl AsyncRead for ProcessRead {
     }
 }
 
-impl Drop for DropHandle {
-    fn drop(&mut self) {
-        self.inner.abort();
-    }
-}
-
 impl Wake for ProcessReadWaker {
     fn wake(self: Arc<Self>) {
         match Arc::try_unwrap(self) {

src/queue/cleanup.rs

@@ -236,7 +236,7 @@ where
 
         let span = tracing::info_span!(parent: current_span, "error-boundary");
 
-        let res = crate::sync::spawn(
+        let res = crate::sync::abort_on_drop(crate::sync::spawn(
             "prune-missing",
             async move {
                 let mut count = count;
@@ -261,7 +261,7 @@ where
                 Ok(count) as Result<u64, Error>
             }
             .instrument(span),
-        )
+        ))
         .await;
 
         match res {

src/queue/process.rs

@@ -1,5 +1,6 @@
 use reqwest_middleware::ClientWithMiddleware;
 use time::Instant;
+use tracing::{Instrument, Span};
 
 use crate::{
     concurrent_processor::ProcessMap,
@@ -135,8 +136,13 @@ where
         let repo = repo.clone();
         let client = client.clone();
 
+        let current_span = Span::current();
+        let span = tracing::info_span!(parent: current_span, "error_boundary");
+
         let config = config.clone();
-        let error_boundary = crate::sync::spawn("ingest-media", async move {
+        let error_boundary = crate::sync::abort_on_drop(crate::sync::spawn(
+            "ingest-media",
+            async move {
                 let stream = crate::stream::from_err(store2.to_stream(&ident, None, None).await?);
 
                 let session = crate::ingest::ingest(
@@ -151,7 +157,9 @@ where
                 .await?;
 
                 Ok(session) as Result<Session, Error>
-        })
+            }
+            .instrument(span),
+        ))
         .await;
 
         store.remove(&unprocessed_identifier).await?;

src/repo/postgres.rs

@@ -33,6 +33,7 @@ use crate::{
     future::{WithMetrics, WithTimeout},
     serde_str::Serde,
     stream::LocalBoxStream,
+    sync::DropHandle,
 };
 
 use self::job_status::JobStatus;
@@ -49,7 +50,7 @@ use super::{
 pub(crate) struct PostgresRepo {
     inner: Arc<Inner>,
     #[allow(dead_code)]
-    notifications: Arc<tokio::task::JoinHandle<()>>,
+    notifications: Arc<DropHandle<()>>,
 }
 
 struct Inner {
@@ -151,7 +152,7 @@ impl PostgresRepo {
             .await
             .map_err(ConnectPostgresError::ConnectForMigration)?;
 
-        let handle = crate::sync::spawn("postgres-migrations", conn);
+        let handle = crate::sync::abort_on_drop(crate::sync::spawn("postgres-migrations", conn));
 
         embedded::migrations::runner()
             .run_async(&mut client)
@@ -199,10 +200,10 @@ impl PostgresRepo {
             upload_notifications: DashMap::new(),
         });
 
-        let handle = crate::sync::spawn(
+        let handle = crate::sync::abort_on_drop(crate::sync::spawn(
             "postgres-delegate-notifications",
             delegate_notifications(rx, inner.clone(), parallelism * 8),
-        );
+        ));
 
         let notifications = Arc::new(handle);
 

src/store/object_store.rs

@@ -298,7 +298,7 @@ impl Store for ObjectStore {
 
                 let object_id2 = object_id.clone();
                 let upload_id2 = upload_id.to_string();
-                let handle = crate::sync::spawn(
+                let handle = crate::sync::abort_on_drop(crate::sync::spawn(
                     "upload-multipart-part",
                     async move {
                         let response = this
@@ -333,7 +333,7 @@ impl Store for ObjectStore {
                         Ok(etag) as Result<String, StoreError>
                     }
                     .instrument(tracing::Span::current()),
-                );
+                ));
 
                 futures.push(handle);
             }

src/stream.rs

@@ -30,7 +30,7 @@ where
 {
     let (tx, rx) = crate::sync::channel(1);
 
-    let handle = crate::sync::spawn("send-stream", async move {
+    let handle = crate::sync::abort_on_drop(crate::sync::spawn("send-stream", async move {
         let stream = std::pin::pin!(stream);
         let mut streamer = stream.into_streamer();
 
@@ -39,7 +39,7 @@ where
                 break;
             }
         }
-    });
+    }));
 
     streem::from_fn(|yiedler| async move {
         let mut stream = rx.into_stream().into_streamer();

src/sync.rs

@@ -1,6 +1,40 @@
 use std::sync::Arc;
 
-use tokio::sync::{Notify, Semaphore};
+use tokio::{
+    sync::{Notify, Semaphore},
+    task::JoinHandle,
+};
+
+pub(crate) struct DropHandle<T> {
+    handle: JoinHandle<T>,
+}
+
+pub(crate) fn abort_on_drop<T>(handle: JoinHandle<T>) -> DropHandle<T> {
+    DropHandle { handle }
+}
+
+impl<T> DropHandle<T> {
+    pub(crate) fn abort(&self) {
+        self.handle.abort();
+    }
+}
+
+impl<T> Drop for DropHandle<T> {
+    fn drop(&mut self) {
+        self.handle.abort();
+    }
+}
+
+impl<T> std::future::Future for DropHandle<T> {
+    type Output = <JoinHandle<T> as std::future::Future>::Output;
+
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        std::pin::Pin::new(&mut self.handle).poll(cx)
+    }
+}
 
 #[track_caller]
 pub(crate) fn channel<T>(bound: usize) -> (flume::Sender<T>, flume::Receiver<T>) {