Module axum 

Axum framework integration.

This module provides helpers for using BrowserPool with Axum.

Setup

Add to your Cargo.toml:

[dependencies]
html2pdf-api = { version = "0.1", features = ["axum-integration"] }
axum = "0.8"
tower = "0.5"

Basic Usage with State

use axum::{
    Router,
    routing::get,
    extract::State,
    response::IntoResponse,
    http::StatusCode,
};
use html2pdf_api::prelude::*;
use std::sync::Arc;

async fn generate_pdf(
    State(pool): State<SharedBrowserPool>,
) -> Result<impl IntoResponse, StatusCode> {
    let pool_guard = pool.lock().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let browser = pool_guard.get().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let tab = browser.new_tab().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    tab.navigate_to("https://example.com").map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    // Generate PDF...
    let pdf_data = tab.print_to_pdf(None).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    Ok((
        [(axum::http::header::CONTENT_TYPE, "application/pdf")],
        pdf_data,
    ))
}

#[tokio::main]
async fn main() {
    // Create and warmup pool
    let pool = BrowserPool::builder()
        .factory(Box::new(ChromeBrowserFactory::with_defaults()))
        .build()
        .expect("Failed to create pool");

    pool.warmup().await.expect("Failed to warmup");

    // Convert to shared state
    let shared_pool = pool.into_shared();

    let app = Router::new()
        .route("/pdf", get(generate_pdf))
        .with_state(shared_pool);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Using Extension Layer

Alternatively, use the Extension layer pattern:

use axum::{
    Router,
    routing::get,
    Extension,
    response::IntoResponse,
};
use html2pdf_api::prelude::*;
use std::sync::Arc;

async fn generate_pdf(
    Extension(pool): Extension<SharedBrowserPool>,
) -> impl IntoResponse {
    let pool_guard = pool.lock().unwrap();
    let browser = pool_guard.get().unwrap();
    // ...
}

#[tokio::main]
async fn main() {
    let pool = BrowserPool::builder()
        .factory(Box::new(ChromeBrowserFactory::with_defaults()))
        .build()
        .expect("Failed to create pool");

    pool.warmup().await.expect("Failed to warmup");

    let shared_pool = pool.into_shared();

    let app = Router::new()
        .route("/pdf", get(generate_pdf))
        .layer(Extension(shared_pool));

    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Using with init_browser_pool

If you have the env-config feature enabled:

use axum::{Router, routing::get};
use html2pdf_api::init_browser_pool;

#[tokio::main]
async fn main() {
    let pool = init_browser_pool().await
        .expect("Failed to initialize browser pool");

    let app = Router::new()
        .route("/pdf", get(generate_pdf))
        .with_state(pool);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Graceful Shutdown

For proper cleanup with graceful shutdown:

use axum::Router;
use html2pdf_api::prelude::*;
use std::sync::Arc;
use tokio::signal;

#[tokio::main]
async fn main() {
    let pool = BrowserPool::builder()
        .factory(Box::new(ChromeBrowserFactory::with_defaults()))
        .build()
        .expect("Failed to create pool");

    pool.warmup().await.expect("Failed to warmup");

    let shared_pool = Arc::new(std::sync::Mutex::new(pool));
    let shutdown_pool = Arc::clone(&shared_pool);

    let app = Router::new()
        .with_state(shared_pool);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
     
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal(shutdown_pool))
        .await
        .unwrap();
}

async fn shutdown_signal(pool: SharedBrowserPool) {
    let ctrl_c = async {
        signal::ctrl_c().await.expect("Failed to listen for ctrl+c");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("Failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    println!("Shutting down...");
    if let Ok(mut pool) = pool.lock() {
        pool.shutdown_async().await;
    }
}

Custom Extractor

For cleaner handler signatures, create a custom extractor:

use axum::{
    async_trait,
    extract::{FromRequestParts, State},
    http::{request::Parts, StatusCode},
};
use html2pdf_api::prelude::*;

pub struct Browser(pub BrowserHandle);

#[async_trait]
impl FromRequestParts<SharedBrowserPool> for Browser {
    type Rejection = StatusCode;

    async fn from_request_parts(
        _parts: &mut Parts,
        state: &SharedBrowserPool,
    ) -> Result<Self, Self::Rejection> {
        let pool = state.lock().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
        let browser = pool.get().map_err(|_| StatusCode::SERVICE_UNAVAILABLE)?;
        Ok(Browser(browser))
    }
}

// Then use in handlers:
async fn generate_pdf(Browser(browser): Browser) -> impl IntoResponse {
    let tab = browser.new_tab().unwrap();
    // ...
}

Traits

BrowserPoolAxumExt

Extension trait for BrowserPool with Axum helpers.

Functions

create_extension

Create an Axum Extension from an existing shared pool.

Type Aliases

BrowserPoolState

Type alias for Axum State extractor with the shared pool.