Improvements

Cargo.toml

@@ -9,7 +9,13 @@ edition = "2021"
 axum = "0"
 libc = "0"
 tokio = { version = "1", features = ["full"] }
+tokio-util = "0.7.12"
 tower = "0"
 tower-http = { version = "0", features = ["trace", "timeout"] }
 tracing = "0"
 tracing-subscriber = { version = "0", features = ["json", "env-filter"] }
+
+[profile.release]
+codegen-units = 1
+lto = "fat"
+panic = "abort"

Dockerfile

@@ -8,8 +8,17 @@ COPY src/ /usr/src/src/
 WORKDIR /usr/src
 
 RUN apk add --no-cache rustup build-base && \
-    rustup-init -qy --profile=minimal && \
+    rustup-init -qy --profile=minimal --default-toolchain=$RUST_VERSION && \
     source "$HOME/.cargo/env" && \
+    mkdir src && echo "fn main() {}" > src/main.rs && \
+    cargo fetch && \
+    cargo build --release --target=x86_64-unknown-linux-musl && \
+    rm src/main.rs
+
+COPY src/ /usr/src/src/
+
+RUN source "$HOME/.cargo/env" && \
+    touch src/main.rs && \
     cargo build --release --target=x86_64-unknown-linux-musl
 
 FROM alpine:$ALPINE_VERSION

src/lib.rs

@@ -0,0 +1,31 @@
+use std::{io, time::Duration};
+
+use axum::{routing, Router};
+use tokio::net::TcpListener;
+use tokio_util::sync::CancellationToken;
+use tower_http::{timeout::TimeoutLayer, trace::TraceLayer};
+
+pub struct Config {
+    pub listen_port: u16,
+}
+
+pub async fn run(config: Config, token: CancellationToken) -> Result<(), io::Error> {
+    let app = Router::new()
+        .route("/", routing::get(|| async { "Hello world!" }))
+        .layer((
+            TraceLayer::new_for_http(),
+            TimeoutLayer::new(Duration::from_secs(10)),
+        ));
+
+    let listener = TcpListener::bind(("0.0.0.0", config.listen_port))
+        .await
+        .unwrap();
+
+    let _ = axum::serve(listener, app)
+        .with_graceful_shutdown(async move {
+            token.cancelled().await;
+        })
+        .await;
+
+    Ok(())
+}

src/main.rs

@@ -1,64 +1,85 @@
-use std::time::Duration;
+use core::panic;
+use std::str::FromStr;
+use std::{env, fmt::Display};
+use tokio::signal::unix::{signal, SignalKind};
+use tokio_util::sync::CancellationToken;
+use tracing_subscriber::{fmt, prelude::*, EnvFilter};
+use web_template::Config;
 
-use axum::{routing, Router};
-use tokio::{
-    net::TcpListener,
-    signal::{self, unix::SignalKind},
-    task::JoinSet,
-};
-use tower_http::{timeout::TimeoutLayer, trace::TraceLayer};
-use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
-
-#[tokio::main(flavor = "current_thread")]
-async fn main() -> Result<(), std::io::Error> {
+#[tokio::main]
+async fn main() -> Result<(), String> {
+    // Enable logging in JSON format
     tracing_subscriber::registry()
-        .with(tracing_subscriber::EnvFilter::from_env("LOG_LEVEL"))
-        .with(tracing_subscriber::fmt::layer().json())
+        .with(fmt::layer().json())
+        .with(EnvFilter::from_env("LOG_LEVEL"))
         .init();
 
-    let app = Router::new()
-        .route("/", routing::get(|| async { "Hello world!" }))
-        .layer((
-            TraceLayer::new_for_http(),
-            TimeoutLayer::new(Duration::from_secs(10)),
-        ));
+    // On panic format the message as JSON before exit
+    std::panic::set_hook(Box::new(|panic_info| {
+        let location = panic_info
+            .location()
+            .map(|loc| format!("{}:{}:{}", loc.file(), loc.line(), loc.column()))
+            .unwrap_or_else(|| "unknown location".to_string());
 
-    let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap();
+        let message = match panic_info.payload().downcast_ref::<String>() {
+            Some(s) => s,
+            None => "unknown message",
+        };
 
-    axum::serve(listener, app)
-        .with_graceful_shutdown(shutdown_signal())
-        .await
+        tracing::error!("Panic occured at {}: {}", location, message)
+    }));
+
+    // Parse environment variables and create a Config struct
+    let config = Config {
+        listen_port: parse_env("LISTEN_PORT", 8080),
+    };
+
+    let (Ok(mut sigterm), Ok(mut sigint)) = (
+        signal(SignalKind::terminate()),
+        signal(SignalKind::interrupt()),
+    ) else {
+        panic!("Failed to install signal handlers");
+    };
+
+    // A token to be used to signal a shutdown request
+    let token = CancellationToken::new();
+
+    let mut webserver_handle = tokio::spawn(web_template::run(config, token.clone()));
+
+    tokio::select! {
+        _ = sigterm.recv() => {
+            tracing::info!("Received SIGTERM");
+            token.cancel();
+        },
+        _ = sigint.recv() => {
+            tracing::info!("Received SIGINT");
+            token.cancel();
+        },
+        _ = &mut webserver_handle => {},
+    };
+
+    let _ = webserver_handle.await;
+    Ok(())
 }
 
-const SIGNALS: &'static [(i32, &'static str)] = &[
-    (libc::SIGTERM, "SIGTERM"),
-    (libc::SIGQUIT, "SIGQUIT"),
-    (libc::SIGINT, "SIGINT"),
-];
+fn parse_env<T>(env_name: &str, default: T) -> T
+where
+    T: Display + FromStr,
+    <T as FromStr>::Err: std::fmt::Display,
+{
+    let Ok(env_value) = env::var(env_name) else {
+        tracing::info!("Environment variable '{env_name}' not set, using default {default}");
+        return default;
+    };
 
-async fn shutdown_signal() {
-    let term_signals = [
-        SignalKind::terminate(),
-        SignalKind::quit(),
-        SignalKind::interrupt(),
-    ];
+    str::parse(&env_value).unwrap_or_else(|e| {
+        let msg = format!(
+            "Environment variable '{}' could not be converted to type {}: {}",
+            env_name,
+            std::any::type_name::<T>(),
+            e
+        );
 
-    let mut futures = JoinSet::new();
-
-    for term_signal in term_signals {
-        futures.spawn(async move {
-            signal::unix::signal(term_signal)
-                .expect("failed to install signal handler")
-                .recv()
-                .await;
-            term_signal.as_raw_value()
-        });
-    }
-
-    let kind = futures.join_next().await;
-
-    if let Some(Ok(kind)) = kind {
-        let name = SIGNALS.iter().find(|e| e.0 == kind).unwrap().1;
-        tracing::info!("{} received, shutting down gracefully", name);
-    }
+        panic!("{}", msg);
+    })
 }