salvo_cache/

lib.rs

#![cfg_attr(test, allow(clippy::unwrap_used))]
//! Response caching middleware for the Salvo web framework.
//!
//! This middleware intercepts HTTP responses and caches them for subsequent
//! requests, reducing server load and improving response times for cacheable
//! content.
//!
//! # What Gets Cached
//!
//! The cache stores the complete response including:
//! - HTTP status code
//! - Response headers
//! - Response body (except for streaming responses)
//!
//! # Key Components
//!
//! - [`CacheIssuer`]: Determines the cache key for each request
//! - [`CacheStore`]: Backend storage for cached responses
//! - [`Cache`]: The middleware handler
//!
//! # Default Implementations
//!
//! - [`RequestIssuer`]: Generates cache keys from the request URI and method
//! - [`MokaStore`]: High-performance concurrent cache backed by [`moka`]
//!
//! # Example
//!
//! ```ignore
//! use std::time::Duration;
//! use salvo_cache::{Cache, MokaStore, RequestIssuer};
//! use salvo_core::prelude::*;
//!
//! let cache = Cache::new(
//!     MokaStore::builder()
//!         .time_to_live(Duration::from_secs(300))  // Cache for 5 minutes
//!         .build(),
//!     RequestIssuer::default(),
//! );
//!
//! let router = Router::new()
//!     .hoop(cache)
//!     .get(my_expensive_handler);
//! ```
//!
//! # Custom Cache Keys
//!
//! Implement [`CacheIssuer`] to customize cache key generation:
//!
//! ```ignore
//! use salvo_cache::CacheIssuer;
//!
//! struct UserBasedIssuer;
//! impl CacheIssuer for UserBasedIssuer {
//!     type Key = String;
//!
//!     async fn issue(&self, req: &mut Request, depot: &Depot) -> Option<Self::Key> {
//!         // Cache per user + path
//!         let user_id = depot.get::<String>("user_id").ok()?;
//!         Some(format!("{}:{}", user_id, req.uri().path()))
//!     }
//! }
//! ```
//!
//! # Skipping Cache
//!
//! By default, only GET requests are cached. Use the `skipper` method to customize:
//!
//! ```ignore
//! let cache = Cache::new(store, issuer)
//!     .skipper(|req, _depot| req.uri().path().starts_with("/api/"));
//! ```
//!
//! # Concurrent Misses
//!
//! Concurrent misses for the same cache key are coalesced. One request populates
//! the cache, and other in-flight requests reuse the generated cache entry when
//! the response is cacheable. At most [`DEFAULT_MAX_IN_FLIGHT`] distinct cache
//! keys are coalesced at once by default; additional misses bypass coalescing and
//! execute normally until an in-flight slot is released.
//!
//! # Limitations
//!
//! - Streaming responses ([`ResBody::Stream`]) cannot be cached
//! - Error responses are not cached
//!
//! Read more: <https://salvo.rs>
#![doc(html_favicon_url = "https://salvo.rs/favicon-32x32.png")]
#![doc(html_logo_url = "https://salvo.rs/images/logo.svg")]
#![cfg_attr(docsrs, feature(doc_cfg))]

use std::borrow::Borrow;
use std::collections::{HashMap, VecDeque};
use std::error::Error as StdError;
use std::fmt::{self, Debug, Formatter};
use std::hash::Hash;
use std::sync::{Arc, Mutex, MutexGuard};

use bytes::Bytes;
use salvo_core::handler::Skipper;
use salvo_core::http::header::{AUTHORIZATION, CACHE_CONTROL, COOKIE, SET_COOKIE, VARY};
use salvo_core::http::{HeaderMap, ResBody, StatusCode};
use salvo_core::{Depot, Error, FlowCtrl, Handler, Request, Response, async_trait, cfg_feature};
use tokio::sync::Notify;

mod skipper;
pub use skipper::MethodSkipper;

cfg_feature! {
    #![feature = "moka-store"]

    pub mod moka_store;
    pub use moka_store::{MokaStore};
}

/// Issues a cache key for a request, deciding whether the request should be cached.
pub trait CacheIssuer: Send + Sync + 'static {
    /// The key type used to identify a cached entry.
    type Key: Hash + Eq + Send + Sync + 'static;
    /// Issue a key for the request. If it returns `None`, the request will not be cached.
    fn issue(
        &self,
        req: &mut Request,
        depot: &Depot,
    ) -> impl Future<Output = Option<Self::Key>> + Send;
}
impl<F, K> CacheIssuer for F
where
    F: Fn(&mut Request, &Depot) -> Option<K> + Send + Sync + 'static,
    K: Hash + Eq + Send + Sync + 'static,
{
    type Key = K;
    async fn issue(&self, req: &mut Request, depot: &Depot) -> Option<Self::Key> {
        self(req, depot)
    }
}

/// Identify cacheable requests by their URI.
///
/// # Caveats
///
/// The generated key is derived only from the request's scheme, authority,
/// path, query, and (optionally) method. It does **not** include content
/// negotiation headers such as `Accept-Encoding` or `Accept`. If the cached
/// responses vary by those headers — for example when a compression middleware
/// also runs — a client may receive a representation encoded for a different
/// request (e.g. a `gzip` body without `Accept-Encoding: gzip`).
///
/// The cache stores the response produced by the handlers *inside* it (`hoop`s
/// run outer-to-inner and the entry is captured on the way back out), so the
/// negotiating middleware must run **outside** the cache — added to the router
/// *before* the cache hoop — so it re-negotiates on every request, including
/// cache hits. Alternatively, use a custom [`CacheIssuer`] that folds the
/// relevant headers into the key.
///
/// Note that a `Vary` response header is **not** a fix here: the store never
/// evaluates `Vary` at lookup time, so a `Vary` response is never cached (in
/// either `cache_private` mode) — it would otherwise be replayed under the same
/// key regardless of the request's negotiation headers. Fold the relevant headers
/// into a custom [`CacheIssuer`] key instead.
#[derive(Clone, Debug)]
pub struct RequestIssuer {
    use_scheme: bool,
    use_authority: bool,
    use_path: bool,
    use_query: bool,
    use_method: bool,
}
impl Default for RequestIssuer {
    fn default() -> Self {
        Self::new()
    }
}
impl RequestIssuer {
    /// Create a new `RequestIssuer`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            use_scheme: true,
            use_authority: true,
            use_path: true,
            use_query: true,
            use_method: true,
        }
    }
    /// Whether to use the request's URI scheme when generating the key.
    #[must_use]
    pub fn use_scheme(mut self, value: bool) -> Self {
        self.use_scheme = value;
        self
    }
    /// Whether to use the request's URI authority when generating the key.
    #[must_use]
    pub fn use_authority(mut self, value: bool) -> Self {
        self.use_authority = value;
        self
    }
    /// Whether to use the request's URI path when generating the key.
    #[must_use]
    pub fn use_path(mut self, value: bool) -> Self {
        self.use_path = value;
        self
    }
    /// Whether to use the request's URI query when generating the key.
    #[must_use]
    pub fn use_query(mut self, value: bool) -> Self {
        self.use_query = value;
        self
    }
    /// Whether to use the request method when generating the key.
    #[must_use]
    pub fn use_method(mut self, value: bool) -> Self {
        self.use_method = value;
        self
    }
}

impl CacheIssuer for RequestIssuer {
    type Key = String;
    async fn issue(&self, req: &mut Request, _depot: &Depot) -> Option<Self::Key> {
        let mut key = String::with_capacity(req.uri().path().len() + 16);
        if self.use_scheme
            && let Some(scheme) = req.uri().scheme_str()
        {
            key.push_str(scheme);
            key.push_str("://");
        }
        if self.use_authority
            && let Some(authority) = req.uri().authority()
        {
            key.push_str(authority.as_str());
        }
        if self.use_path {
            key.push_str(req.uri().path());
        }
        if self.use_query
            && let Some(query) = req.uri().query()
        {
            key.push('?');
            key.push_str(query);
        }
        if self.use_method {
            key.push('|');
            key.push_str(req.method().as_str());
        }
        Some(key)
    }
}

/// Store cache.
pub trait CacheStore: Send + Sync + 'static {
    /// Error type for CacheStore.
    type Error: StdError + Sync + Send + 'static;
    /// Key
    type Key: Hash + Eq + Send + Clone + 'static;
    /// Get the cache item from the store.
    fn load_entry<Q>(&self, key: &Q) -> impl Future<Output = Option<CachedEntry>> + Send
    where
        Self::Key: Borrow<Q>,
        Q: Hash + Eq + Sync;
    /// Save the cache item to the store.
    fn save_entry(
        &self,
        key: Self::Key,
        data: CachedEntry,
    ) -> impl Future<Output = Result<(), Self::Error>> + Send;
}

fn mutex_lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    match mutex.lock() {
        Ok(guard) => guard,
        Err(poisoned) => poisoned.into_inner(),
    }
}

/// Default maximum number of distinct cache keys tracked for concurrent miss coalescing.
pub const DEFAULT_MAX_IN_FLIGHT: usize = 1024;

struct InFlight<K> {
    entries: Mutex<HashMap<K, Arc<Flight>>>,
    max_entries: usize,
}

impl<K> InFlight<K> {
    fn new(max_entries: usize) -> Self {
        Self {
            entries: Mutex::new(HashMap::new()),
            max_entries,
        }
    }
}

impl<K> Default for InFlight<K> {
    fn default() -> Self {
        Self::new(DEFAULT_MAX_IN_FLIGHT)
    }
}

impl<K> InFlight<K>
where
    K: Hash + Eq + Clone,
{
    fn enter(self: &Arc<Self>, key: K) -> FlightPermit<K> {
        let mut entries = mutex_lock(&self.entries);
        if let Some(flight) = entries.get(&key) {
            return FlightPermit::Follower(flight.clone());
        }
        if entries.len() >= self.max_entries {
            return FlightPermit::Bypass;
        }

        let flight = Arc::new(Flight::default());
        entries.insert(key.clone(), flight.clone());
        FlightPermit::Leader(FlightGuard {
            key: Some(key),
            flight,
            in_flight: self.clone(),
        })
    }
}

impl<K> InFlight<K>
where
    K: Hash + Eq,
{
    fn remove(&self, key: &K, flight: &Arc<Flight>) {
        let mut entries = mutex_lock(&self.entries);
        if entries
            .get(key)
            .is_some_and(|current| Arc::ptr_eq(current, flight))
        {
            entries.remove(key);
        }
    }
}

enum FlightPermit<K>
where
    K: Hash + Eq,
{
    Leader(FlightGuard<K>),
    Follower(Arc<Flight>),
    Bypass,
}

struct FlightGuard<K>
where
    K: Hash + Eq,
{
    key: Option<K>,
    flight: Arc<Flight>,
    in_flight: Arc<InFlight<K>>,
}

impl<K> FlightGuard<K>
where
    K: Hash + Eq,
{
    fn finish(mut self, entry: Option<CachedEntry>) {
        self.complete(entry);
    }

    fn complete(&mut self, entry: Option<CachedEntry>) {
        if let Some(key) = self.key.take() {
            self.flight.finish(entry);
            self.in_flight.remove(&key, &self.flight);
        }
    }
}

impl<K> Drop for FlightGuard<K>
where
    K: Hash + Eq,
{
    fn drop(&mut self) {
        self.complete(None);
    }
}

#[derive(Default)]
struct Flight {
    state: Mutex<FlightState>,
    notify: Notify,
}

#[derive(Default)]
struct FlightState {
    done: bool,
    entry: Option<CachedEntry>,
}

impl Flight {
    async fn wait(&self) {
        loop {
            let notified = self.notify.notified();
            if mutex_lock(&self.state).done {
                return;
            }
            notified.await;
        }
    }

    fn finish(&self, entry: Option<CachedEntry>) {
        *mutex_lock(&self.state) = FlightState { done: true, entry };
        self.notify.notify_waiters();
    }

    fn entry(&self) -> Option<CachedEntry> {
        mutex_lock(&self.state).entry.clone()
    }
}

/// `CachedBody` is used to save the response body to `CacheStore`.
///
/// [`ResBody`] has a Stream type, which is not `Send + Sync`, so we need to convert it to
/// `CachedBody`. If the response's body is [`ResBody::Stream`], it will not be cached.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub enum CachedBody {
    /// No body.
    None,
    /// Single bytes body.
    Once(Bytes),
    /// Chunks body.
    Chunks(VecDeque<Bytes>),
}
impl TryFrom<&ResBody> for CachedBody {
    type Error = Error;
    fn try_from(body: &ResBody) -> Result<Self, Self::Error> {
        match body {
            ResBody::None => Ok(Self::None),
            ResBody::Once(bytes) => Ok(Self::Once(bytes.to_owned())),
            ResBody::Chunks(chunks) => Ok(Self::Chunks(chunks.to_owned())),
            _ => Err(Error::other("unsupported body type")),
        }
    }
}
impl From<CachedBody> for ResBody {
    fn from(body: CachedBody) -> Self {
        match body {
            CachedBody::None => Self::None,
            CachedBody::Once(bytes) => Self::Once(bytes),
            CachedBody::Chunks(chunks) => Self::Chunks(chunks),
        }
    }
}

/// Cached entry which will be stored in the cache store.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct CachedEntry {
    /// Response status.
    pub status: Option<StatusCode>,
    /// Response headers.
    pub headers: HeaderMap,
    /// Response body.
    ///
    /// *Notice: If the response's body is streaming, it will be ignored and not cached.
    pub body: CachedBody,
}
impl CachedEntry {
    /// Create a new `CachedEntry`.
    pub fn new(status: Option<StatusCode>, headers: HeaderMap, body: CachedBody) -> Self {
        Self {
            status,
            headers,
            body,
        }
    }

    /// Get the response status.
    pub fn status(&self) -> Option<StatusCode> {
        self.status
    }

    /// Get the response headers.
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    /// Get the response body.
    ///
    /// *Notice: If the response's body is streaming, it will be ignored and not cached.
    pub fn body(&self) -> &CachedBody {
        &self.body
    }
}

/// Cache middleware.
///
/// # Example
///
/// ```
/// use std::time::Duration;
///
/// use salvo_cache::{Cache, MokaStore, RequestIssuer};
/// use salvo_core::Router;
///
/// let cache = Cache::new(
///     MokaStore::builder()
///         .time_to_live(Duration::from_secs(60))
///         .build(),
///     RequestIssuer::default(),
/// );
/// let router = Router::new().hoop(cache);
/// ```
#[non_exhaustive]
pub struct Cache<S, I>
where
    S: CacheStore,
{
    /// Cache store.
    pub store: S,
    /// Cache issuer.
    pub issuer: I,
    /// Skipper.
    pub skipper: Box<dyn Skipper>,
    cache_private: bool,
    in_flight: Arc<InFlight<S::Key>>,
}
impl<S, I> Debug for Cache<S, I>
where
    S: CacheStore + Debug,
    I: Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("Cache")
            .field("store", &self.store)
            .field("issuer", &self.issuer)
            .finish()
    }
}

impl<S, I> Cache<S, I>
where
    S: CacheStore,
{
    /// Create a new `Cache`.
    #[inline]
    #[must_use]
    pub fn new(store: S, issuer: I) -> Self {
        let skipper = MethodSkipper::new().skip_all().skip_get(false);
        Self {
            store,
            issuer,
            skipper: Box::new(skipper),
            cache_private: false,
            in_flight: Arc::new(InFlight::default()),
        }
    }
    /// Sets skipper and returns a new `Cache`.
    #[inline]
    #[must_use]
    pub fn skipper(mut self, skipper: impl Skipper) -> Self {
        self.skipper = Box::new(skipper);
        self
    }

    /// Allow caching requests or responses that contain private-user cache signals.
    ///
    /// By default, requests with `Authorization` or `Cookie`, and responses with `Set-Cookie`,
    /// `Vary`, or `Cache-Control: private`, are not cached. Enabling this lifts those
    /// privacy restrictions. `Cache-Control: no-store`/`no-cache` responses are never
    /// cached regardless of this setting.
    #[inline]
    #[must_use]
    pub fn cache_private(mut self, cache_private: bool) -> Self {
        self.cache_private = cache_private;
        self
    }

    /// Sets the maximum number of distinct cache keys tracked for concurrent miss coalescing.
    ///
    /// When this limit is reached, misses for new keys bypass coalescing and execute normally.
    /// Existing in-flight keys can still be followed. Set to `0` to disable miss coalescing.
    #[inline]
    #[must_use]
    pub fn max_in_flight(mut self, max_in_flight: usize) -> Self {
        self.in_flight = Arc::new(InFlight::new(max_in_flight));
        self
    }
}

async fn call_next_and_cache<S>(
    store: &S,
    key: S::Key,
    cache_private: bool,
    req: &mut Request,
    depot: &mut Depot,
    res: &mut Response,
    ctrl: &mut FlowCtrl,
) -> Option<CachedEntry>
where
    S: CacheStore,
{
    // Snapshot the headers set by outer middleware *before* the inner handler
    // runs, so only the headers the handler itself adds or changes are cached.
    // This keeps per-request outer headers (e.g. `x-request-id`) out of the
    // entry while still capturing headers the handler overwrites (e.g. CORS
    // headers under `CallNext::After`).
    let headers_before = res.headers().clone();
    ctrl.call_next(req, depot, res).await;
    let cached_data = cached_response(res, &headers_before, cache_private)?;
    if let Err(e) = store.save_entry(key, cached_data.clone()).await {
        tracing::error!(error = ?e, "cache failed");
    }
    Some(cached_data)
}

fn cached_response(
    res: &Response,
    headers_before: &HeaderMap,
    cache_private: bool,
) -> Option<CachedEntry> {
    if res.body.is_stream() || res.body.is_error() {
        return None;
    }
    // Only cache successful responses. Use the *effective* status: a missing
    // status code is sent as `200 OK` when a body is present but as `404` when
    // the body is `None` (mirrors `Response::into_hyper`). Caching error or
    // redirect statuses (e.g. a transient `500`, an auth-dependent `401`/`403`,
    // or an empty unmatched `404`) would replay them to every client for the
    // whole TTL.
    let effective_status = res.status_code.unwrap_or(if res.body.is_none() {
        StatusCode::NOT_FOUND
    } else {
        StatusCode::OK
    });
    if !effective_status.is_success() {
        return None;
    }
    // `no-store`/`no-cache` forbid storing or reusing without revalidation in
    // *any* cache, so they apply even in private-cache mode.
    if response_disallows_caching(res.headers()) {
        return None;
    }
    if !cache_private && response_has_private_cache_headers(res.headers()) {
        return None;
    }
    // The entry can only record headers still present on the response, so a
    // header the handler *removed* (one set by an outer middleware before the
    // cache, now gone) cannot be represented. Replaying such an entry on a hit
    // would leave the stale outer header in place, making hits differ from the
    // cached miss. Skip caching instead of serving an inconsistent response.
    if headers_before
        .keys()
        .any(|name| !res.headers().contains_key(name))
    {
        return None;
    }
    let headers = handler_response_headers(headers_before, res.headers());
    let body = match TryInto::<CachedBody>::try_into(&res.body) {
        Ok(body) => body,
        Err(e) => {
            tracing::error!(error = ?e, "cache failed");
            return None;
        }
    };
    Some(CachedEntry::new(res.status_code, headers, body))
}

/// Headers the inner handler contributed, i.e. names whose values changed
/// between `before` (set by outer middleware) and `after` (the final response).
/// Headers left untouched by the handler are excluded so the cache does not
/// replay a stale per-request value over the fresh one on later hits.
fn handler_response_headers(before: &HeaderMap, after: &HeaderMap) -> HeaderMap {
    let mut headers = HeaderMap::new();
    for name in after.keys() {
        let after_values = after.get_all(name).iter().collect::<Vec<_>>();
        let before_values = before.get_all(name).iter().collect::<Vec<_>>();
        if after_values != before_values {
            for value in after_values {
                headers.append(name.clone(), value.clone());
            }
        }
    }
    headers
}

fn request_has_private_cache_headers(req: &Request) -> bool {
    req.headers().contains_key(AUTHORIZATION) || req.headers().contains_key(COOKIE)
}

fn response_has_private_cache_headers(headers: &HeaderMap) -> bool {
    headers.contains_key(SET_COOKIE) || cache_control_contains(headers, "private")
}

/// Directives/headers that forbid caching regardless of whether this is a shared
/// or private cache. `no-store` bans storing the response anywhere; `no-cache`
/// requires revalidation before reuse, which this middleware does not perform, so
/// both are treated as non-cacheable rather than served blindly.
///
/// `Vary` is also handled here: this middleware never folds the varied request
/// headers into the cache key, so it cannot tell two representations apart. Caching
/// a `Vary` response (even in private mode) would replay one client's representation
/// (e.g. a `gzip` body, or a specific `Accept-Language`) to clients that negotiated
/// differently, so such responses are not cached at all.
fn response_disallows_caching(headers: &HeaderMap) -> bool {
    cache_control_contains(headers, "no-store")
        || cache_control_contains(headers, "no-cache")
        || headers.contains_key(VARY)
}

fn cache_control_contains(headers: &HeaderMap, directive: &str) -> bool {
    headers.get_all(CACHE_CONTROL).iter().any(|value| {
        value.to_str().ok().is_some_and(|value| {
            value.split(',').any(|part| {
                let part = part.trim();
                part.eq_ignore_ascii_case(directive)
                    || part
                        .split_once('=')
                        .is_some_and(|(name, _)| name.trim().eq_ignore_ascii_case(directive))
            })
        })
    })
}

#[async_trait]
impl<S, I> Handler for Cache<S, I>
where
    S: CacheStore<Key = I::Key>,
    I: CacheIssuer,
    I::Key: Clone,
{
    async fn handle(
        &self,
        req: &mut Request,
        depot: &mut Depot,
        res: &mut Response,
        ctrl: &mut FlowCtrl,
    ) {
        if self.skipper.skipped(req, depot)
            || (!self.cache_private && request_has_private_cache_headers(req))
        {
            ctrl.call_next(req, depot, res).await;
            return;
        }
        let Some(key) = self.issuer.issue(req, depot).await else {
            // No cache key means "do not cache this request"; still run the rest
            // of the chain so the handler executes instead of returning an empty
            // response.
            ctrl.call_next(req, depot, res).await;
            return;
        };
        let Some(cache) = self.store.load_entry(&key).await else {
            match self.in_flight.enter(key.clone()) {
                FlightPermit::Leader(guard) => {
                    let cached_data = call_next_and_cache(
                        &self.store,
                        key,
                        self.cache_private,
                        req,
                        depot,
                        res,
                        ctrl,
                    )
                    .await;
                    guard.finish(cached_data);
                }
                FlightPermit::Follower(flight) => {
                    flight.wait().await;
                    if let Some(cache) = flight.entry() {
                        respond_from_cache(res, cache);
                        ctrl.skip_rest();
                    } else {
                        call_next_and_cache(
                            &self.store,
                            key,
                            self.cache_private,
                            req,
                            depot,
                            res,
                            ctrl,
                        )
                        .await;
                    }
                }
                FlightPermit::Bypass => {
                    call_next_and_cache(
                        &self.store,
                        key,
                        self.cache_private,
                        req,
                        depot,
                        res,
                        ctrl,
                    )
                    .await;
                }
            }
            return;
        };
        respond_from_cache(res, cache);
        ctrl.skip_rest();
    }
}

fn respond_from_cache(res: &mut Response, cache: CachedEntry) {
    let CachedEntry {
        status,
        headers,
        body,
    } = cache;
    if let Some(status) = status {
        res.status_code(status);
    }
    // Merge the cached headers into the response rather than replacing the whole
    // map, so headers set by outer middleware (request id, CORS, security
    // headers) on *this* request are preserved. The entry only holds the headers
    // the cached handler itself produced (see `handler_response_headers`), so a
    // cached header fully replaces any same-named value — restoring handler
    // overrides — while untouched outer headers are left as their fresh value.
    let res_headers = res.headers_mut();
    for name in headers.keys() {
        res_headers.remove(name);
        // Re-insert every cached value, keeping multi-valued headers such as
        // `Set-Cookie` intact.
        for value in headers.get_all(name) {
            res_headers.append(name.clone(), value.clone());
        }
    }
    *res.body_mut() = body.into();
}

#[cfg(test)]
mod tests {
    use std::collections::VecDeque;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    use bytes::Bytes;
    use salvo_core::http::HeaderMap;
    use salvo_core::prelude::*;
    use salvo_core::test::{ResponseExt, TestClient};
    use time::OffsetDateTime;

    use super::*;

    #[handler]
    async fn cached() -> String {
        format!(
            "Hello World, my birth time is {}",
            OffsetDateTime::now_utc()
        )
    }

    #[derive(Debug)]
    struct SlowCached {
        calls: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl Handler for SlowCached {
        async fn handle(
            &self,
            _req: &mut Request,
            _depot: &mut Depot,
            res: &mut Response,
            _ctrl: &mut FlowCtrl,
        ) {
            let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
            res.render(format!("backend call {call}"));
        }
    }

    #[tokio::test]
    async fn test_cache() {
        let cache = Cache::new(
            MokaStore::builder()
                .time_to_live(std::time::Duration::from_secs(5))
                .build(),
            RequestIssuer::default(),
        );
        let router = Router::new().hoop(cache).goal(cached);
        let service = Service::new(router);

        let mut res = TestClient::get("http://127.0.0.1:5801")
            .send(&service)
            .await;
        assert_eq!(res.status_code.unwrap(), StatusCode::OK);

        let content0 = res.take_string().await.unwrap();

        let mut res = TestClient::get("http://127.0.0.1:5801")
            .send(&service)
            .await;
        assert_eq!(res.status_code.unwrap(), StatusCode::OK);

        let content1 = res.take_string().await.unwrap();
        assert_eq!(content0, content1);

        tokio::time::sleep(tokio::time::Duration::from_secs(6)).await;
        let mut res = TestClient::post("http://127.0.0.1:5801")
            .send(&service)
            .await;
        let content2 = res.take_string().await.unwrap();

        assert_ne!(content0, content2);
    }

    #[handler]
    async fn server_error(res: &mut Response) {
        res.status_code(StatusCode::INTERNAL_SERVER_ERROR);
        res.render(format!("error at {}", OffsetDateTime::now_utc()));
    }

    #[tokio::test]
    async fn test_cache_skips_non_success_status() {
        let cache = Cache::new(
            MokaStore::builder()
                .time_to_live(std::time::Duration::from_secs(5))
                .build(),
            RequestIssuer::default(),
        );
        let router = Router::new().hoop(cache).goal(server_error);
        let service = Service::new(router);

        let mut res = TestClient::get("http://127.0.0.1:5802")
            .send(&service)
            .await;
        assert_eq!(res.status_code.unwrap(), StatusCode::INTERNAL_SERVER_ERROR);
        let content0 = res.take_string().await.unwrap();

        // A non-success response must not be cached, so the backend runs again
        // and produces a fresh (different) body.
        let mut res = TestClient::get("http://127.0.0.1:5802")
            .send(&service)
            .await;
        let content1 = res.take_string().await.unwrap();
        assert_ne!(content0, content1);
    }

    #[tokio::test]
    async fn test_cache_issuer_none_runs_handler() {
        let cache = Cache::new(
            MokaStore::builder()
                .time_to_live(std::time::Duration::from_secs(5))
                .build(),
            // Issuer that never produces a key: requests must still be handled.
            |_req: &mut Request, _depot: &Depot| Option::<String>::None,
        );
        let router = Router::new().hoop(cache).goal(cached);
        let service = Service::new(router);

        let mut res = TestClient::get("http://127.0.0.1:5803")
            .send(&service)
            .await;
        assert_eq!(res.status_code.unwrap(), StatusCode::OK);
        let body = res.take_string().await.unwrap();
        assert!(body.contains("Hello World"));
    }

    #[test]
    fn test_cached_response_skips_empty_unmatched_404() {
        // No status code + empty body is sent as `404 NOT_FOUND`
        // (see `Response::into_hyper`), so its effective status is non-success
        // and it must not be cached.
        let res = Response::new();
        assert!(res.status_code.is_none() && res.body.is_none());
        assert!(cached_response(&res, &HeaderMap::new(), false).is_none());
    }

    #[test]
    fn test_cached_response_caches_default_success() {
        // No status code but a body present is sent as `200 OK`, so it is cached.
        let mut res = Response::new();
        res.render("ok");
        assert!(res.status_code.is_none());
        assert!(cached_response(&res, &HeaderMap::new(), false).is_some());
    }

    // Tests for RequestIssuer
    #[test]
    fn test_request_issuer_new() {
        let issuer = RequestIssuer::new();
        assert!(issuer.use_scheme);
        assert!(issuer.use_authority);
        assert!(issuer.use_path);
        assert!(issuer.use_query);
        assert!(issuer.use_method);
    }

    #[test]
    fn test_request_issuer_default() {
        let issuer = RequestIssuer::default();
        assert!(issuer.use_scheme);
        assert!(issuer.use_authority);
        assert!(issuer.use_path);
        assert!(issuer.use_query);
        assert!(issuer.use_method);
    }

    #[test]
    fn test_request_issuer_use_scheme() {
        let issuer = RequestIssuer::new().use_scheme(false);
        assert!(!issuer.use_scheme);
        assert!(issuer.use_authority);
    }

    #[test]
    fn test_request_issuer_use_authority() {
        let issuer = RequestIssuer::new().use_authority(false);
        assert!(issuer.use_scheme);
        assert!(!issuer.use_authority);
    }

    #[test]
    fn test_request_issuer_use_path() {
        let issuer = RequestIssuer::new().use_path(false);
        assert!(!issuer.use_path);
    }

    #[test]
    fn test_request_issuer_use_query() {
        let issuer = RequestIssuer::new().use_query(false);
        assert!(!issuer.use_query);
    }

    #[test]
    fn test_request_issuer_use_method() {
        let issuer = RequestIssuer::new().use_method(false);
        assert!(!issuer.use_method);
    }

    #[test]
    fn test_request_issuer_chain() {
        let issuer = RequestIssuer::new()
            .use_scheme(false)
            .use_authority(false)
            .use_path(true)
            .use_query(false)
            .use_method(true);
        assert!(!issuer.use_scheme);
        assert!(!issuer.use_authority);
        assert!(issuer.use_path);
        assert!(!issuer.use_query);
        assert!(issuer.use_method);
    }

    #[test]
    fn test_request_issuer_debug() {
        let issuer = RequestIssuer::new();
        let debug_str = format!("{issuer:?}");
        assert!(debug_str.contains("RequestIssuer"));
        assert!(debug_str.contains("use_scheme"));
    }

    #[test]
    fn test_request_issuer_clone() {
        let issuer = RequestIssuer::new().use_scheme(false);
        let cloned = issuer.clone();
        assert_eq!(issuer.use_scheme, cloned.use_scheme);
        assert_eq!(issuer.use_authority, cloned.use_authority);
    }

    // Tests for CachedBody
    #[test]
    fn test_cached_body_none() {
        let body = CachedBody::None;
        assert_eq!(body, CachedBody::None);
    }

    #[test]
    fn test_cached_body_once() {
        let bytes = Bytes::from("test data");
        let body = CachedBody::Once(bytes.clone());
        assert_eq!(body, CachedBody::Once(bytes));
    }

    #[test]
    fn test_cached_body_chunks() {
        let mut chunks = VecDeque::new();
        chunks.push_back(Bytes::from("chunk1"));
        chunks.push_back(Bytes::from("chunk2"));
        let body = CachedBody::Chunks(chunks.clone());
        assert_eq!(body, CachedBody::Chunks(chunks));
    }

    #[test]
    fn test_cached_body_try_from_res_body_none() {
        let res_body = ResBody::None;
        let result: Result<CachedBody, _> = (&res_body).try_into();
        assert_eq!(result.unwrap(), CachedBody::None);
    }

    #[test]
    fn test_cached_body_try_from_res_body_once() {
        let bytes = Bytes::from("test");
        let res_body = ResBody::Once(bytes.clone());
        let result: Result<CachedBody, _> = (&res_body).try_into();
        assert_eq!(result.unwrap(), CachedBody::Once(bytes));
    }

    #[test]
    fn test_cached_body_try_from_res_body_chunks() {
        let mut chunks = VecDeque::new();
        chunks.push_back(Bytes::from("chunk1"));
        chunks.push_back(Bytes::from("chunk2"));
        let res_body = ResBody::Chunks(chunks.clone());
        let result: Result<CachedBody, _> = (&res_body).try_into();
        assert_eq!(result.unwrap(), CachedBody::Chunks(chunks));
    }

    #[test]
    fn test_cached_body_into_res_body_none() {
        let cb = CachedBody::None;
        let res_body: ResBody = cb.into();
        assert!(matches!(res_body, ResBody::None));
    }

    #[test]
    fn test_cached_body_into_res_body_once() {
        let bytes = Bytes::from("test");
        let cb = CachedBody::Once(bytes.clone());
        let res_body: ResBody = cb.into();
        assert!(matches!(res_body, ResBody::Once(b) if b == bytes));
    }

    #[test]
    fn test_cached_body_into_res_body_chunks() {
        let mut chunks = VecDeque::new();
        chunks.push_back(Bytes::from("chunk1"));
        let cb = CachedBody::Chunks(chunks);
        let res_body: ResBody = cb.into();
        assert!(matches!(res_body, ResBody::Chunks(_)));
    }

    #[test]
    fn test_cached_body_debug() {
        let body = CachedBody::None;
        let debug_str = format!("{body:?}");
        assert!(debug_str.contains("None"));

        let body = CachedBody::Once(Bytes::from("test"));
        let debug_str = format!("{body:?}");
        assert!(debug_str.contains("Once"));
    }

    #[test]
    fn test_cached_body_clone() {
        let body = CachedBody::Once(Bytes::from("test"));
        let cloned = body.clone();
        assert_eq!(body, cloned);
    }

    // Tests for CachedEntry
    #[test]
    fn test_cached_entry_new() {
        let entry = CachedEntry::new(Some(StatusCode::OK), HeaderMap::new(), CachedBody::None);
        assert_eq!(entry.status, Some(StatusCode::OK));
        assert!(entry.headers.is_empty());
        assert_eq!(entry.body, CachedBody::None);
    }

    #[test]
    fn test_cached_entry_status() {
        let entry = CachedEntry::new(
            Some(StatusCode::NOT_FOUND),
            HeaderMap::new(),
            CachedBody::None,
        );
        assert_eq!(entry.status(), Some(StatusCode::NOT_FOUND));
    }

    #[test]
    fn test_cached_entry_status_none() {
        let entry = CachedEntry::new(None, HeaderMap::new(), CachedBody::None);
        assert_eq!(entry.status(), None);
    }

    #[test]
    fn test_cached_entry_headers() {
        let mut headers = HeaderMap::new();
        headers.insert("Content-Type", "application/json".parse().unwrap());
        let entry = CachedEntry::new(Some(StatusCode::OK), headers.clone(), CachedBody::None);
        assert_eq!(entry.headers().len(), 1);
        assert!(entry.headers().contains_key("Content-Type"));
    }

    #[test]
    fn test_cached_entry_body() {
        let body = CachedBody::Once(Bytes::from("test body"));
        let entry = CachedEntry::new(Some(StatusCode::OK), HeaderMap::new(), body.clone());
        assert_eq!(entry.body(), &body);
    }

    #[test]
    fn test_cached_entry_debug() {
        let entry = CachedEntry::new(Some(StatusCode::OK), HeaderMap::new(), CachedBody::None);
        let debug_str = format!("{entry:?}");
        assert!(debug_str.contains("CachedEntry"));
        assert!(debug_str.contains("status"));
    }

    #[test]
    fn handler_response_headers_excludes_untouched_outer_headers() {
        // An outer middleware set `x-request-id` before the handler ran; the
        // handler only added `content-type`. The per-request `x-request-id` must
        // be left out of the entry so later hits keep their own fresh value.
        let mut before = HeaderMap::new();
        before.insert("x-request-id", "req-1".parse().unwrap());
        let mut after = before.clone();
        after.insert("content-type", "text/plain".parse().unwrap());

        let stored = handler_response_headers(&before, &after);
        assert!(!stored.contains_key("x-request-id"));
        assert_eq!(stored.get("content-type").unwrap(), "text/plain");
    }

    #[test]
    fn cached_response_skips_when_handler_removes_outer_header() {
        // An outer middleware set `x-default` before the cache; the handler
        // removed it. The entry cannot represent that deletion, so the response
        // must not be cached (otherwise hits would keep the stale `x-default`).
        let mut before = HeaderMap::new();
        before.insert("x-default", "from-outer".parse().unwrap());

        let mut res = Response::new();
        res.body(ResBody::Once(Bytes::from_static(b"cached")));
        // `res` does not carry `x-default`, i.e. the handler dropped it.
        assert!(cached_response(&res, &before, false).is_none());
        assert!(cached_response(&res, &before, true).is_none());

        // Sanity: with no pre-cache header to drop, the same response caches.
        assert!(cached_response(&res, &HeaderMap::new(), false).is_some());
    }

    #[test]
    fn handler_response_headers_includes_handler_overrides() {
        // CORS `CallNext::After` writes a default before the handler runs; the
        // handler overwrites it. The overridden value must be cached so hits
        // reproduce it instead of the pre-cache default.
        let mut before = HeaderMap::new();
        before.insert(
            "access-control-allow-origin",
            "https://default.example".parse().unwrap(),
        );
        let mut after = HeaderMap::new();
        after.insert(
            "access-control-allow-origin",
            "https://handler.example".parse().unwrap(),
        );

        let stored = handler_response_headers(&before, &after);
        assert_eq!(
            stored.get("access-control-allow-origin").unwrap(),
            "https://handler.example"
        );
    }

    #[test]
    fn respond_from_cache_overrides_handler_headers_but_keeps_outer() {
        // The entry only carries handler-produced headers. On a hit they must
        // override same-named pre-cache defaults, while outer headers absent
        // from the entry (a fresh per-request `x-request-id`) are preserved.
        let mut entry_headers = HeaderMap::new();
        entry_headers.insert(
            "access-control-allow-origin",
            "https://handler.example".parse().unwrap(),
        );
        entry_headers.insert("content-type", "text/plain".parse().unwrap());
        let entry = CachedEntry::new(
            Some(StatusCode::OK),
            entry_headers,
            CachedBody::Once(Bytes::from_static(b"cached")),
        );

        let mut res = Response::new();
        res.headers_mut()
            .insert("x-request-id", "fresh-for-this-request".parse().unwrap());
        res.headers_mut().insert(
            "access-control-allow-origin",
            "https://default.example".parse().unwrap(),
        );

        respond_from_cache(&mut res, entry);

        assert_eq!(
            res.headers().get("x-request-id").unwrap(),
            "fresh-for-this-request"
        );
        assert_eq!(
            res.headers().get("access-control-allow-origin").unwrap(),
            "https://handler.example"
        );
        assert_eq!(res.headers().get("content-type").unwrap(), "text/plain");
    }

    #[test]
    fn respond_from_cache_restores_multi_valued_headers() {
        let mut cached_headers = HeaderMap::new();
        cached_headers.append(SET_COOKIE, "a=1".parse().unwrap());
        cached_headers.append(SET_COOKIE, "b=2".parse().unwrap());
        let entry = CachedEntry::new(Some(StatusCode::OK), cached_headers, CachedBody::None);

        let mut res = Response::new();
        respond_from_cache(&mut res, entry);

        let cookies: Vec<_> = res
            .headers()
            .get_all(SET_COOKIE)
            .iter()
            .map(|v| v.to_str().unwrap().to_owned())
            .collect();
        assert_eq!(cookies, vec!["a=1", "b=2"]);
    }

    #[test]
    fn test_cached_entry_clone() {
        let entry = CachedEntry::new(
            Some(StatusCode::OK),
            HeaderMap::new(),
            CachedBody::Once(Bytes::from("test")),
        );
        let cloned = entry.clone();
        assert_eq!(entry.status, cloned.status);
        assert_eq!(entry.body, cloned.body);
    }

    // Tests for Cache
    #[test]
    fn test_cache_new() {
        let cache = Cache::new(MokaStore::<String>::new(100), RequestIssuer::default());
        assert!(format!("{cache:?}").contains("Cache"));
    }

    #[test]
    fn in_flight_limit_bypasses_new_keys_when_full() {
        let in_flight = Arc::new(InFlight::new(1));
        let FlightPermit::Leader(first) = in_flight.enter("a") else {
            panic!("first key should lead an in-flight request");
        };

        assert!(matches!(in_flight.enter("a"), FlightPermit::Follower(_)));
        assert!(matches!(in_flight.enter("b"), FlightPermit::Bypass));

        drop(first);
        assert!(matches!(in_flight.enter("b"), FlightPermit::Leader(_)));
    }

    #[test]
    fn cache_max_in_flight_can_disable_coalescing() {
        let cache =
            Cache::new(MokaStore::<String>::new(100), RequestIssuer::default()).max_in_flight(0);

        assert!(matches!(
            cache.in_flight.enter("uncached".to_owned()),
            FlightPermit::Bypass
        ));
    }

    #[test]
    fn cached_response_skips_private_cache_headers_by_default() {
        // Privacy heuristics only restrict a *shared* cache; a private cache
        // (`cache_private(true)`) may still store these.
        for (name, value) in [(SET_COOKIE, "sid=abc"), (CACHE_CONTROL, "private")] {
            let mut res = Response::new();
            res.body(ResBody::Once(Bytes::from_static(b"cached")));
            res.headers_mut().insert(name, value.parse().unwrap());
            assert!(cached_response(&res, &HeaderMap::new(), false).is_none());
            assert!(cached_response(&res, &HeaderMap::new(), true).is_some());
        }
    }

    #[test]
    fn cached_response_never_stores_vary() {
        // The store never folds the varied request headers into the key, so a
        // `Vary` response cannot be told apart from another representation and is
        // not cached in either mode (otherwise one client's representation would be
        // replayed to clients that negotiated differently).
        let mut res = Response::new();
        res.body(ResBody::Once(Bytes::from_static(b"cached")));
        res.headers_mut()
            .insert(VARY, "accept-language".parse().unwrap());
        assert!(cached_response(&res, &HeaderMap::new(), false).is_none());
        assert!(cached_response(&res, &HeaderMap::new(), true).is_none());
    }

    #[test]
    fn cached_response_never_stores_no_store_or_no_cache() {
        // `no-store`/`no-cache` forbid caching in any cache, so they must be
        // honored even in private-cache mode.
        for value in ["max-age=60, no-store", "no-cache", "public, no-cache"] {
            let mut res = Response::new();
            res.body(ResBody::Once(Bytes::from_static(b"cached")));
            res.headers_mut()
                .insert(CACHE_CONTROL, value.parse().unwrap());
            assert!(cached_response(&res, &HeaderMap::new(), false).is_none());
            assert!(cached_response(&res, &HeaderMap::new(), true).is_none());
        }
    }

    #[tokio::test]
    async fn authorization_requests_are_not_cached_by_default() {
        let calls = Arc::new(AtomicUsize::new(0));
        let cache = Cache::new(
            MokaStore::builder()
                .time_to_live(std::time::Duration::from_secs(60))
                .build(),
            RequestIssuer::default(),
        );
        let router = Arc::new(Router::new().hoop(cache).goal(SlowCached {
            calls: calls.clone(),
        }));

        for _ in 0..2 {
            let mut res = TestClient::get("http://127.0.0.1:5801")
                .add_header(AUTHORIZATION, "Bearer token", true)
                .send(router.clone())
                .await;
            assert_eq!(res.status_code, Some(StatusCode::OK));
            let _ = res.take_string().await.unwrap();
        }

        assert_eq!(calls.load(Ordering::SeqCst), 2);
    }

    #[test]
    fn test_cache_debug() {
        let cache = Cache::new(MokaStore::<String>::new(100), RequestIssuer::default());
        let debug_str = format!("{cache:?}");
        assert!(debug_str.contains("Cache"));
        assert!(debug_str.contains("store"));
        assert!(debug_str.contains("issuer"));
    }

    #[tokio::test]
    async fn test_cache_same_path_same_content() {
        let cache = Cache::new(
            MokaStore::builder()
                .time_to_live(std::time::Duration::from_secs(60))
                .build(),
            RequestIssuer::default(),
        );
        let router = Router::new().hoop(cache).goal(cached);
        let service = Service::new(router);

        let mut res1 = TestClient::get("http://127.0.0.1:5801/same-path")
            .send(&service)
            .await;
        let content1 = res1.take_string().await.unwrap();

        let mut res2 = TestClient::get("http://127.0.0.1:5801/same-path")
            .send(&service)
            .await;
        let content2 = res2.take_string().await.unwrap();

        // Same path should return cached content
        assert_eq!(content1, content2);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn test_cache_coalesces_concurrent_misses() {
        let calls = Arc::new(AtomicUsize::new(0));
        let cache = Cache::new(
            MokaStore::builder()
                .time_to_live(std::time::Duration::from_secs(60))
                .build(),
            RequestIssuer::default(),
        );
        let router = Arc::new(Router::new().hoop(cache).goal(SlowCached {
            calls: calls.clone(),
        }));
        let barrier = Arc::new(tokio::sync::Barrier::new(16));

        let mut tasks = Vec::new();
        for _ in 0..16 {
            let router = router.clone();
            let barrier = barrier.clone();
            tasks.push(tokio::spawn(async move {
                barrier.wait().await;
                let mut res = TestClient::get("http://127.0.0.1:5801").send(router).await;
                res.take_string().await.unwrap()
            }));
        }

        let mut bodies = Vec::new();
        for task in tasks {
            bodies.push(task.await.unwrap());
        }

        assert_eq!(calls.load(Ordering::SeqCst), 1);
        assert!(bodies.iter().all(|body| body == &bodies[0]));
    }
}