main spawns the REAL main

This commit is contained in:
Toshit Chawda 2024-11-24 16:47:30 -08:00
parent 5a48c10cd9
commit 1b7f5a10c0
No known key found for this signature in database
GPG key ID: 91480ED99E2B3D9D
8 changed files with 122 additions and 90 deletions

View file

@ -10,6 +10,7 @@ async-trait = { version = "0.1.81", optional = true }
bytes = "1.7.1"
cfg-if = "1.0.0"
clap = { version = "4.5.16", features = ["cargo", "derive"] }
console-subscriber = { version = "0.4.1", optional = true }
dashmap = "6.0.1"
ed25519-dalek = { version = "2.1.1", features = ["pem"] }
env_logger = "0.11.5"
@ -50,6 +51,7 @@ toml = ["dep:toml"]
twisp = ["dep:pty-process", "dep:libc", "dep:async-trait", "dep:shell-words"]
speed-limit = ["dep:async-speed-limit"]
tokio-console = ["dep:console-subscriber", "tokio/tracing"]
[build-dependencies]
vergen-git2 = { version = "1.0.0", features = ["rustc"] }

View file

@ -290,9 +290,7 @@ pub async fn handle_wisp(stream: WispResult, is_v2: bool, id: String) -> anyhow:
let event: Arc<Event> = Event::new().into();
let mux_id = id.clone();
set.spawn(tokio::task::unconstrained(fut.map(move |x| {
debug!("wisp client id {:?} multiplexor result {:?}", mux_id, x)
})));
set.spawn(fut.map(move |x| debug!("wisp client id {:?} multiplexor result {:?}", mux_id, x)));
let ping_mux = mux.clone();
let ping_event = event.clone();

View file

@ -4,7 +4,7 @@
use std::{fs::read_to_string, net::IpAddr};
use anyhow::Context;
use anyhow::{Context, Result};
use clap::Parser;
use config::{validate_config_cache, Cli, Config, RuntimeFlavor};
use dashmap::DashMap;
@ -114,7 +114,7 @@ lazy_static! {
static JEMALLOCATOR: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
#[doc(hidden)]
fn main() -> anyhow::Result<()> {
fn main() -> Result<()> {
if CLI.default_config {
println!("{}", Config::default().ser()?);
return Ok(());
@ -135,88 +135,100 @@ fn main() -> anyhow::Result<()> {
builder.enable_all();
let rt = builder.build()?;
rt.block_on(async {
validate_config_cache().await;
rt.block_on(async move {
tokio::spawn(async_main()).await??;
Ok(())
})
}
info!(
"listening on {:?} with runtime flavor {:?} and socket transport {:?}",
CONFIG.server.bind, CONFIG.server.runtime, CONFIG.server.transport
);
#[doc(hidden)]
async fn async_main() -> Result<()> {
#[cfg(feature = "tokio-console")]
console_subscriber::init();
trace!("CLI: {:#?}", &*CLI);
trace!("CONFIG: {:#?}", &*CONFIG);
trace!("RESOLVER: {:?}", &*RESOLVER);
validate_config_cache().await;
tokio::spawn(async {
let mut sig = signal(SignalKind::user_defined1()).unwrap();
while sig.recv().await.is_some() {
match generate_stats() {
Ok(stats) => info!("Stats:\n{}", stats),
Err(err) => error!("error while creating stats {:?}", err),
}
info!(
"listening on {:?} with runtime flavor {:?} and socket transport {:?}",
CONFIG.server.bind, CONFIG.server.runtime, CONFIG.server.transport
);
trace!("CLI: {:#?}", &*CLI);
trace!("CONFIG: {:#?}", &*CONFIG);
trace!("RESOLVER: {:?}", &*RESOLVER);
tokio::spawn(async {
let mut sig = signal(SignalKind::user_defined1()).unwrap();
while sig.recv().await.is_some() {
match generate_stats() {
Ok(stats) => info!("Stats:\n{}", stats),
Err(err) => error!("error while creating stats {:?}", err),
}
});
}
});
let mut listener = ServerListener::new(&CONFIG.server.bind)
.await
.with_context(|| format!("failed to bind to address {}", CONFIG.server.bind.1))?;
let mut listener = ServerListener::new(&CONFIG.server.bind)
.await
.with_context(|| format!("failed to bind to address {}", CONFIG.server.bind.1))?;
if let Some(bind_addr) = CONFIG
.server
.stats_endpoint
.as_ref()
.and_then(|x| x.get_bindaddr())
{
info!("stats server listening on {:?}", bind_addr);
let mut stats_listener = ServerListener::new(&bind_addr).await.with_context(|| {
format!("failed to bind to address {} for stats server", bind_addr.1)
})?;
if let Some(bind_addr) = CONFIG
.server
.stats_endpoint
.as_ref()
.and_then(|x| x.get_bindaddr())
{
info!("stats server listening on {:?}", bind_addr);
let mut stats_listener = ServerListener::new(&bind_addr).await.with_context(|| {
format!("failed to bind to address {} for stats server", bind_addr.1)
})?;
tokio::spawn(async move {
loop {
match stats_listener.accept().await {
Ok((stream, _)) => {
tokio::spawn(async move {
loop {
match stats_listener.accept().await {
Ok((stream, _)) => {
tokio::spawn(async move {
if let Err(e) = route_stats(stream).await {
error!("error while routing stats client: {:?}", e);
}
}
Err(e) => error!("error while accepting stats client: {:?}", e),
});
}
Err(e) => error!("error while accepting stats client: {:?}", e),
}
});
}
let stats_endpoint = CONFIG
.server
.stats_endpoint
.as_ref()
.and_then(|x| x.get_endpoint());
loop {
let stats_endpoint = stats_endpoint.clone();
match listener.accept().await {
Ok((stream, client_id)) => {
tokio::spawn(async move {
let res = route::route(stream, stats_endpoint, move |stream, maybe_ip| {
let client_id = if let Some(ip) = maybe_ip {
format!("{} ({})", client_id, ip)
} else {
client_id
};
trace!("routed {:?}: {}", client_id, stream);
handle_stream(stream, client_id)
})
.await;
if let Err(e) = res {
error!("error while routing client: {:?}", e);
}
});
}
Err(e) => error!("error while accepting client: {:?}", e),
}
});
}
let stats_endpoint = CONFIG
.server
.stats_endpoint
.as_ref()
.and_then(|x| x.get_endpoint());
loop {
let stats_endpoint = stats_endpoint.clone();
match listener.accept().await {
Ok((stream, client_id)) => {
tokio::spawn(async move {
let res = route::route(stream, stats_endpoint, move |stream, maybe_ip| {
let client_id = if let Some(ip) = maybe_ip {
format!("{} ({})", client_id, ip)
} else {
client_id
};
trace!("routed {:?}: {}", client_id, stream);
handle_stream(stream, client_id)
})
.await;
if let Err(e) = res {
error!("error while routing client: {:?}", e);
}
});
}
Err(e) => error!("error while accepting client: {:?}", e),
}
})
}
}
#[doc(hidden)]