Update TaskManager to support nesting (#794)

* Update TaskManager to support nesting * Refactor to support nested TaskManager better
helium · Apr 23, 2024 · 992d80c · 992d80c
1 parent fe059f9
commit 992d80c
Show file tree

Hide file tree

Showing 13 changed files with 194 additions and 64 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -44,6 +44,7 @@ tokio = { version = "1", default-features = false, features = [
   "fs",
   "macros",
   "signal",
+  "sync",
   "rt-multi-thread",
   "rt",
   "process",

diff --git a/boost_manager/src/main.rs b/boost_manager/src/main.rs
@@ -139,6 +139,7 @@ impl Server {
             .add_task(watcher)
             .add_task(updater)
             .add_task(purger)
+            .build()
             .start()
             .await
     }

diff --git a/ingest/src/server_iot.rs b/ingest/src/server_iot.rs
@@ -393,6 +393,7 @@ pub async fn grpc_server(settings: &Settings) -> Result<()> {
         .add_task(beacon_report_sink_server)
         .add_task(witness_report_sink_server)
         .add_task(grpc_server)
+        .build()
         .start()
         .await
 }
diff --git a/ingest/src/server_mobile.rs b/ingest/src/server_mobile.rs
@@ -454,6 +454,7 @@ pub async fn grpc_server(settings: &Settings) -> Result<()> {
         .add_task(invalidated_radio_threshold_report_sink_server)
         .add_task(coverage_object_report_sink_server)
         .add_task(grpc_server)
+        .build()
         .start()
         .await
 }
diff --git a/ingest/tests/iot_ingest.rs b/ingest/tests/iot_ingest.rs
@@ -28,7 +28,11 @@ async fn initialize_session_and_send_beacon_and_witness() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -77,7 +81,11 @@ async fn stream_stops_after_incorrectly_signed_init_request() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -111,7 +119,11 @@ async fn stream_stops_after_incorrectly_signed_beacon() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -148,7 +160,11 @@ async fn stream_stops_after_incorrect_beacon_pubkey() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -188,7 +204,11 @@ async fn stream_stops_after_incorrectly_signed_witness() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -225,7 +245,11 @@ async fn stream_stops_after_incorrect_witness_pubkey() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -265,7 +289,11 @@ async fn stream_stop_if_client_attempts_to_initiliaze_2nd_session() {
         .run_until(async move {
             tokio::task::spawn_local(async move {
                 let server = create_test_server(port, beacon_client, witness_client, None, None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let pub_key = generate_keypair();
@@ -316,7 +344,11 @@ async fn stream_stops_if_init_not_sent_within_timeout() {
             tokio::task::spawn_local(async move {
                 let server =
                     create_test_server(port, beacon_client, witness_client, Some(500), None);
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let mut client = connect_and_stream(port).await;
@@ -338,7 +370,11 @@ async fn stream_stops_on_session_timeout() {
             tokio::task::spawn_local(async move {
                 let server =
                     create_test_server(port, beacon_client, witness_client, Some(500), Some(900));
-                TaskManager::builder().add_task(server).start().await
+                TaskManager::builder()
+                    .add_task(server)
+                    .build()
+                    .start()
+                    .await
             });
 
             let mut client = connect_and_stream(port).await;

diff --git a/iot_config/src/main.rs b/iot_config/src/main.rs
@@ -123,6 +123,7 @@ impl Daemon {
         TaskManager::builder()
             .add_task(grpc_server)
             .add_task(db_cleaner)
+            .build()
             .start()
             .await
     }

diff --git a/iot_packet_verifier/src/daemon.rs b/iot_packet_verifier/src/daemon.rs
@@ -207,6 +207,7 @@ impl Cmd {
             .add_task(verifier_daemon)
             .add_task(burner)
             .add_task(report_files_server)
+            .build()
             .start()
             .await
     }

diff --git a/iot_verifier/src/main.rs b/iot_verifier/src/main.rs
@@ -318,6 +318,7 @@ impl Server {
             .add_task(pk_loader_server)
             .add_task(entropy_loader_server)
             .add_task(rewarder)
+            .build()
             .start()
             .await
     }

diff --git a/mobile_config/src/main.rs b/mobile_config/src/main.rs
@@ -107,7 +107,11 @@ impl Daemon {
             hex_boosting_svc,
         };
 
-        TaskManager::builder().add_task(grpc_server).start().await
+        TaskManager::builder()
+            .add_task(grpc_server)
+            .build()
+            .start()
+            .await
     }
 }
 

diff --git a/mobile_packet_verifier/src/daemon.rs b/mobile_packet_verifier/src/daemon.rs
@@ -186,6 +186,7 @@ impl Cmd {
             .add_task(reports_server)
             .add_task(event_id_purger)
             .add_task(daemon)
+            .build()
             .start()
             .await
     }

diff --git a/mobile_verifier/src/cli/server.rs b/mobile_verifier/src/cli/server.rs
@@ -389,6 +389,7 @@ impl Cmd {
             .add_task(radio_threshold_ingest_server)
             .add_task(invalidated_radio_threshold_ingest_server)
             .add_task(data_session_ingestor)
+            .build()
             .start()
             .await
     }

diff --git a/price/src/main.rs b/price/src/main.rs
@@ -110,6 +110,7 @@ impl Server {
             .add_task(hnt_price_generator)
             .add_task(mobile_price_generator)
             .add_task(iot_price_generator)
+            .build()
             .start()
             .await
     }