scaffold_service.rs

//! Generate a minimal production-shaped axum 0.8 service: Cargo.toml + src/main.rs with shared
//! state, an AppError that maps to HTTP, health/readiness, tracing, timeout and body-limit
//! layers, a TaskTracker + CancellationToken background task, graceful shutdown and tests.
//! The output compiles, its tests pass, and it passes this skill's scanners.
//!
//! Build: rustc -O --edition 2021 scripts/scaffold_service.rs -o /tmp/scaffold_service
//! Usage: scaffold_service [--name NAME] [--file main|cargo] [--out DIR [--force]]

use std::path::Path;
use std::process::exit;

const USAGE: &str = "usage: scaffold_service [--name NAME] [--file main|cargo] [--out DIR [--force]]

Print (or write) a minimal axum 0.8 / tokio 1 service:
  Cargo.toml    axum 0.8, tokio 1, tower 0.5, tower-http 0.6, tokio-util 0.7, tracing, thiserror 2
  src/main.rs   AppState, AppError -> IntoResponse, /healthz, /readyz, /items/{id}, TraceLayer,
                TimeoutLayer, RequestBodyLimitLayer, a TaskTracker background task stopped by a
                CancellationToken, graceful shutdown on Ctrl+C/SIGTERM, oneshot router tests

options:
  --name NAME   crate name (lowercase, digits, '-' or '_'; default: my-service)
  --file F      print only `main` (src/main.rs) or `cargo` (Cargo.toml)
  --out DIR     write DIR/Cargo.toml and DIR/src/main.rs instead of printing
  --force       overwrite existing files with --out
Exit: 0 done, 2 usage error or files exist without --force.";

const CARGO: &str = r#"[package]
name = "{name}"
version = "0.1.0"
edition = "2021"
rust-version = "1.80"

[dependencies]
axum = "0.8"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
thiserror = "2"
tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "signal", "time", "sync"] }
tokio-util = { version = "0.7", features = ["rt"] }
tower = "0.5"
tower-http = { version = "0.6", features = ["trace", "timeout", "limit"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

[dev-dependencies]
tokio = { version = "1", features = ["test-util"] }
tower = { version = "0.5", features = ["util"] }
"#;

const MAIN: &str = r#"//! {name}: axum 0.8 service skeleton generated by scaffold_service.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

use axum::extract::{Path, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::{Json, Router};
use serde::Serialize;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use tower::ServiceBuilder;
use tower_http::limit::RequestBodyLimitLayer;
use tower_http::timeout::TimeoutLayer;
use tower_http::trace::TraceLayer;

/// Cheap to clone: everything shared lives behind one Arc.
#[derive(Clone)]
pub struct AppState {
    inner: Arc<Inner>,
}

struct Inner {
    ready: AtomicBool,
}

impl AppState {
    pub fn new() -> Self {
        Self { inner: Arc::new(Inner { ready: AtomicBool::new(false) }) }
    }

    pub fn set_ready(&self, ready: bool) {
        self.inner.ready.store(ready, Ordering::Release);
    }

    pub fn is_ready(&self) -> bool {
        self.inner.ready.load(Ordering::Acquire)
    }
}

impl Default for AppState {
    fn default() -> Self {
        Self::new()
    }
}

/// Every handler error. Client errors keep their message; server errors are logged, not leaked.
#[derive(Debug, thiserror::Error)]
pub enum AppError {
    #[error("{0} not found")]
    NotFound(String),
    #[error("service unavailable")]
    Unavailable,
    #[error("internal error")]
    Internal(#[source] Box<dyn std::error::Error + Send + Sync>),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let status = match &self {
            AppError::NotFound(_) => StatusCode::NOT_FOUND,
            AppError::Unavailable => StatusCode::SERVICE_UNAVAILABLE,
            AppError::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR,
        };
        if status.is_server_error() {
            tracing::error!(error = ?self, "request failed");
        }
        (status, Json(serde_json::json!({ "error": self.to_string() }))).into_response()
    }
}

#[derive(Debug, Serialize)]
pub struct Item {
    pub id: u64,
    pub name: String,
}

async fn healthz() -> &'static str {
    "ok"
}

async fn readyz(State(state): State<AppState>) -> Result<&'static str, AppError> {
    if state.is_ready() { Ok("ready") } else { Err(AppError::Unavailable) }
}

async fn get_item(Path(id): Path<u64>) -> Result<Json<Item>, AppError> {
    if id == 0 {
        return Err(AppError::NotFound(format!("item {id}")));
    }
    Ok(Json(Item { id, name: format!("item-{id}") }))
}

pub fn router(state: AppState) -> Router {
    Router::new()
        .route("/healthz", get(healthz))
        .route("/readyz", get(readyz))
        .route("/items/{id}", get(get_item))
        .layer(
            ServiceBuilder::new()
                // Outermost first. The body limit must wrap the timeout: TimeoutLayer needs a
                // `Default` response body, which RequestBodyLimitLayer's body type is not.
                .layer(TraceLayer::new_for_http())
                .layer(RequestBodyLimitLayer::new(1024 * 1024))
                .layer(TimeoutLayer::with_status_code(StatusCode::REQUEST_TIMEOUT, Duration::from_secs(10))),
        )
        .with_state(state)
}

/// Resolves on Ctrl+C or SIGTERM and cancels `token` so background tasks stop too.
async fn shutdown_signal(token: CancellationToken) {
    let ctrl_c = async {
        if let Err(err) = tokio::signal::ctrl_c().await {
            tracing::error!(error = %err, "cannot listen for Ctrl+C");
        }
    };
    #[cfg(unix)]
    let terminate = async {
        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
            Ok(mut sig) => {
                sig.recv().await;
            }
            Err(err) => tracing::error!(error = %err, "cannot listen for SIGTERM"),
        }
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();
    tokio::select! {
        _ = ctrl_c => {}
        _ = terminate => {}
    }
    tracing::info!("shutdown signal received");
    token.cancel();
}

fn spawn_background(tracker: &TaskTracker, token: CancellationToken) {
    tracker.spawn(async move {
        let mut tick = tokio::time::interval(Duration::from_secs(60));
        loop {
            tokio::select! {
                _ = token.cancelled() => break,
                _ = tick.tick() => tracing::debug!("background tick"),
            }
        }
        tracing::info!("background task stopped");
    });
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info,tower_http=info".into()),
        )
        .init();

    let token = CancellationToken::new();
    let tracker = TaskTracker::new();
    spawn_background(&tracker, token.child_token());

    let state = AppState::new();
    let addr = std::env::var("BIND_ADDR").unwrap_or_else(|_| "0.0.0.0:3000".to_string());
    let listener = TcpListener::bind(&addr).await?;
    tracing::info!(%addr, "listening");
    state.set_ready(true);

    axum::serve(listener, router(state.clone()))
        .with_graceful_shutdown(shutdown_signal(token.clone()))
        .await?;

    // Drain: stop accepting work, then give background tasks a bounded time to finish.
    state.set_ready(false);
    token.cancel();
    tracker.close();
    if tokio::time::timeout(Duration::from_secs(10), tracker.wait()).await.is_err() {
        tracing::warn!("background tasks did not stop within 10s");
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::{to_bytes, Body};
    use axum::http::Request;
    use tower::ServiceExt; // oneshot

    async fn call(app: Router, uri: &str) -> (StatusCode, String) {
        let res = app.oneshot(Request::get(uri).body(Body::empty()).expect("request")).await.expect("infallible");
        let status = res.status();
        let body = to_bytes(res.into_body(), 64 * 1024).await.expect("body");
        (status, String::from_utf8_lossy(&body).into_owned())
    }

    #[tokio::test]
    async fn healthz_is_ok() {
        assert_eq!(call(router(AppState::new()), "/healthz").await, (StatusCode::OK, "ok".into()));
    }

    #[tokio::test]
    async fn readyz_follows_state() {
        let state = AppState::new();
        assert_eq!(call(router(state.clone()), "/readyz").await.0, StatusCode::SERVICE_UNAVAILABLE);
        state.set_ready(true);
        assert_eq!(call(router(state), "/readyz").await.0, StatusCode::OK);
    }

    #[tokio::test]
    async fn missing_item_is_json_404() {
        let (status, body) = call(router(AppState::new()), "/items/0").await;
        assert_eq!(status, StatusCode::NOT_FOUND);
        assert!(body.contains("item 0 not found"), "{body}");
    }

    #[tokio::test(start_paused = true)]
    async fn background_task_stops_on_cancel() {
        let (tracker, token) = (TaskTracker::new(), CancellationToken::new());
        spawn_background(&tracker, token.clone());
        token.cancel();
        tracker.close();
        tokio::time::timeout(Duration::from_secs(1), tracker.wait()).await.expect("task stopped");
    }
}
"#;

fn die(msg: &str) -> ! {
    eprintln!("error: {msg}");
    exit(2);
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.iter().any(|a| a == "-h" || a == "--help") {
        println!("{USAGE}");
        return;
    }
    let (mut name, mut file, mut out, mut force) = ("my-service".to_string(), None, None, false);
    let mut it = args.iter();
    while let Some(a) = it.next() {
        let mut val = || it.next().cloned().unwrap_or_else(|| die(&format!("{a} needs a value")));
        match a.as_str() {
            "--name" => name = val(),
            "--file" => file = Some(val()),
            "--out" => out = Some(val()),
            "--force" => force = true,
            other => die(&format!("unknown argument {other}\n\n{USAGE}")),
        }
    }
    let valid = name.chars().next().is_some_and(|c| c.is_ascii_lowercase())
        && name.chars().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_');
    if !valid {
        die(&format!("--name must be lowercase letters, digits, '-' or '_', got '{name}'"));
    }
    let cargo = CARGO.replace("{name}", &name);
    let main = MAIN.replace("{name}", &name);
    match (out, file.as_deref()) {
        (Some(dir), _) => {
            let targets = [(Path::new(&dir).join("Cargo.toml"), cargo), (Path::new(&dir).join("src/main.rs"), main)];
            if !force {
                if let Some((p, _)) = targets.iter().find(|(p, _)| p.exists()) {
                    die(&format!("{} exists; pass --force to overwrite", p.display()));
                }
            }
            for (p, text) in &targets {
                if let Some(parent) = p.parent() {
                    std::fs::create_dir_all(parent).unwrap_or_else(|e| die(&format!("cannot create {}: {e}", parent.display())));
                }
                std::fs::write(p, text).unwrap_or_else(|e| die(&format!("cannot write {}: {e}", p.display())));
                println!("wrote {}", p.display());
            }
        }
        (None, Some("main")) => print!("{main}"),
        (None, Some("cargo")) => print!("{cargo}"),
        (None, Some(other)) => die(&format!("--file must be main or cargo, got '{other}'")),
        (None, None) => print!("# Cargo.toml\n{cargo}\n// src/main.rs\n{main}"),
    }
}