Tech Wiki

TOPICSSERIES

[Rust Zero to Production 30] Integrate a Tokio Background Worker with an Axum Application

An Axum handler and a Tokio worker are easy to build in isolation. The hard part is giving them one lifecycle. While request handlers read shared state, the worker runs HTTP checks and persists outcomes. Once shutdown begins, the application must reject new work without losing jobs it already accepted.

This fixture uses Rust 2024, rustc and Cargo 1.98.1, Axum 0.8.9, Tokio 1.53.1, tokio-util 0.7.16, and SQLx 0.8.6. It keeps the feature set small so the startup order and shutdown contract remain visible.

1. Define the boundaries first

Directly coupling HTTP and persistence implementations makes failure and shutdown tests needlessly difficult. HttpChecker receives the endpoint and a CancellationToken, then returns a domain outcome. Repository owns endpoint loading and result persistence.

#[async_trait]
pub trait HttpChecker: Send + Sync {
    async fn check(&self, endpoint: &Endpoint, cancel: CancellationToken) -> CheckOutcome;
}

#[async_trait]
pub trait Repository: Send + Sync {
    async fn endpoints(&self) -> Result<Vec<Endpoint>, AppError>;
    async fn save_result(&self, endpoint_id: &str, outcome: &CheckOutcome) -> Result<(), AppError>;
    async fn results(&self) -> Result<Vec<(String, String, Option<i64>, Option<String>)>, AppError>;
}

A checker reports either a successful status code or a failure string. The worker persists both, so a DNS failure or cancellation does not disappear. Test checkers reproduce success, failure, blocking, and cancellation without opening a network connection.

2. Bounded channels define admission policy

The admission channel before the scheduler and the work channel before the worker use the same bounded capacity. try_enqueue never waits. When the buffer has no room, it returns QueueError::Full, leaving the API layer free to reject the request or ask the caller to retry.

pub fn try_enqueue(&self, endpoint_id: &str) -> Result<(), QueueError> {
    let endpoint = self
        .endpoints
        .get(endpoint_id)
        .cloned()
        .ok_or(QueueError::UnknownEndpoint)?;
    let guard = self.admission.lock().expect("admission lock poisoned");
    let sender = guard.as_ref().ok_or(QueueError::Closed)?;
    sender.try_send(endpoint).map_err(|error| match error {
        mpsc::error::TrySendError::Full(_) => QueueError::Full,
        mpsc::error::TrySendError::Closed(_) => QueueError::Closed,
    })
}

Capacity is more than a tuning knob. When producers outrun the consumer, the limit returns pressure to the caller instead of turning memory into a queue. The blocking-checker test holds the first job in flight, buffers a second, and verifies that a third submission returns Full.

3. Give startup a single owner

Application::start reads endpoints from a migrated repository, builds shared state, and then spawns the worker followed by the scheduler. Handler state is complete before either background task starts.

startup_steps: vec![
    StartupStep::MigrationsApplied,
    StartupStep::EndpointsLoaded,
    StartupStep::SharedStateBuilt,
    StartupStep::WorkerSpawned,
    StartupStep::SchedulerSpawned,
],

The sequence is executable behavior, not a comment. An integration test checks it directly. If migration or initial loading fails, startup does not return a half-built application with one background task still running.

4. Keep Axum state focused on the API

The /status handler returns the endpoint count and number of completed checks. Router::with_state receives an Arc<SharedState>, and the handler extracts it with State. Queue senders and join handles stay inside Application rather than leaking into handler state.

async fn status(State(state): State<Arc<SharedState>>) -> Json<StatusBody> {
    Json(StatusBody {
        endpoint_count: state.endpoint_count,
        completed_checks: state.completed.load(Ordering::Acquire),
    })
}

pub fn router(&self) -> Router {
    Router::new()
        .route("/status", get(status))
        .with_state(Arc::clone(&self.shared))
}

That split keeps request handling separate from process lifecycle control. The API observes state. Application coordinates startup, admission, and shutdown.

5. Pass cancellation through the HTTP seam

The worker gives each check a child token. Calling cancel_checks lets the current checker and later checkers observe cancellation. A token does not forcibly stop work; the checker must await cancelled() or include it in a select! branch.

The fixture's cancellable checker returns Failure { error: "cancelled" } after receiving the signal. Its test waits until the worker persists that outcome and increments the completion counter. Cancellation is therefore not treated as permission to lose a result.

6. Shut down by closing admission, draining, then joining

Graceful shutdown removes the admission sender first. It signals the scheduler, which still forwards every item already in its channel. Once the scheduler drops the work sender, the worker drains the work queue and eventually receives None. The application awaits both join handles in a fixed order.

pub async fn shutdown(self) -> Result<ShutdownReport, AppError> {
    self.admission
        .lock()
        .expect("admission lock poisoned")
        .take();
    self.scheduler_cancel.cancel();
    self.scheduler.await??;
    let mut join_order = vec!["scheduler"];
    self.worker.await??;
    join_order.push("worker");
    Ok(ShutdownReport {
        drained_jobs: self.shared.completed.load(Ordering::Acquire),
        scheduler_joined: true,
        worker_joined: true,
        join_order,
    })
}

This path deliberately does not cancel the HTTP-check token. Graceful shutdown promises to drain accepted work. A production service can add a second phase that calls cancel_checks only after its drain deadline expires.

7. Re-run the contract

Run these commands from the repository root.

cd examples/article-30-axum-tokio-background-worker
cargo fmt --all -- --check
cargo check --all-targets --all-features
cargo clippy --all-targets --all-features -- -D warnings
cargo test --all-features
cargo test --release --all-features

Six integration tests cover startup order and Axum state, successful persistence, failure persistence, check cancellation, queue backpressure, and deterministic drain and join behavior. They need neither an external HTTP server nor a shared database file.

This is not a universal shutdown policy. A service still has to choose a drain deadline for slow HTTP requests, assign retry ownership after persistence failures, and set concurrency across multiple workers. Decide when admission closes and which accepted jobs must finish first. The Axum and worker boundaries follow from that contract.

Full source code

The complete runnable source for this article is available in the Chapter 30 project on GitHub.

Sources


3 responses

  1. […] Next articleIntegrate a Tokio Background Worker with an Axum Application […]

  2. […] Previous articleIntegrate a Tokio Background Worker with an Axum Application […]

  3. […] Previous articleIntegrate a Tokio Background Worker with an Axum Application […]

Leave a Reply

Your email address will not be published. Required fields are marked *

Tech Wiki

Built with WordPress · Learn in public.