v0.3.11-pre.005

crates/ksp-worker-raw-transaction-ingest-lib/unit_tests/runtime.rs

@@ -1,5 +1,23 @@
 // file: crates/ksp-worker-raw-transaction-ingest-lib/unit_tests/runtime.rs
-// version: 1
+// version: 2
+
+struct ActiveTaskGuard {
+    active: std::sync::Arc<std::sync::atomic::AtomicUsize>,
+}
+
+impl ActiveTaskGuard {
+    fn new(active: std::sync::Arc<std::sync::atomic::AtomicUsize>) -> Self {
+        active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
+        return Self { active };
+    }
+}
+
+impl std::ops::Drop for ActiveTaskGuard {
+    fn drop(&mut self) {
+        self.active.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
+        return;
+    }
+}
 
 fn settings(network: &str) -> std::option::Option<crate::RawTransactionIngestSettings> {
     let network = match ksp_store_lib::RawNetworkId::new(network) {
@@ -13,6 +31,16 @@ fn settings(network: &str) -> std::option::Option<crate::RawTransactionIngestSet
     return std::option::Option::Some(crate::RawTransactionIngestSettings::with_defaults(network, worker_id));
 }
 
+async fn wait_for_active_count(active: &std::sync::Arc<std::sync::atomic::AtomicUsize>, expected: usize) -> bool {
+    for _ in 0..64 {
+        if active.load(std::sync::atomic::Ordering::Acquire) == expected {
+            return true;
+        }
+        tokio::task::yield_now().await;
+    }
+    return active.load(std::sync::atomic::Ordering::Acquire) == expected;
+}
+
 #[test]
 fn pre_004_current_runtime_is_required_before_spawn() {
     let result = super::current_runtime_handle();
@@ -125,3 +153,101 @@ async fn pre_004_dropping_last_control_handle_causes_private_runtime_exit() {
         }
     }
 }
+
+#[tokio::test(flavor = "current_thread")]
+async fn pre_005_supervisor_joins_all_cooperative_source_tasks_before_terminal() {
+    let settings = match settings("mainnet") {
+        std::option::Option::Some(value) => value,
+        std::option::Option::None => return,
+    };
+    let active = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let source_active = std::sync::Arc::clone(&active);
+    let handle = match super::start_foundation_with_source_spawner(
+        settings,
+        tokio::runtime::Handle::current(),
+        std::option::Option::None,
+        move |children, stop_receiver| {
+            for _ in 0..3 {
+                let active = std::sync::Arc::clone(&source_active);
+                let mut child_stop = stop_receiver.clone();
+                let _abort_handle = children.spawn(async move {
+                    let _guard = ActiveTaskGuard::new(active);
+                    loop {
+                        let changed = child_stop.changed().await;
+                        if changed.is_err() || *child_stop.borrow() {
+                            return;
+                        }
+                    }
+                });
+            }
+        },
+    ) {
+        std::result::Result::Ok(value) => value,
+        std::result::Result::Err(_) => return,
+    };
+    assert!(wait_for_active_count(&active, 3).await);
+    assert!(handle.request_stop());
+    let terminal = handle.wait_terminal().await;
+    let terminal = match terminal {
+        std::result::Result::Ok(value) => value,
+        std::result::Result::Err(_) => return,
+    };
+    assert_eq!(terminal, ksp_worker_api::WorkerState::Stopped);
+    assert_eq!(active.load(std::sync::atomic::Ordering::Acquire), 0);
+    return;
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn pre_005_supervisor_reaps_completed_source_task_and_still_joins_live_child_on_stop() {
+    let settings = match settings("mainnet") {
+        std::option::Option::Some(value) => value,
+        std::option::Option::None => return,
+    };
+    let active = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let completed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let source_active = std::sync::Arc::clone(&active);
+    let source_completed = std::sync::Arc::clone(&completed);
+    let handle = match super::start_foundation_with_source_spawner(
+        settings,
+        tokio::runtime::Handle::current(),
+        std::option::Option::None,
+        move |children, stop_receiver| {
+            let completed = std::sync::Arc::clone(&source_completed);
+            let _completed_abort_handle = children.spawn(async move {
+                completed.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
+                return;
+            });
+            let active = std::sync::Arc::clone(&source_active);
+            let mut child_stop = stop_receiver.clone();
+            let _active_abort_handle = children.spawn(async move {
+                let _guard = ActiveTaskGuard::new(active);
+                loop {
+                    let changed = child_stop.changed().await;
+                    if changed.is_err() || *child_stop.borrow() {
+                        return;
+                    }
+                }
+            });
+        },
+    ) {
+        std::result::Result::Ok(value) => value,
+        std::result::Result::Err(_) => return,
+    };
+    assert!(wait_for_active_count(&active, 1).await);
+    for _ in 0..64 {
+        if completed.load(std::sync::atomic::Ordering::Acquire) == 1 {
+            break;
+        }
+        tokio::task::yield_now().await;
+    }
+    assert_eq!(completed.load(std::sync::atomic::Ordering::Acquire), 1);
+    assert!(handle.request_stop());
+    let terminal = handle.wait_terminal().await;
+    let terminal = match terminal {
+        std::result::Result::Ok(value) => value,
+        std::result::Result::Err(_) => return,
+    };
+    assert_eq!(terminal, ksp_worker_api::WorkerState::Stopped);
+    assert_eq!(active.load(std::sync::atomic::Ordering::Acquire), 0);
+    return;
+}