m

Author: yinqiwen

src/main.rs

@@ -60,6 +60,12 @@ struct Args {
     #[clap(default_value = "5", long)]
     concurrent: usize,
 
+    #[clap(default_value = "2", long)]
+    threads: usize,
+
+    #[clap(default_value = "1024*1024", long)]
+    thread_stack_size: usize,
+
     #[clap(default_value = "mydomain.io", long)]
     tls_host: String,
 
@@ -97,24 +103,7 @@ fn rcgen(tls_host: &String) {
     }
 }
 
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
-    let args: Args = Args::parse();
-    if args.log.is_empty() {
-        tracing_subscriber::fmt::init();
-    } else {
-        let file_appender = tracing_appender::rolling::daily("./", args.log.as_str());
-        //let (non_blocking_appender, _guard) = tracing_appender::non_blocking(file_appender);
-        tracing_subscriber::fmt().with_writer(file_appender).init();
-        tokio::spawn(utils::clean_rotate_logs(format!("./{}", args.log.as_str())));
-    }
-    tracing::info!("{args:?}");
-
-    if args.rcgen {
-        rcgen(&args.tls_host);
-        return Ok(());
-    }
-
+async fn service_main(args: &Args) -> anyhow::Result<()> {
     let recorder = utils::MetricsLogRecorder::new(Duration::from_secs(10));
     metrics::set_boxed_recorder(Box::new(recorder)).unwrap();
 
@@ -192,3 +181,30 @@ async fn main() -> anyhow::Result<()> {
     }
     Ok(())
 }
+
+fn main() {
+    let args: Args = Args::parse();
+    if args.log.is_empty() {
+        tracing_subscriber::fmt::init();
+    } else {
+        let file_appender = tracing_appender::rolling::daily("./", args.log.as_str());
+        //let (non_blocking_appender, _guard) = tracing_appender::non_blocking(file_appender);
+        tracing_subscriber::fmt().with_writer(file_appender).init();
+        tokio::spawn(utils::clean_rotate_logs(format!("./{}", args.log.as_str())));
+    }
+    tracing::info!("{args:?}");
+
+    if args.rcgen {
+        rcgen(&args.tls_host);
+        return;
+    }
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(args.threads)
+        .enable_all()
+        .thread_stack_size(args.thread_stack_size)
+        .build()
+        .unwrap()
+        .block_on(async {
+            let _ = service_main(&args).await;
+        });
+}