instrument the per-hash tasks as well

Author: rklaehn

src/store/fs.rs

@@ -570,9 +570,12 @@ impl Actor {
 }
 
 trait HashSpecificCommand: HashSpecific + Send + 'static {
+    /// Handle the command on success by spawning a task into the per-hash context.
     fn handle(self, ctx: HashContext) -> impl Future<Output = ()> + Send + 'static;
 
-    fn on_error(self) -> impl Future<Output = ()> + Send + 'static;
+    /// Opportunity to send an error if spawning fails due to the task being busy (inbox full)
+    /// or dead (e.g. panic in one of the running tasks).
+    fn on_error(self, arg: SpawnArg<EmParams>) -> impl Future<Output = ()> + Send + 'static;
 
     async fn spawn(
         self,
@@ -581,25 +584,24 @@ trait HashSpecificCommand: HashSpecific + Send + 'static {
     ) where
         Self: Sized,
     {
+        let span = tracing::Span::current();
         let task = manager
-            .spawn_boxed(
-                self.hash(),
-                Box::new(|x| {
-                    Box::pin(async move {
-                        match x {
-                            SpawnArg::Active(state) => {
-                                self.handle(state).await;
-                            }
-                            SpawnArg::Busy => {
-                                self.on_error().await;
-                            }
-                            SpawnArg::Dead => {
-                                self.on_error().await;
-                            }
+            .spawn(self.hash(), |arg| {
+                async move {
+                    match arg {
+                        SpawnArg::Active(state) => {
+                            self.handle(state).await;
                         }
-                    })
-                }),
-            )
+                        SpawnArg::Busy => {
+                            self.on_error(arg).await;
+                        }
+                        SpawnArg::Dead => {
+                            self.on_error(arg).await;
+                        }
+                    }
+                }
+                .instrument(span)
+            })
             .await;
         if let Some(task) = task {
             tasks.spawn(task);
@@ -611,31 +613,31 @@ impl HashSpecificCommand for ObserveMsg {
     async fn handle(self, ctx: HashContext) {
         observe(self, ctx).await
     }
-    async fn on_error(self) {}
+    async fn on_error(self, _arg: SpawnArg<EmParams>) {}
 }
 impl HashSpecificCommand for ExportPathMsg {
     async fn handle(self, ctx: HashContext) {
         export_path(self, ctx).await
     }
-    async fn on_error(self) {}
+    async fn on_error(self, _arg: SpawnArg<EmParams>) {}
 }
 impl HashSpecificCommand for ExportBaoMsg {
     async fn handle(self, ctx: HashContext) {
         export_bao(self, ctx).await
     }
-    async fn on_error(self) {}
+    async fn on_error(self, _arg: SpawnArg<EmParams>) {}
 }
 impl HashSpecificCommand for ExportRangesMsg {
     async fn handle(self, ctx: HashContext) {
         export_ranges(self, ctx).await
     }
-    async fn on_error(self) {}
+    async fn on_error(self, _arg: SpawnArg<EmParams>) {}
 }
 impl HashSpecificCommand for ImportBaoMsg {
     async fn handle(self, ctx: HashContext) {
         import_bao(self, ctx).await
     }
-    async fn on_error(self) {}
+    async fn on_error(self, _arg: SpawnArg<EmParams>) {}
 }
 impl HashSpecific for (TempTag, ImportEntryMsg) {
     fn hash(&self) -> Hash {
@@ -647,7 +649,7 @@ impl HashSpecificCommand for (TempTag, ImportEntryMsg) {
         let (tt, cmd) = self;
         finish_import(cmd, tt, ctx).await
     }
-    async fn on_error(self) {}
+    async fn on_error(self, _arg: SpawnArg<EmParams>) {}
 }
 
 struct RtWrapper(Option<tokio::runtime::Runtime>);

src/store/fs/bao_file.rs

@@ -285,6 +285,15 @@ fn read_size(size_file: &File) -> io::Result<u64> {
 }
 
 /// The storage for a bao file. This can be either in memory or on disk.
+///
+/// The two initial states `Initial` and `Loading` are used to coordinate the
+/// loading of the entry from the metadata database. Once that is complete,
+/// you should never see these states again.
+///
+/// From the remaining states you can get into `Poisoned` if there is an
+/// IO error during an operation.
+///
+/// `Poisioned` is also used once the handle is persisted and no longer usable.
 #[derive(derive_more::From, Default)]
 pub(crate) enum BaoFileStorage {
     /// Initial state, we don't know anything yet.
@@ -311,13 +320,8 @@ pub(crate) enum BaoFileStorage {
     ///
     /// Writing to this is a no-op, since it is already complete.
     Complete(CompleteStorage),
-    /// We will get into that state if there is an io error in the middle of an operation
-    ///
-    /// Also, when the handle is dropped we will poison the storage, so poisoned
-    /// can be seen when the handle is revived during the drop.
-    ///
-    /// BaoFileHandleWeak::upgrade() will return None if the storage is poisoned,
-    /// treat it as dead.
+    /// We will get into that state if there is an io error in the middle of an operation,
+    /// or after the handle is persisted and no longer usable.
     Poisoned,
 }
 

src/store/fs/util/entity_manager.rs

@@ -373,24 +373,25 @@ mod main_actor {
         }
 
         /// Friendly version of `spawn_boxed` that does the boxing
-        pub async fn spawn<F, Fut>(&mut self, id: P::EntityId, f: F, tasks: &mut JoinSet<()>)
+        #[must_use = "this function may return a future that must be spawned by the caller"]
+        pub async fn spawn<F, Fut>(
+            &mut self,
+            id: P::EntityId,
+            f: F,
+        ) -> Option<impl Future<Output = ()> + Send + 'static>
         where
             F: FnOnce(SpawnArg<P>) -> Fut + Send + 'static,
             Fut: Future<Output = ()> + Send + 'static,
         {
-            let task = self
-                .spawn_boxed(
-                    id,
-                    Box::new(|x| {
-                        Box::pin(async move {
-                            f(x).await;
-                        })
-                    }),
-                )
-                .await;
-            if let Some(task) = task {
-                tasks.spawn(task);
-            }
+            self.spawn_boxed(
+                id,
+                Box::new(|x| {
+                    Box::pin(async move {
+                        f(x).await;
+                    })
+                }),
+            )
+            .await
         }
 
         #[must_use = "this function may return a future that must be spawned by the caller"]