Background Jobs & Tasks

use std::sync::Arc;
use std::time::Duration;

use axum::{Json, Router, extract::State, http::StatusCode, routing::post};
use serde::Deserialize;
use tokio::sync::mpsc;
use tokio_util::{sync::CancellationToken, task::TaskTracker};

#[derive(Debug, Deserialize)]
struct ReviewPosted {
    member_email: String,
    title: String,
}

#[derive(Clone)]
struct AppState {
    jobs: mpsc::Sender<ReviewPosted>,
}

async fn post_review(State(state): State<AppState>, Json(job): Json<ReviewPosted>) -> StatusCode {
    match state.jobs.try_send(job) {
        Ok(()) => StatusCode::ACCEPTED,
        Err(mpsc::error::TrySendError::Full(_)) => StatusCode::SERVICE_UNAVAILABLE, // shed load, client retries
        Err(mpsc::error::TrySendError::Closed(_)) => StatusCode::SERVICE_UNAVAILABLE, // shutting down
    }
}

async fn send_thank_you(job: &ReviewPosted) -> Result<(), String> {
    tokio::time::sleep(Duration::from_millis(100)).await; // pretend to call the mail API
    println!("thanked {} for reviewing {}", job.member_email, job.title);
    Ok(())
}

/// N workers share one receiver; each takes the next job when it is free.
fn start_workers(rx: mpsc::Receiver<ReviewPosted>, workers: usize, tracker: &TaskTracker, cancel: CancellationToken) {
    let rx = Arc::new(tokio::sync::Mutex::new(rx)); // async mutex: held across recv().await on purpose
    for id in 0..workers {
        let (rx, cancel) = (Arc::clone(&rx), cancel.clone());
        tracker.spawn(async move {
            loop {
                let job = tokio::select! {
                    _ = cancel.cancelled() => break,
                    job = async { rx.lock().await.recv().await } => match job {
                        Some(job) => job,
                        None => break, // every sender dropped
                    },
                };
                if let Err(err) = send_thank_you(&job).await {
                    eprintln!("worker {id}: job {job:?} failed: {err}"); // never fail silently
                }
            }
        });
    }
}

#[tokio::main]
async fn main() {
    let (tx, rx) = mpsc::channel(100); // at most 100 queued jobs
    let tracker = TaskTracker::new();
    let cancel = CancellationToken::new();
    start_workers(rx, 4, &tracker, cancel.clone());

    let app = Router::new().route("/api/v1/reviews", post(post_review)).with_state(AppState { jobs: tx });
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app)
        .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.expect("signal handler") })
        .await
        .unwrap();

    // after the server stopped accepting requests: stop the workers, then wait for them
    cancel.cancel();
    tracker.close();
    tracker.wait().await;
}

$ curl -s -o /dev/null -w '%{http_code}\n' -H 'content-type: application/json' \
    -d '{"member_email":"ana@example.com","title":"Dune"}' localhost:3000/api/v1/reviews
202
thanked ana@example.com for reviewing Dune

async fn checksum_upload(bytes: Vec<u8>) -> u64 {
    tokio::task::spawn_blocking(move || {
        // stand-in for image decoding/resizing: pure CPU work
        bytes.iter().fold(0u64, |acc, b| acc.rotate_left(5) ^ u64::from(*b))
    })
    .await
    .expect("blocking task panicked")
}

#[tokio::main]
async fn main() {
    let sum = checksum_upload(vec![42; 5_000_000]).await;
    println!("checksum {sum:x}");
}

use tokio_cron_scheduler::{Job, JobScheduler};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut scheduler = JobScheduler::new().await?;

    // sec min hour day-of-month month day-of-week (UTC)
    scheduler
        .add(Job::new_async("0 0 7 * * *", |_id, _scheduler| {
            Box::pin(async {
                println!("sending loan-due reminders"); // query loans due in 2 days, enqueue e-mails
            })
        })?)
        .await?;

    // every 2 seconds, to see it work
    scheduler.add(Job::new_async("*/2 * * * * *", |_id, _scheduler| Box::pin(async { println!("tick") }))?).await?;

    scheduler.start().await?;
    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
    scheduler.shutdown().await?;
    Ok(())
}

use std::time::Duration;

use sqlx::PgPool;

/// Claims one due job, runs it, marks it done -- all in one transaction, so a crash releases the claim.
async fn work_one(pool: &PgPool) -> sqlx::Result<bool> {
    let mut tx = pool.begin().await?;
    let job = sqlx::query!(
        r#"SELECT id, kind, payload FROM jobs
           WHERE done_at IS NULL AND run_at <= now()
           ORDER BY run_at
           LIMIT 1
           FOR UPDATE SKIP LOCKED"#
    )
    .fetch_optional(&mut *tx)
    .await?;
    let Some(job) = job else { return Ok(false) };

    println!("running job {} ({}): {}", job.id, job.kind, job.payload);
    // on failure: UPDATE jobs SET attempts = attempts + 1, run_at = now() + backoff ... and commit

    sqlx::query!("UPDATE jobs SET done_at = now() WHERE id = $1", job.id).execute(&mut *tx).await?;
    tx.commit().await?;
    Ok(true)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?;
    sqlx::query!("INSERT INTO jobs (kind, payload) VALUES ('loan-reminder', '{\"loan_id\": 1}')").execute(&pool).await?;

    // a worker loop: drain due jobs, then sleep a little (or LISTEN/NOTIFY for instant wake-ups)
    for _ in 0..3 {
        while work_one(&pool).await? {}
        tokio::time::sleep(Duration::from_millis(200)).await;
    }
    Ok(())
}