Prevent deadlock due to lack of approve_loaded()

src/units/zone_server.rs

@@ -46,8 +46,8 @@ use crate::manager::Terminated;
 use crate::manager::record_zone_event;
 use crate::util::AbortOnDrop;
 use crate::zone::{
-    HistoricalEvent, SignedZoneVersionState, SigningTrigger, UnsignedZoneVersionState, ZoneHandle,
-    ZoneVersionReviewState,
+    HistoricalEvent, SignedZoneVersionState, SigningTrigger, UnsignedZoneVersionState, Zone,
+    ZoneHandle, ZoneVersionReviewState,
 };
 
 /// The source of a zone server.
@@ -367,7 +367,7 @@ impl ZoneServer {
                             "[{unit_name}]: Cannot promote unsigned zone '{zone_name}' to the signable set of zones: {err}"
                         );
                     } else {
-                        self.on_unsigned_zone_approved(center, zone_name, zone_serial);
+                        self.on_unsigned_zone_approved(center, &zone, zone_serial);
                     }
                 }
                 Source::Signed => {
@@ -526,21 +526,29 @@ impl ZoneServer {
     fn on_unsigned_zone_approved(
         &self,
         center: &Arc<Center>,
-        zone_name: Name<Bytes>,
+        zone: &Arc<Zone>,
         zone_serial: Serial,
     ) {
-        record_zone_event(
-            center,
-            &zone_name,
-            HistoricalEvent::UnsignedZoneReview {
-                status: crate::api::ZoneReviewStatus::Approved,
-            },
-            Some(zone_serial),
-        );
+        {
+            let mut zone_state = zone.state.lock().unwrap();
+            zone_state.record_event(
+                HistoricalEvent::UnsignedZoneReview {
+                    status: crate::api::ZoneReviewStatus::Approved,
+                },
+                Some(zone_serial),
+            );
+            ZoneHandle {
+                zone,
+                state: &mut zone_state,
+                center,
+            }
+            .storage()
+            .approve_loaded();
+        }
         info!("[CC]: Instructing zone signer to sign the approved zone");
         center.signer.on_sign_zone(
             center,
-            zone_name,
+            zone.name.clone(),
             Some(zone_serial),
             SigningTrigger::ZoneChangesApproved,
         );
@@ -637,24 +645,14 @@ impl ZoneServer {
                     }
 
                     version.review = new_review_state;
-
-                    if matches!(decision, ZoneReviewDecision::Approve) {
-                        ZoneHandle {
-                            zone: &zone,
-                            state: &mut zone_state,
-                            center,
-                        }
-                        .storage()
-                        .approve_loaded();
-                    }
                 }
                 if matches!(decision, ZoneReviewDecision::Approve) {
                     info!(
                         "Unsigned zone '{zone_name}' with serial {zone_serial} has been approved."
                     );
                     match Self::promote_zone_to_signable(center.clone(), &zone_name) {
                         Ok(()) => {
-                            self.on_unsigned_zone_approved(center, zone_name, zone_serial);
+                            self.on_unsigned_zone_approved(center, &zone, zone_serial);
                         }
                         Err(err) => {
                             error!(

src/util.rs

@@ -15,6 +15,7 @@ use tokio::{
     task::{AbortHandle, JoinHandle},
     time::Instant,
 };
+use tracing::{Instrument, trace};
 
 //----------- AbortOnDrop ------------------------------------------------------
 
@@ -49,6 +50,71 @@ impl AbortOnDrop {
     }
 }
 
+//----------- BackgroundTasks --------------------------------------------------
+
+/// A monitor for background tasks.
+///
+/// This allows tracking background tasks for a particular purpose.
+///
+/// If a [`BackgroundTasks`] is dropped, all associated tasks are aborted.
+#[derive(Default)]
+pub struct BackgroundTasks {
+    /// The underlying tasks.
+    tasks: foldhash::HashMap<tokio::task::Id, JoinHandle<()>>,
+}
+
+impl BackgroundTasks {
+    /// Spawn a new background task.
+    pub fn spawn<F>(&mut self, span: tracing::Span, f: F)
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        let task = tokio::task::spawn(f.instrument(span));
+        let other = self.tasks.insert(task.id(), task);
+        assert!(
+            other.is_none(),
+            "Task IDs for two live 'JoinHandles' never collide"
+        );
+    }
+
+    /// Spawn a new blocking background task.
+    pub fn spawn_blocking<F>(&mut self, span: tracing::Span, f: F)
+    where
+        F: FnOnce() + Send + 'static,
+    {
+        let task = tokio::task::spawn_blocking(move || span.in_scope(f));
+        let other = self.tasks.insert(task.id(), task);
+        assert!(
+            other.is_none(),
+            "Task IDs for two live 'JoinHandles' never collide"
+        );
+    }
+
+    /// Mark the running task as finished.
+    #[tracing::instrument(level = "trace", skip_all)]
+    pub fn finish(&mut self) {
+        let id = tokio::task::id();
+        match self.tasks.remove(&id) {
+            Some(handle) => {
+                trace!("Detaching background task");
+                std::mem::drop(handle);
+            }
+            None => {
+                trace!("Inconsistency: unknown task");
+            }
+        }
+    }
+}
+
+impl Drop for BackgroundTasks {
+    fn drop(&mut self) {
+        // Abort all tasks.
+        self.tasks.drain().for_each(|(_, handle)| handle.abort());
+    }
+}
+
+//------------------------------------------------------------------------------
+
 /// Force a [`Future`] to evaluate synchronously.
 pub fn force_future<F: IntoFuture>(future: F) -> F::Output {
     let waker = std::task::Waker::noop();