Browse Source

peer: manually configure tokio

Justin Tracey 2 months ago
parent
commit
0c064bd0fd
1 changed files with 20 additions and 8 deletions
  1. 20 8
      src/bin/mgen-peer.rs

+ 20 - 8
src/bin/mgen-peer.rs

@@ -1,5 +1,8 @@
 // Code specific to the peer in the p2p mode.
+
 use mgen::{log, updater::Updater, MessageHeader, SerializedMessage};
+
+use futures::future::try_join_all;
 use rand_xoshiro::{rand_core::SeedableRng, Xoshiro256PlusPlus};
 use serde::Deserialize;
 use std::collections::HashMap;
@@ -434,24 +437,33 @@ fn process_config(
     Ok(())
 }
 
-#[tokio::main]
-async fn main() -> Result<(), Box<dyn std::error::Error>> {
+async fn main_worker() -> Result<(), Box<dyn std::error::Error>> {
+    #[cfg(feature = "tracing")]
+    console_subscriber::init();
+
     let mut args = std::env::args();
     let _ = args.next();
-    let hosts_file = std::fs::read_to_string(args.next().unwrap())?;
+    let hosts_file = args.next().expect("missing hosts file arg");
+    let hosts_file = std::fs::read_to_string(hosts_file).expect("could not find hosts file");
     let hosts_map = parse_hosts_file(&hosts_file);
 
     let mut handles = vec![];
-
     for config_file in args.flat_map(|a| glob::glob(a.as_str()).unwrap()) {
         let yaml_s = std::fs::read_to_string(config_file?)?;
         let config: Config = serde_yaml::from_str(&yaml_s)?;
         process_config(config, &hosts_map, &mut handles)?;
     }
 
-    let handles: futures::stream::FuturesUnordered<_> = handles.into_iter().collect();
-    for handle in handles {
-        handle.await??;
-    }
+    try_join_all(handles).await?;
     Ok(())
 }
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(2)
+        .enable_all()
+        .disable_lifo_slot()
+        .build()
+        .unwrap()
+        .block_on(main_worker())
+}