refactor main (#102)

* replace joinset with tasktracker

* restructure `main`

Author: vilgotf

Cargo.lock

@@ -13,7 +13,7 @@ hyper-rustls = { version = "0.27", default-features = false, features = ["webpki
 hyper-hickory = { version = "0.8", features = ["tokio"] }
 hyper-util = { version = "0.1", default-features = false, features = ["tokio", "server", "http1", "http2"] }
 tokio = { version = "1.43", features = ["rt-multi-thread", "macros", "signal"] }
-tokio-util = { version = "0.7.8", default-features = false, features = ["time"] }
+tokio-util = { version = "0.7.8", features = ["rt", "time"] }
 tracing = "0.1"
 tracing-subscriber = "0.3"
 twilight-http-ratelimiting = "0.17"

Cargo.toml

@@ -22,13 +22,14 @@ use std::{
     convert::Infallible,
     env,
     error::Error,
-    net::{Ipv4Addr, SocketAddr},
+    net::{Ipv4Addr, SocketAddrV4},
     pin::pin,
     str::FromStr,
     sync::Arc,
     time::{Duration, Instant},
 };
-use tokio::{net::TcpListener, task::JoinSet};
+use tokio::net::TcpListener;
+use tokio_util::task::TaskTracker;
 use tracing::{error, info, trace};
 use twilight_http_ratelimiting::{Endpoint, Method, RateLimitHeaders, RateLimiter};
 
@@ -55,10 +56,7 @@ static METRIC_KEY: LazyLock<Cow<str>> = LazyLock::new(|| {
 async fn main() -> Result<(), Box<dyn Error>> {
     tracing_subscriber::fmt::init();
 
-    let host = parse_env("HOST")?.unwrap_or(Ipv4Addr::UNSPECIFIED);
-    let port = parse_env("PORT")?.unwrap_or(80);
-
-    let https_connector = {
+    let client = {
         let mut http_connector = TokioHickoryResolver::default().into_http_connector();
         http_connector.enforce_http(false);
 
@@ -67,45 +65,42 @@ async fn main() -> Result<(), Box<dyn Error>> {
             .https_only()
             .enable_http1();
 
-        if env::var_os("DISABLE_HTTP2").is_some() {
+        let https_connector = if env::var_os("DISABLE_HTTP2").is_some() {
             builder.wrap_connector(http_connector)
         } else {
             builder.enable_http2().wrap_connector(http_connector)
-        }
+        };
+
+        Client::builder(TokioExecutor::new()).build(https_connector)
     };
 
-    let client: Client<_, Incoming> = Client::builder(TokioExecutor::new()).build(https_connector);
+    #[cfg(feature = "metrics")]
+    let handle = PrometheusBuilder::new()
+        .idle_timeout(
+            MetricKindMask::COUNTER | MetricKindMask::HISTOGRAM,
+            Some(Duration::from_secs(
+                parse_env("METRIC_TIMEOUT")?.unwrap_or(300),
+            )),
+        )
+        .install_recorder()
+        .expect("installed once");
+
     let ratelimiter_map = Arc::new(RatelimiterMap::new(
         env::var("DISCORD_TOKEN")?,
         Duration::from_secs(parse_env("CLIENT_DECAY_TIMEOUT")?.unwrap_or(3600)),
         parse_env("CLIENT_CACHE_MAX_SIZE")?,
     ));
 
-    let address = SocketAddr::from((host, port));
-
-    #[cfg(feature = "metrics")]
-    let handle: Arc<PrometheusHandle>;
-
-    #[cfg(feature = "metrics")]
-    {
-        let timeout = parse_env("METRIC_TIMEOUT")?.unwrap_or(300);
-        let recorder = PrometheusBuilder::new()
-            .idle_timeout(
-                MetricKindMask::COUNTER | MetricKindMask::HISTOGRAM,
-                Some(Duration::from_secs(timeout)),
-            )
-            .build_recorder();
-        handle = Arc::new(recorder.handle());
-        metrics::set_global_recorder(Box::new(recorder))
-            .expect("Failed to create metrics receiver!");
-    }
+    let host = parse_env("HOST")?.unwrap_or(Ipv4Addr::UNSPECIFIED);
+    let port = parse_env("PORT")?.unwrap_or(80);
+    let address = SocketAddrV4::new(host, port);
 
     let listener = TcpListener::bind(&address).await?;
     let mut shutdown_signal = pin!(shutdown_signal());
 
     info!("Listening on http://{}", address);
 
-    let mut tasks = JoinSet::new();
+    let tracker = TaskTracker::new();
 
     loop {
         tokio::select! {
@@ -115,53 +110,35 @@ async fn main() -> Result<(), Box<dyn Error>> {
                     continue;
                 };
 
-
-                let ratelimiter_map = ratelimiter_map.clone();
-                // Cloning a hyper client is fairly cheap by design
+                let ratelimiter_map = Arc::clone(&ratelimiter_map);
                 let client = client.clone();
-
                 #[cfg(feature = "metrics")]
                 let handle = handle.clone();
 
-                tasks.spawn(async move {
-                    trace!("Connection from: {:?}", addr);
-
-                    let service_fn = service::service_fn(move |incoming: Request<Incoming>| {
-                        let token = incoming
-                            .headers()
-                            .get("authorization")
-                            .and_then(|value| value.to_str().ok());
-                        let (ratelimiter, token) = ratelimiter_map.get_or_insert(token);
-                        let client = client.clone();
-
+                let service_fn = service::service_fn(move |request| {
+                    let token = request
+                        .headers()
+                        .get(header::AUTHORIZATION)
+                        .and_then(|value| value.to_str().ok());
+                    let (ratelimiter, token) = ratelimiter_map.get_or_insert(token);
+                    let client = client.clone();
+                    #[cfg(feature = "metrics")]
+                    let handle = handle.clone();
+
+                    async move {
                         #[cfg(feature = "metrics")]
-                        {
-                            let handle = handle.clone();
-
-                            async move {
-                                Ok::<_, Infallible>({
-                                    if incoming.uri().path() == "/metrics" {
-                                        handle_metrics(handle)
-                                    } else {
-                                        handle_request(client, ratelimiter, token, incoming)
-                                            .await
-                                            .unwrap_or_else(|err| err.as_response())
-                                    }
-                                })
-                            }
+                        if request.uri().path() == "/metrics" {
+                            return Ok::<_, Infallible>(handle_metrics(handle));
                         }
 
-                        #[cfg(not(feature = "metrics"))]
-                        {
-                            async move {
-                                Ok::<_, Infallible>(
-                                    handle_request(client, ratelimiter, token, incoming)
-                                        .await
-                                        .unwrap_or_else(|err| err.as_response()),
-                                )
-                            }
-                        }
-                    });
+                        Ok::<_, Infallible>(handle_request(client, ratelimiter, token, request)
+                            .await
+                            .unwrap_or_else(|err| err.as_response()))
+                    }
+                });
+
+                tracker.spawn(async move {
+                    trace!("Connection from: {:?}", addr);
 
                     let result = Builder::new(TokioExecutor::new())
                         .serve_connection(TokioIo::new(stream), service_fn)
@@ -180,7 +157,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
         }
     }
 
-    while tasks.join_next().await.is_some() {}
+    tracker.close();
+    info!("waiting for {} task(s) to finish", tracker.len());
+    tracker.wait().await;
 
     Ok(())
 }
@@ -308,7 +287,7 @@ async fn handle_request(
 }
 
 #[cfg(feature = "metrics")]
-fn handle_metrics(handle: Arc<PrometheusHandle>) -> Response<BoxBody<Bytes, hyper::Error>> {
+fn handle_metrics(handle: PrometheusHandle) -> Response<BoxBody<Bytes, hyper::Error>> {
     Response::builder()
         .header(
             header::CONTENT_TYPE,
@@ -379,10 +358,10 @@ fn parse_headers(
     }
 }
 
-fn parse_env<T>(key: &str) -> Result<Option<T>, Box<dyn Error>>
+fn parse_env<F>(key: &str) -> Result<Option<F>, Box<dyn Error>>
 where
-    T: FromStr,
-    <T as FromStr>::Err: Error + 'static,
+    F: FromStr,
+    <F as FromStr>::Err: Error + 'static,
 {
     match env::var(key) {
         Ok(s) => match s.parse() {