Hardswap Invariants in tests + Quota fixes

src/maestro/control_plane.rs

@@ -120,6 +120,19 @@ impl ProcessControlPlane {
         Ok(())
     }
 
+    /// Registers work that must finish once accepted, even after shutdown cancellation starts.
+    pub(crate) fn spawn_completion<F>(&self, future: F) -> Result<(), F>
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        let Some(registration) = self.inner.admission.try_register() else {
+            return Err(future);
+        };
+        self.inner.tasks.spawn(future);
+        drop(registration);
+        Ok(())
+    }
+
     /// Closes task admission, cancels all owned work, and joins it within the deadline.
     pub(crate) async fn shutdown(&self, timeout: Duration) -> bool {
         let deadline = tokio::time::Instant::now() + timeout;
@@ -214,4 +227,31 @@ mod tests {
 
         assert!(scope.shutdown(Duration::from_secs(1)).await);
     }
+
+    #[tokio::test]
+    async fn shutdown_waits_for_accepted_completion_without_cancelling_it() {
+        let scope = ProcessControlPlane::new();
+        let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+        let completed = Arc::new(AtomicBool::new(false));
+        let completed_task = completed.clone();
+        assert!(
+            scope
+                .spawn_completion(async move {
+                    let _ = release_rx.await;
+                    completed_task.store(true, Ordering::Release);
+                })
+                .is_ok()
+        );
+
+        let shutdown_scope = scope.clone();
+        let shutdown =
+            tokio::spawn(async move { shutdown_scope.shutdown(Duration::from_secs(1)).await });
+        tokio::task::yield_now().await;
+        assert!(!shutdown.is_finished());
+        assert!(!completed.load(Ordering::Acquire));
+
+        release_tx.send(()).unwrap();
+        assert!(shutdown.await.unwrap());
+        assert!(completed.load(Ordering::Acquire));
+    }
 }