Instruction file imported from SaynaAI/sayna (.cursor/rules/axum.mdc). Copyright stays with the author.

description: Comprehensive best practices for Axum 0.8+ development with WebSocket support globs: *.rs alwaysApply: true

Axum Best Practices: A Comprehensive Guide (0.8+)

This guide provides best practices for developing applications using the Axum framework in Rust, enhanced with WebSocket support for building atomic web services. It covers code organization, common patterns, performance, security, WebSocket specifics, testing, pitfalls, and tooling.

Note: This project uses Axum 0.8.x (released January 2025). Minimum Rust version: 1.75.

1. Code Organization and Structure

A clear project layout improves maintainability and scalability. Consider the following structure:

project_root/
├── src/
│   ├── main.rs             # Application entry point (sets up Router + server)
│   ├── lib.rs              # Shared library for core logic
│   ├── routes/             # Route definitions
│   │   ├── mod.rs
│   │   ├── api.rs          # REST endpoints
│   │   ├── ws.rs           # WebSocket routes and upgrade handlers
│   ├── handlers/           # Request handlers and WebSocket message handlers
│   │   ├── mod.rs
│   │   ├── api_handlers.rs
│   │   ├── ws_handlers.rs
│   ├── services/           # Business logic services
│   │   ├── mod.rs
│   │   ├── user_service.rs
│   │   ├── chat_service.rs # WebSocket message routing
│   ├── middleware/         # Tower middleware layers
│   │   ├── mod.rs
│   │   ├── auth.rs
│   │   ├── logging.rs
│   ├── state/              # Shared application state
│   │   ├── mod.rs          # Defines AppState struct (use `State<AppState>` extractor)
│   ├── errors/             # Error types and conversion to responses
│   │   ├── mod.rs
│   │   ├── app_error.rs    # Implements `IntoResponse`
│   ├── utils/              # Utility functions (e.g., DB pool setup)
│   │   ├── mod.rs
│   │   ├── db.rs
├── tests/                  # Integration tests
│   ├── api_tests.rs        # Tests for HTTP endpoints
│   └── ws_tests.rs         # WebSocket integration tests
├── Dockerfile              # Containerization
├── Cargo.toml
└── .env

1.1 Naming Conventions

• Modules & Files: lowercase with underscores (e.g., ws_handlers.rs, user_service.rs).
• Structs & Enums: CamelCase (e.g., AppState, ChatMessage).
• Functions: snake_case reflecting action (e.g., create_user, handle_ws_message).

1.2 Modular Boundaries

• Router vs. Handler: Keep route declarations (Router::route) in routes/ and logic in handlers/ or services/.
• State Management: Use State<AppState> for compile-time checked shared state (DB pools, config, caches). Reserve Extension only for per-request data injected by middleware (e.g., authenticated user info).
• Separation of Concerns: Handlers focus on HTTP/WebSocket framing; services manage business logic.

1.3 State vs Extension (Important)

• Prefer State<T> over Extension<T> for application-wide state:
  • State is type-checked at compile time; Extension gives runtime errors if type mismatches
  • State is faster than Extension
  • Use .with_state(app_state) when building the router
• Use Extension<T> only for:
  • Per-request data set by middleware (e.g., Extension<AuthenticatedUser>)
  • Dynamic data that varies per request

2. Common Patterns and Anti-patterns

2.1 Axum-Specific Patterns

• Extractors: Use built‑in extractors (Path, Query, Json, State) for typed data parsing.
• Tower Middleware: Leverage tower-http layers for CORS, compression, and tracing.
• Error Handling: Define custom error types implementing IntoResponse to standardize responses.
• State Management: Use State<AppState> extractor. Wrap shared resources in Arc within AppState and use tokio::sync types (e.g., RwLock, broadcast) for interior mutability.

2.2 Anti-patterns to Avoid

• Blocking in Handlers: Never perform blocking I/O; use spawn_blocking or async drivers.
• Global Mutable State: Avoid static mutable variables; prefer injected state.
• Large Handlers: Break complex flows into smaller service functions.
• Ignoring WebSocket Backpressure: Always handle send/receive delays to prevent buffer exhaustion.

3. Performance Considerations

• Asynchronous Ecosystem: Use async database clients (sqlx, tokio-postgres).
• Connection Pooling: Configure DB pools for optimal concurrency.
• HTTP/2 & Keep-Alive: Enable keep-alive and HTTP/2 via Hyper configuration in Server::builder.
• Compression & Caching: Add CompressionLayer and Cache middlewares from tower-http.
• Efficient Serialization: Use serde with #[serde(crate = "serde")] and consider bincode or MessagePack for WS when binary.

4. Security Best Practices

• CORS: Permit only trusted origins with CorsLayer.
• TLS: Terminate TLS at reverse proxy or configure Hyper with rustls.
• Input Validation: Rely on extractor types and manual checks for complex data.
• Authentication/Authorization: Use JWT or session-based auth; protect WS upgrades by validating tokens in the handshake.
• Rate Limiting: Apply RateLimitLayer to APIs and WS endpoints.

5. WebSocket Best Practices

• Upgrade Endpoint: Define a route using get(ws_handler) and call WebSocketUpgrade::on_upgrade.
• Message Loop: Spawn a task handling ws.on_upgrade, splitting WebSocket into Sender/Receiver.
• Heartbeats: Implement periodic ping/pong to detect stale connections.
• Channel Broadcasting: Use tokio::sync::broadcast or mpsc for multi-client message distribution.
• Graceful Shutdown: Listen for shutdown signals and close WS streams properly.
• Binary vs. Text: Choose appropriate frame types and validate payload sizes.

6. Testing Approaches

• Unit Tests: Test services and error types in isolation.
• HTTP Integration: Import tower::ServiceExt and use .oneshot() method:

  use tower::ServiceExt; // Required for oneshot()
  
  let app = Router::new()
      .route("/", get(handler))
      .with_state(app_state); // MUST call with_state() before oneshot()
  
  let response = app
      .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
      .await
      .unwrap();

  Note: oneshot() only works on Router<()>, so .with_state() must be called first for stateful routers.
• WebSocket Tests: Employ tokio-tungstenite to open in-memory WS connections and assert message flows.
• Property-Based Testing: Integrate proptest for fuzzing JSON payloads.
• Test Organization: Mirror src/ structure under tests/ for clarity.

7. Common Pitfalls and Gotchas

• Unhandled Close Frames: Always handle Close messages to avoid panics.
• Leaking Tasks: Ensure spawned tasks are bounded or tied to connection lifetime.
• Shared State Locks: Avoid long-held locks inside async contexts.
• Version Mismatches: Keep Axum, Tower, and Hyper versions compatible.
• Middleware Order: Layer order matters—e.g., CORS before auth.

8. Tooling and Environment

• Cargo & Rustup: Keep toolchains updated.
• Clippy & Rustfmt: Enforce linting and formatting via pre-commit hooks.
• Tokio Console: Use tokio-console for runtime diagnostics.
• Tracing: Integrate tracing with tracing-subscriber for structured logs.
• IDE Support: Leverage Rust Analyzer in VS Code or IntelliJ Rust.

9. Deployment Best Practices

• Containerization: Build minimal Docker images with scratch or distroless.
• Reverse Proxy: Offload TLS and load balancing to Nginx or Traefik.
• Auto-Scaling: Configure pods/instances to handle WS connections gracefully.
• Observability: Export metrics via tower-http metrics layer and scrape with Prometheus.
• CI/CD: Automate builds, tests, and docker pushes in GitHub Actions or GitLab CI.

10. Axum 0.8.x Breaking Changes

When migrating from older versions or reading older tutorials, note these changes:

• Path Syntax: Use /{param} and /{*wildcard} instead of /:param and /*wildcard
• async_trait: Import from async-trait crate directly, not from axum (for custom extractors)
• Option Extractors: Now requires OptionalFromRequestParts or OptionalFromRequest trait
• WebSocket Messages: Now use Bytes and Utf8Bytes instead of Vec<u8> and String
• State over Extension: State<T> is strongly preferred over Extension<T> for app state (since 0.6)

By following these Axum-specific guidelines—including WebSocket nuances—you can create atomic, resilient, and maintainable Rust web services tailored to modern requirements.