v0.3.11-pre.006

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

@@ -1,5 +1,5 @@
 // file: crates/ksp-worker-raw-transaction-ingest-lib/src/runtime.rs
-// version: 2
+// version: 3
 
 /// Runtime-neutral boxed future resolving after one RAW transaction ingest Worker has fully reached a terminal lifecycle state.
 pub type RawTransactionIngestTerminalFuture<'a> =
@@ -85,6 +85,19 @@ fn current_runtime_handle() -> ksp_core_lib::Result<tokio::runtime::Handle> {
     };
 }
 
+async fn drain_admission(settings: &crate::RawTransactionIngestSettings, admission: &mut crate::RawTransactionAdmission) -> bool {
+    let mut clean = true;
+    admission.close();
+    loop {
+        let received = admission.receive(settings.network()).await;
+        match received {
+            std::result::Result::Ok(std::option::Option::Some(_acquisition)) => {},
+            std::result::Result::Ok(std::option::Option::None) => return clean,
+            std::result::Result::Err(_) => clean = false,
+        }
+    }
+}
+
 async fn drain_children(children: &mut tokio::task::JoinSet<()>) -> bool {
     let mut clean = true;
     while let std::option::Option::Some(joined) = children.join_next().await {
@@ -119,13 +132,16 @@ fn finish_stopped(lifecycle: &mut ksp_worker_api::WorkerLifecycle, sender: &toki
 }
 
 async fn run_supervisor<Spawner>(
+    settings: crate::RawTransactionIngestSettings,
     mut lifecycle: ksp_worker_api::WorkerLifecycle,
     _store_guard: std::option::Option<std::sync::Arc<ksp_store_lib::Store>>,
     mut stop_receiver: tokio::sync::watch::Receiver<bool>,
     terminal_sender: tokio::sync::watch::Sender<ksp_worker_api::WorkerState>,
     source_spawner: Spawner,
 ) where
-    Spawner: FnOnce(&mut tokio::task::JoinSet<()>, tokio::sync::watch::Receiver<bool>) + std::marker::Send + 'static,
+    Spawner: FnOnce(&mut tokio::task::JoinSet<()>, tokio::sync::watch::Receiver<bool>, tokio::sync::mpsc::Sender<crate::RawTransactionIngress>)
+        + std::marker::Send
+        + 'static,
 {
     if *stop_receiver.borrow() {
         finish_stopped(&mut lifecycle, &terminal_sender);
@@ -137,9 +153,12 @@ async fn run_supervisor<Spawner>(
     }
     terminal_sender.send_replace(lifecycle.state());
     let mut children = tokio::task::JoinSet::new();
-    source_spawner(&mut children, stop_receiver.clone());
-    let supervised_clean = supervise_until_stop(&mut stop_receiver, &mut children).await;
-    if !supervised_clean {
+    let (mut admission, admission_sender) = crate::RawTransactionAdmission::new(settings.admission_queue_capacity());
+    source_spawner(&mut children, stop_receiver.clone(), admission_sender);
+    let supervised_clean = supervise_until_stop(&settings, &mut stop_receiver, &mut children, &mut admission).await;
+    let admission_clean = drain_admission(&settings, &mut admission).await;
+    let children_clean = drain_children(&mut children).await;
+    if !supervised_clean || !admission_clean || !children_clean {
         finish_faulted(&mut lifecycle, &terminal_sender);
         return;
     }
@@ -152,7 +171,7 @@ fn start_foundation(
     runtime: tokio::runtime::Handle,
     store_guard: std::option::Option<std::sync::Arc<ksp_store_lib::Store>>,
 ) -> ksp_core_lib::Result<crate::RawTransactionIngestHandle> {
-    return start_foundation_with_source_spawner(settings, runtime, store_guard, |_children, _stop_receiver| {});
+    return start_foundation_with_source_spawner(settings, runtime, store_guard, |_children, _stop_receiver, _admission_sender| {});
 }
 
 fn start_foundation_with_source_spawner<Spawner>(
@@ -162,7 +181,9 @@ fn start_foundation_with_source_spawner<Spawner>(
     source_spawner: Spawner,
 ) -> ksp_core_lib::Result<crate::RawTransactionIngestHandle>
 where
-    Spawner: FnOnce(&mut tokio::task::JoinSet<()>, tokio::sync::watch::Receiver<bool>) + std::marker::Send + 'static,
+    Spawner: FnOnce(&mut tokio::task::JoinSet<()>, tokio::sync::watch::Receiver<bool>, tokio::sync::mpsc::Sender<crate::RawTransactionIngress>)
+        + std::marker::Send
+        + 'static,
 {
     let kind = match ksp_worker_api::WorkerKindCode::new(crate::RAW_TRANSACTION_INGEST_WORKER_KIND_CODE) {
         std::result::Result::Ok(value) => value,
@@ -176,17 +197,23 @@ where
     let (stop_sender, stop_receiver) = tokio::sync::watch::channel(false);
     let (terminal_sender, terminal_receiver) = tokio::sync::watch::channel(lifecycle.state());
     let handle = crate::RawTransactionIngestHandle { stop_sender, stop_token, terminal_receiver };
-    std::mem::drop(runtime.spawn(run_supervisor(lifecycle, store_guard, stop_receiver, terminal_sender, source_spawner)));
+    std::mem::drop(runtime.spawn(run_supervisor(settings, lifecycle, store_guard, stop_receiver, terminal_sender, source_spawner)));
     return std::result::Result::Ok(handle);
 }
 
-async fn supervise_until_stop(stop_receiver: &mut tokio::sync::watch::Receiver<bool>, children: &mut tokio::task::JoinSet<()>) -> bool {
+async fn supervise_until_stop(
+    settings: &crate::RawTransactionIngestSettings,
+    stop_receiver: &mut tokio::sync::watch::Receiver<bool>,
+    children: &mut tokio::task::JoinSet<()>,
+    admission: &mut crate::RawTransactionAdmission,
+) -> bool {
+    let mut admission_open = true;
     let mut clean = true;
     loop {
         if *stop_receiver.borrow() {
             break;
         }
-        if children.is_empty() {
+        if children.is_empty() && !admission_open {
             let changed = stop_receiver.changed().await;
             if changed.is_err() || *stop_receiver.borrow() {
                 break;
@@ -200,15 +227,21 @@ async fn supervise_until_stop(stop_receiver: &mut tokio::sync::watch::Receiver<b
                     break;
                 }
             }
-            joined = children.join_next() => {
+            joined = children.join_next(), if !children.is_empty() => {
                 if let std::option::Option::Some(std::result::Result::Err(_)) = joined {
                     clean = false;
                 }
             }
+            received = admission.receive(settings.network()), if admission_open => {
+                match received {
+                    std::result::Result::Ok(std::option::Option::Some(_acquisition)) => {},
+                    std::result::Result::Ok(std::option::Option::None) => admission_open = false,
+                    std::result::Result::Err(_) => clean = false,
+                }
+            }
         }
     }
-    let drained_clean = drain_children(children).await;
-    return clean && drained_clean;
+    return clean;
 }
 
 fn validate_store_network(settings: &crate::RawTransactionIngestSettings, store_network: &ksp_store_lib::RawNetworkId) -> ksp_core_lib::Result<()> {