feat: Janitor, utilities for cleaning up with tokio

util/tests/janitor.rs

@@ -0,0 +1,85 @@
+#![cfg(feature = "tokio")]
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use tokio::time::sleep;
+
+use rosenpass_util::tokio::janitor::{enter_janitor, spawn_cleanup_job, try_spawn_daemon};
+
+#[tokio::test]
+async fn janitor_demo() -> anyhow::Result<()> {
+    let count = Arc::new(AtomicUsize::new(0));
+
+    // Make sure the program has access to an ambient janitor
+    {
+        let count = count.clone();
+        enter_janitor(async move {
+            let _drop_guard = AsyncDropDemo::new(count.clone()).await;
+
+            // Start a background job
+            {
+                let count = count.clone();
+                try_spawn_daemon(async move {
+                    for _ in 0..17 {
+                        count.fetch_add(1, Ordering::Relaxed);
+                        sleep(Duration::from_micros(200)).await;
+                    }
+                    Ok(())
+                })?;
+            }
+
+            // Start another
+            {
+                let count = count.clone();
+                try_spawn_daemon(async move {
+                    for _ in 0..6 {
+                        count.fetch_add(100, Ordering::Relaxed);
+                        sleep(Duration::from_micros(800)).await;
+                    }
+                    Ok(())
+                })?;
+            }
+
+            // Note how this function just starts a couple background jobs, but exits immediately
+
+            anyhow::Ok(())
+        })
+    }
+    .await;
+
+    // At this point, all background jobs have finished, now we can check the result of all our
+    // additions
+    assert_eq!(count.load(Ordering::Acquire), 41617);
+
+    Ok(())
+}
+
+/// Demo of how janitor can be used to implement async destructors
+struct AsyncDropDemo {
+    count: Arc<AtomicUsize>,
+}
+
+impl AsyncDropDemo {
+    async fn new(count: Arc<AtomicUsize>) -> Self {
+        count.fetch_add(1000, Ordering::Relaxed);
+        sleep(Duration::from_micros(50)).await;
+        AsyncDropDemo { count }
+    }
+}
+
+impl Drop for AsyncDropDemo {
+    fn drop(&mut self) {
+        let count = self.count.clone();
+        // This necessarily uses the panicking variant;
+        // we use spawn_cleanup_job because this makes more semantic sense in this context
+        spawn_cleanup_job(async move {
+            for _ in 0..4 {
+                count.fetch_add(10000, Ordering::Relaxed);
+                sleep(Duration::from_micros(800)).await;
+            }
+            Ok(())
+        })
+    }
+}