rust-tokio-async-service-builder

0. BUILD    -> once: for t in blocking_in_async await_holding task_hygiene service_hygiene scaffold_service; do
                 rustc -O --edition 2021 "$SKILL/scripts/$t.rs" -o "/tmp/$t"; done
1. SCAFFOLD -> new service: /tmp/scaffold_service --name orders-api --out ./orders-api
2. LINT     -> for t in blocking_in_async await_holding task_hygiene service_hygiene; do /tmp/$t src/; done
3. SHAPE    -> runtime (table below), AppState (Arc, std Mutex, pools, clients built once), router + layers
4. OWN      -> every spawn goes into a TaskTracker/JoinSet; every loop selects on a CancellationToken
5. BOUND    -> bounded channels, Semaphore around upstreams, timeouts on every outbound call and request
6. FAIL     -> AppError: IntoResponse; no unwrap in handlers; readiness reflects dependencies
7. OBSERVE  -> tracing + EnvFilter (+ JSON in prod), TraceLayer, #[instrument], tokio-console when stuck
8. VERIFY   -> cargo test with tower::ServiceExt::oneshot and start_paused; cargo clippy -- -D warnings;
               every scanner with --strict in CI

fn main() -> anyhow::Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .max_blocking_threads(64) // spawn_blocking pool, default 512
        .thread_name("api-worker")
        .enable_all() // I/O + time drivers
        .build()?;
    rt.block_on(run())
}

async fn run() -> anyhow::Result<()> {
    Ok(())
}

use std::{sync::Arc, time::Duration};

use axum::{
    Json, Router,
    extract::{Path, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::get,
};
use tokio::net::TcpListener;
use tower::ServiceBuilder;
use tower_http::{timeout::TimeoutLayer, trace::TraceLayer};

struct AppState {
    db: sqlx::PgPool,
    http: reqwest::Client,
}

#[derive(Debug, thiserror::Error)]
enum AppError {
    #[error("not found")]
    NotFound,
    #[error(transparent)]
    Db(#[from] sqlx::Error),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let status = match &self {
            AppError::NotFound => StatusCode::NOT_FOUND,
            AppError::Db(err) => {
                tracing::error!(error = %err, "database error");
                StatusCode::INTERNAL_SERVER_ERROR
            }
        };
        (status, self.to_string()).into_response()
    }
}

#[derive(serde::Serialize, sqlx::FromRow)]
struct Order {
    id: i64,
    sku: String,
}

async fn get_order(
    State(state): State<Arc<AppState>>,
    Path(id): Path<i64>,
) -> Result<Json<Order>, AppError> {
    let order = sqlx::query_as::<_, Order>("SELECT id, sku FROM orders WHERE id = $1")
        .bind(id)
        .fetch_optional(&state.db)
        .await?
        .ok_or(AppError::NotFound)?;
    Ok(Json(order))
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt().with_env_filter(tracing_subscriber::EnvFilter::from_default_env()).init();
    let db = sqlx::postgres::PgPoolOptions::new()
        .max_connections(20)
        .acquire_timeout(Duration::from_secs(3))
        .connect(&std::env::var("DATABASE_URL")?)
        .await?;
    let http = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .connect_timeout(Duration::from_secs(3))
        .build()?;
    let state = Arc::new(AppState { db, http });

    let app = Router::new()
        .route("/orders/{id}", get(get_order)) // axum 0.8 syntax; "/:id" panics
        .layer(
            ServiceBuilder::new()
                .layer(TraceLayer::new_for_http())
                .layer(TimeoutLayer::with_status_code(StatusCode::REQUEST_TIMEOUT, Duration::from_secs(10))),
        )
        .with_state(state); // Router<Arc<AppState>> -> Router<()>, which axum::serve needs

    let listener = TcpListener::bind("0.0.0.0:3000").await?;
    axum::serve(listener, app).with_graceful_shutdown(shutdown_signal()).await?;
    Ok(())
}

async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c().await.expect("install Ctrl+C handler");
    };
    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("install SIGTERM handler")
            .recv()
            .await;
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();
    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}

// Wrong: guard alive across .await -> "future cannot be sent between threads safely"
// let mut cache = state.cache.lock().unwrap();
// let v = fetch(&state.http, &key).await?;
// cache.insert(key, v);

// Right: lock, copy out, unlock, await, lock again
let cached = state.cache.lock().unwrap().get(&key).cloned();
let value = match cached {
    Some(v) => v,
    None => {
        let v = fetch(&state.http, &key).await?;
        state.cache.lock().unwrap().insert(key.clone(), v.clone());
        v
    }
};

use std::time::Duration;
use tokio_util::{sync::CancellationToken, task::TaskTracker};

async fn run_workers(token: CancellationToken) {
    let tracker = TaskTracker::new();
    for id in 0..4 {
        let token = token.child_token();
        tracker.spawn(async move {
            let mut tick = tokio::time::interval(Duration::from_secs(5));
            loop {
                tokio::select! {
                    biased; // check cancellation first, deterministically
                    _ = token.cancelled() => break,
                    _ = tick.tick() => tracing::debug!(worker = id, "tick"),
                }
            }
        });
    }
    tracker.close(); // no more spawns; wait() can now complete
    token.cancelled().await;
    if tokio::time::timeout(Duration::from_secs(10), tracker.wait()).await.is_err() {
        tracing::warn!("workers did not stop in 10s");
    }
}

// Short blocking call (file, sync library, bcrypt): the blocking pool
let hash = tokio::task::spawn_blocking(move || bcrypt::hash(password, 12)).await??;

// Parallel CPU: rayon pool, result back through a oneshot
let (tx, rx) = tokio::sync::oneshot::channel();
rayon::spawn(move || {
    let sum: u64 = items.par_iter().map(|x| expensive(x)).sum(); // use rayon::prelude::*
    let _ = tx.send(sum);
});
let sum = rx.await?;

use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};

fn init_tracing(json: bool) {
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info,tower_http=info"));
    let registry = tracing_subscriber::registry().with(filter);
    if json {
        registry.with(tracing_subscriber::fmt::layer().json()).init(); // feature "json"
    } else {
        registry.with(tracing_subscriber::fmt::layer()).init();
    }
}

#[tracing::instrument(skip(state), err)] // records `id`; logs the error on Err
async fn load_order(state: &AppState, id: i64) -> Result<Order, AppError> { /* ... */ }

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

    #[tokio::test]
    async fn unknown_order_is_404() {
        let app = router(test_state());
        let res = app.oneshot(Request::get("/orders/999").body(Body::empty()).unwrap()).await.unwrap();
        assert_eq!(res.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test(start_paused = true)] // tokio feature "test-util"; time auto-advances when idle
    async fn retry_gives_up_after_deadline() {
        let started = tokio::time::Instant::now();
        let res = with_retry(|| async { Err::<(), _>("down") }).await;
        assert!(res.is_err());
        // 4 backoff sleeps: 100 + 200 + 400 + 800 ms of virtual time, the test runs instantly
        assert!(started.elapsed() >= Duration::from_millis(1500));
    }
}

/tmp/blocking_in_async src/                        # files and/or directories (*.rs, skips target/)
/tmp/await_holding src/main.rs --json              # machine-readable findings with fix
cat handler.rs | /tmp/service_hygiene -            # stdin
/tmp/task_hygiene src/ --strict                    # exit 1 on any finding (CI gate)
/tmp/blocking_in_async src/ --select TK001,TK002 --ignore TK018
/tmp/service_hygiene --list-rules                  # each scanner lists its rules with why and fix

/tmp/scaffold_service --name orders-api            # print both files
/tmp/scaffold_service --name orders-api --file main   # only src/main.rs (or --file cargo)
/tmp/scaffold_service --name orders-api --out ./orders-api [--force]