runtime flavors

server/src/config.rs

@@ -59,6 +59,18 @@ pub enum SocketTransport {
 	LengthDelimitedLe,
 }
 
+#[derive(Serialize, Deserialize, Default, Debug)]
+#[serde(rename_all = "lowercase")]
+pub enum RuntimeFlavor {
+	/// Single-threaded tokio runtime.
+	SingleThread,
+	/// Multi-threaded tokio runtime.
+	#[default]
+	MultiThread,
+	/// Multi-threaded tokio runtime with an alternate work in progress scheduler.
+	MultiThreadAlt,
+}
+
 pub type BindAddr = (SocketType, String);
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -106,6 +118,8 @@ pub struct ServerConfig {
 
 	/// Server log level.
 	pub log_level: LevelFilter,
+	/// Runtime type.
+	pub runtime: RuntimeFlavor,
 }
 
 #[derive(Serialize, Deserialize, PartialEq, Eq)]
@@ -324,6 +338,7 @@ impl Default for ServerConfig {
 			max_message_size: 64 * 1024,
 
 			log_level: LevelFilter::Info,
+			runtime: RuntimeFlavor::default(),
 		}
 	}
 }

server/src/main.rs

@@ -5,7 +5,7 @@ use std::{collections::HashMap, fs::read_to_string, net::IpAddr};
 
 use anyhow::Context;
 use clap::Parser;
-use config::{validate_config_cache, Cli, Config};
+use config::{validate_config_cache, Cli, Config, RuntimeFlavor};
 use dashmap::DashMap;
 use handle::{handle_wisp, handle_wsproxy};
 use hickory_resolver::{
@@ -17,7 +17,10 @@ use listener::ServerListener;
 use log::{error, info};
 use route::{route_stats, ServerRouteResult};
 use serde::Serialize;
-use tokio::signal::unix::{signal, SignalKind};
+use tokio::{
+	runtime,
+	signal::unix::{signal, SignalKind},
+};
 use uuid::Uuid;
 use wisp_mux::{ConnectPacket, StreamType};
 
@@ -199,8 +202,7 @@ fn handle_stream(stream: ServerRouteResult, id: String) {
 #[global_allocator]
 static JEMALLOCATOR: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-#[tokio::main(flavor = "multi_thread")]
-async fn main() -> anyhow::Result<()> {
+fn main() -> anyhow::Result<()> {
 	if CLI.default_config {
 		println!("{}", Config::default().ser()?);
 		return Ok(());
@@ -211,11 +213,21 @@ async fn main() -> anyhow::Result<()> {
 		.parse_default_env()
 		.init();
 
+	let mut builder: runtime::Builder = match CONFIG.server.runtime {
+		RuntimeFlavor::SingleThread => runtime::Builder::new_current_thread(),
+		RuntimeFlavor::MultiThread => runtime::Builder::new_multi_thread(),
+		RuntimeFlavor::MultiThreadAlt => runtime::Builder::new_multi_thread_alt(),
+	};
+
+	builder.enable_all();
+	let rt = builder.build()?;
+
+	rt.block_on(async {
 		validate_config_cache().await;
 
 		info!(
-		"listening on {:?} with socket transport {:?}",
-		CONFIG.server.bind, CONFIG.server.transport
+			"listening on {:?} with runtime flavor {:?} and socket transport {:?}",
+			CONFIG.server.bind, CONFIG.server.runtime, CONFIG.server.transport
 		);
 
 		tokio::spawn(async {
@@ -232,7 +244,8 @@ async fn main() -> anyhow::Result<()> {
 		if CONFIG.server.enable_stats_endpoint {
 			if let Some(bind_addr) = CONFIG.server.stats_endpoint.get_bindaddr() {
 				info!("stats server listening on {:?}", bind_addr);
-			let mut stats_listener = ServerListener::new(&bind_addr).await.with_context(|| {
+				let mut stats_listener =
+					ServerListener::new(&bind_addr).await.with_context(|| {
 						format!("failed to bind to address {} for stats server", bind_addr.1)
 					})?;
 
@@ -270,4 +283,5 @@ async fn main() -> anyhow::Result<()> {
 				Err(e) => error!("error while accepting client: {:?}", e),
 			}
 		}
+	})
 }