cached 3.1.1

/*
Stale-while-revalidate: serve the expired value immediately and refresh it off the
critical path. Composed from existing pieces, with no dedicated attribute.

Two primitives do the work:

1. `cache_peek_with_expiry_status(&key) -> (Option<V>, bool)` returns the stored value
   TOGETHER with whether it has expired, and (being the peek variant) removes nothing,
   updates no recency, and counts no hit or miss. An expired entry reads as
   `(Some(value), true)` rather than `None`, which is what makes it servable.

2. `{fn}_prime_cache(args)` recomputes and stores under the same key. It runs the
   function body BEFORE taking the cache write lock, so a refresh never blocks the
   readers being served the stale value.

The cache static generated by the macro carries the cached function's own visibility,
so both are reachable from the call site.

Spawning is left to the caller on purpose: `cached` has no runtime dependency, so it
cannot pick tokio or smol on your behalf. That also means you decide the policy -
whether to rate-limit refreshes, drop them under load, or run them on a dedicated pool.

The static's shape depends on the macro, on whether the function is async, and on
`sync_writes`, so the sections below spell out each one:

    sync  #[cached]            LazyLock<parking_lot::RwLock<Store>>   .read()
    async #[cached]            LazyLock<async_lock::RwLock<Store>>    .read().await
    async #[concurrent_cached] OnceCell<Store>                        .get() (no guard)
    sync_writes = "by_key"     ..wrapped in KeyedCache<Lock, _>       derefs to the lock

Sections 1 to 3 cover the read side per shape; section 4 collapses concurrent refreshes
onto one caller; section 5 combines that with `sync_writes = "by_key"` so the cold path
deduplicates too, which is per-key write synchronization that serves stale instead of
blocking.

Run:
    cargo run --example stale_while_revalidate --features "proc_macro,time_stores,async"
*/

use std::collections::HashSet;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Mutex, PoisonError};
use std::time::Duration;

use cached::macros::{cached, concurrent_cached};
use cached::{CloneCached, ConcurrentCloneCached};

/// Counts real executions of a function body, so the output distinguishes a served
/// value from a recompute.
static COMPUTES: AtomicUsize = AtomicUsize::new(0);

fn compute_count() -> usize {
    COMPUTES.load(Ordering::SeqCst)
}

// ============================================================================
// 1. Synchronous `#[cached]`
//
// The static is a `parking_lot::RwLock`, so the read guard is taken without
// `.await`. Drop it before spawning: the refresh will want the write lock.
// ============================================================================

#[cached(ttl_secs = 1, key = "String", convert = r#"{ id.to_string() }"#)]
fn sync_lookup(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    std::thread::sleep(Duration::from_millis(200));
    format!("{id}-v{n}")
}

fn sync_lookup_swr(id: &str) -> String {
    let (value, expired) = {
        let cache = SYNC_LOOKUP.read();
        // The key must be built exactly as the `convert` block builds it. If the two
        // ever drift you silently peek a different entry, so keep them side by side.
        cache.cache_peek_with_expiry_status(&id.to_string())
    };

    match value {
        // Live entry: nothing to do.
        Some(v) if !expired => v,
        // Expired entry: serve it now, recompute on a thread we own.
        Some(stale) => {
            let owned = id.to_string();
            std::thread::spawn(move || {
                sync_lookup_prime_cache(&owned);
            });
            stale
        }
        // Nothing cached yet. There is no stale value to serve, so this call waits.
        None => sync_lookup(id),
    }
}

// ============================================================================
// 2. Asynchronous `#[cached]`
//
// Identical shape, except the static is an `async_lock::RwLock`, so the read is
// `.read().await`. Bind the peek in its own scope so the guard is dropped before
// the `.await` on the spawn - holding it across an await point would serialize
// every reader against the refresh.
// ============================================================================

#[cached(ttl_secs = 1, key = "String", convert = r#"{ id.to_string() }"#)]
async fn async_lookup(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    tokio::time::sleep(Duration::from_millis(200)).await;
    format!("{id}-v{n}")
}

async fn async_lookup_swr(id: &str) -> String {
    let (value, expired) = {
        let cache = ASYNC_LOOKUP.read().await;
        cache.cache_peek_with_expiry_status(&id.to_string())
    };

    match value {
        Some(v) if !expired => v,
        Some(stale) => {
            let owned = id.to_string();
            tokio::spawn(async move {
                async_lookup_prime_cache(&owned).await;
            });
            stale
        }
        None => async_lookup(id).await,
    }
}

// ============================================================================
// 3. Asynchronous `#[concurrent_cached]`
//
// The sharded stores synchronize internally, so the static is the store itself
// behind a `OnceCell` (initialized by the first call) with no outer lock: reach
// through the cell and peek directly.
// ============================================================================

#[concurrent_cached(ttl_secs = 1, key = "String", convert = r#"{ id.to_string() }"#)]
async fn sharded_lookup(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    tokio::time::sleep(Duration::from_millis(200)).await;
    format!("{id}-v{n}")
}

async fn sharded_lookup_swr(id: &str) -> String {
    // `get()` is `None` only before the first call has initialized the cell, which
    // is indistinguishable from a miss for our purposes.
    let peeked = SHARDED_LOOKUP
        .get()
        .map(|cache| cache.cache_peek_with_expiry_status(&id.to_string()));

    match peeked {
        Some((Some(v), false)) => v,
        Some((Some(stale), true)) => {
            let owned = id.to_string();
            tokio::spawn(async move {
                sharded_lookup_prime_cache(&owned).await;
            });
            stale
        }
        _ => sharded_lookup(id).await,
    }
}

// ============================================================================
// 4. Collapsing concurrent refreshes
//
// Every caller that observes the same stale entry spawns its own refresh, so N
// concurrent readers cause N recomputes. `sync_writes = "by_key"` does NOT help
// here: its per-key lock covers the store write, not the function body (that is
// what keeps a refresh from blocking readers). When the recompute is expensive,
// track the keys already being refreshed and let the first caller win.
//
// The claim is an RAII guard rather than a `release_refresh(id)` call at the end
// of the spawned task, because that call is skipped when the refresh body panics
// or the task is aborted. The key would then stay claimed for the life of the
// process, and since the peek deliberately never removes an expired entry, it
// would be served stale forever with no caller able to recompute it: strictly
// worse than not deduplicating at all. `main` asserts the guard.
// ============================================================================

static REFRESHING: Mutex<Option<HashSet<String>>> = Mutex::new(None);

/// An in-flight refresh claim on one key, released when dropped.
struct RefreshClaim {
    key: String,
}

impl RefreshClaim {
    fn key(&self) -> &str {
        &self.key
    }
}

impl Drop for RefreshClaim {
    fn drop(&mut self) {
        release_refresh(&self.key);
    }
}

/// Returns a claim if this caller is the one that must refresh `key`, or `None` if
/// a refresh is already in flight.
fn claim_refresh(key: &str) -> Option<RefreshClaim> {
    let mut guard = REFRESHING.lock().expect("refresh set poisoned");
    if guard
        .get_or_insert_with(HashSet::new)
        .insert(key.to_string())
    {
        Some(RefreshClaim {
            key: key.to_string(),
        })
    } else {
        None
    }
}

fn release_refresh(key: &str) {
    // Reached from `Drop`, including while a panic unwinds, where a second panic
    // would abort the process. So recover a poisoned lock instead of unwrapping it.
    let mut guard = REFRESHING.lock().unwrap_or_else(PoisonError::into_inner);
    if let Some(set) = guard.as_mut() {
        set.remove(key);
    }
}

async fn async_lookup_swr_deduped(id: &str) -> String {
    let (value, expired) = {
        let cache = ASYNC_LOOKUP.read().await;
        cache.cache_peek_with_expiry_status(&id.to_string())
    };

    match value {
        Some(v) if !expired => v,
        Some(stale) => {
            if let Some(claim) = claim_refresh(id) {
                tokio::spawn(async move {
                    // `claim` lives until this task ends, so a caller arriving
                    // mid-refresh sees the claim rather than starting a second one.
                    // It is dropped, and the key released, on completion, on an
                    // unwind out of the refresh, and on cancellation alike.
                    async_lookup_prime_cache(claim.key()).await;
                });
            }
            stale
        }
        None => async_lookup(id).await,
    }
}

/// Silences the default panic handler for its lifetime, restoring the previous hook on
/// drop rather than at a fixed point after some `.await`. A plain `take_hook` /
/// `set_hook(previous)` pair only restores if nothing between the two panics; here that
/// gap is meant to stay empty (only the awaited outcome is expected to panic, inside the
/// spawned task, not in `main`'s own thread), but a guard costs nothing and does not rely
/// on that staying true.
type PanicHook = dyn Fn(&std::panic::PanicHookInfo<'_>) + Sync + Send;

struct SilencedPanic {
    previous: Option<Box<PanicHook>>,
}

impl SilencedPanic {
    fn new() -> Self {
        let previous = std::panic::take_hook();
        std::panic::set_hook(Box::new(|_| {}));
        Self {
            previous: Some(previous),
        }
    }
}

impl Drop for SilencedPanic {
    fn drop(&mut self) {
        if let Some(previous) = self.previous.take() {
            std::panic::set_hook(previous);
        }
    }
}

/// Set to make the next `flaky_lookup` body panic, so `main` can drive a refresh
/// that fails partway through.
static PANIC_NEXT_REFRESH: AtomicBool = AtomicBool::new(false);

#[cached(ttl_secs = 2, key = "String", convert = r#"{ id.to_string() }"#)]
async fn flaky_lookup(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    assert!(
        !PANIC_NEXT_REFRESH.swap(false, Ordering::SeqCst),
        "simulated refresh failure"
    );
    format!("{id}-v{n}")
}

/// The spawn shape above, reduced to the claim and the refresh, so `main` can await
/// the outcome. Returns `None` when a refresh is already in flight.
fn spawn_refresh(id: &str) -> Option<tokio::task::JoinHandle<()>> {
    let claim = claim_refresh(id)?;
    Some(tokio::spawn(async move {
        flaky_lookup_prime_cache(claim.key()).await;
    }))
}

/// A body slow enough that `main` can reliably abort its task mid-flight, to prove the
/// RAII claim also releases when the task is CANCELLED rather than unwound by a panic.
/// `claim` lives inside the spawned future's own captured state, so `JoinHandle::abort`
/// dropping that future in place runs `RefreshClaim::drop` exactly the same as an unwind
/// or a normal return does.
#[cached(ttl_secs = 2, key = "String", convert = r#"{ id.to_string() }"#)]
async fn stuck_lookup(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    tokio::time::sleep(Duration::from_secs(5)).await;
    format!("{id}-v{n}")
}

/// The spawn shape above, over `stuck_lookup`, so `main` can abort the handle mid-flight.
fn spawn_stuck_refresh(id: &str) -> Option<tokio::task::JoinHandle<()>> {
    let claim = claim_refresh(id)?;
    Some(tokio::spawn(async move {
        stuck_lookup_prime_cache(claim.key()).await;
    }))
}

// ============================================================================
// 5. Single-flight revalidation
//
// Section 4 collapses the refreshes but leaves the cold path alone: with no
// cached value at all, concurrent first callers each run the function. Adding
// `sync_writes = "by_key"` closes that, and the two policies compose into one
// rule per state:
//
//   cold  (nothing cached) -> callers deduplicate and WAIT; there is nothing to
//                             serve, and `by_key` makes exactly one of them run
//                             the body while the rest reuse its result.
//   stale (expired entry)  -> callers never wait. One claims the refresh, and
//                             every caller (claimant included) returns the stale
//                             value immediately.
//
// That is per-key write synchronization that degrades to serving stale instead
// of blocking, which is what a stale-while-revalidate policy usually means.
//
// Note that `sync_writes = "by_key"` wraps the cache static in a `KeyedCache`,
// which derefs to the same lock, so the peek below is unchanged.
// ============================================================================

#[cached(
    ttl_secs = 1,
    key = "String",
    convert = r#"{ id.to_string() }"#,
    sync_writes = "by_key"
)]
async fn single_flight(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    tokio::time::sleep(Duration::from_millis(200)).await;
    format!("{id}-v{n}")
}

async fn single_flight_swr(id: &str) -> String {
    let (value, expired) = {
        let cache = SINGLE_FLIGHT.read().await;
        cache.cache_peek_with_expiry_status(&id.to_string())
    };

    match value {
        Some(v) if !expired => v,
        Some(stale) => {
            // The same RAII claim as section 4: released by drop, so a refresh that
            // panics or is cancelled cannot wedge the key.
            if let Some(claim) = claim_refresh(id) {
                tokio::spawn(async move {
                    single_flight_prime_cache(claim.key()).await;
                });
            }
            stale
        }
        // The only blocking path, and only because there is nothing to serve.
        None => single_flight(id).await,
    }
}

/// Set to make the next `flaky_single_flight` body panic. Section 4's `flaky_lookup`
/// already proves the RAII claim releases under `#[cached]`; this is the same proof
/// under `sync_writes = "by_key"`, the other call site the fix touched, so a
/// regression specific to that site (say, a `by_key` bucket lock held across the
/// panic instead of released with the claim) would not go unnoticed.
static PANIC_NEXT_SINGLE_FLIGHT_REFRESH: AtomicBool = AtomicBool::new(false);

#[cached(
    ttl_secs = 2,
    key = "String",
    convert = r#"{ id.to_string() }"#,
    sync_writes = "by_key"
)]
async fn flaky_single_flight(id: &str) -> String {
    let n = COMPUTES.fetch_add(1, Ordering::SeqCst) + 1;
    assert!(
        !PANIC_NEXT_SINGLE_FLIGHT_REFRESH.swap(false, Ordering::SeqCst),
        "simulated single-flight refresh failure"
    );
    format!("{id}-v{n}")
}

/// `spawn_refresh`'s shape, over `flaky_single_flight`.
fn spawn_flaky_single_flight_refresh(id: &str) -> Option<tokio::task::JoinHandle<()>> {
    let claim = claim_refresh(id)?;
    Some(tokio::spawn(async move {
        flaky_single_flight_prime_cache(claim.key()).await;
    }))
}

#[tokio::main]
async fn main() {
    // --- 1. synchronous -----------------------------------------------------
    println!("sync `#[cached]`");
    let first = sync_lookup_swr("a");
    println!("  cold miss  -> {first} ({} computes)", compute_count());

    let hit = sync_lookup_swr("a");
    println!("  fresh hit  -> {hit} ({} computes)", compute_count());

    std::thread::sleep(Duration::from_millis(1_100));
    let started = std::time::Instant::now();
    let stale = sync_lookup_swr("a");
    println!(
        "  stale hit  -> {stale} returned in {}ms, refresh spawned",
        started.elapsed().as_millis()
    );
    assert_eq!(stale, first, "the stale hit must serve the previous value");

    std::thread::sleep(Duration::from_millis(400));
    let refreshed = sync_lookup_swr("a");
    println!("  after bg   -> {refreshed} ({} computes)", compute_count());
    assert_ne!(
        refreshed, first,
        "the background refresh must have replaced it"
    );

    // --- 2. asynchronous ----------------------------------------------------
    println!("\nasync `#[cached]`");
    let first = async_lookup_swr("b").await;
    println!("  cold miss  -> {first}");

    tokio::time::sleep(Duration::from_millis(1_100)).await;
    let started = std::time::Instant::now();
    let stale = async_lookup_swr("b").await;
    println!(
        "  stale hit  -> {stale} returned in {}ms, refresh spawned",
        started.elapsed().as_millis()
    );
    assert_eq!(stale, first);

    tokio::time::sleep(Duration::from_millis(400)).await;
    println!("  after bg   -> {}", async_lookup_swr("b").await);

    // --- 3. sharded ---------------------------------------------------------
    println!("\nasync `#[concurrent_cached]`");
    let first = sharded_lookup_swr("c").await;
    println!("  cold miss  -> {first}");

    tokio::time::sleep(Duration::from_millis(1_100)).await;
    let stale = sharded_lookup_swr("c").await;
    println!("  stale hit  -> {stale}, refresh spawned");
    assert_eq!(stale, first);

    tokio::time::sleep(Duration::from_millis(400)).await;
    println!("  after bg   -> {}", sharded_lookup_swr("c").await);

    // --- 4. deduplicated refresh -------------------------------------------
    println!("\ncollapsing concurrent refreshes");
    tokio::time::sleep(Duration::from_millis(1_100)).await;

    let before = compute_count();
    let mut handles = Vec::new();
    for _ in 0..8 {
        handles.push(tokio::spawn(async { async_lookup_swr_deduped("b").await }));
    }
    for handle in handles {
        handle.await.expect("refresh task panicked");
    }
    tokio::time::sleep(Duration::from_millis(400)).await;
    println!(
        "  8 concurrent stale readers -> {} recompute(s)",
        compute_count() - before
    );
    assert_eq!(
        compute_count() - before,
        1,
        "the in-flight guard must collapse the refreshes to one"
    );

    // The claim is RAII, so a refresh that panics still releases its key. Were it
    // released by a statement at the end of the task instead, `g` would stay claimed
    // for the life of the process and could never be refreshed again.
    let cold = flaky_lookup("g").await;
    PANIC_NEXT_REFRESH.store(true, Ordering::SeqCst);
    let outcome = {
        let _silence = SilencedPanic::new(); // the panic below is deliberate
        let failing = spawn_refresh("g").expect("the first caller must claim the refresh");
        failing.await
    };
    assert!(outcome.is_err(), "the refresh task must have panicked");

    let retry = spawn_refresh("g").expect("a panicking refresh must not wedge the key");
    retry.await.expect("the retry must not panic");
    let stored = {
        let cache = FLAKY_LOOKUP.read().await;
        cache.cache_peek_with_expiry_status(&"g".to_string()).0
    };
    println!("  a panicking refresh released its claim, retry stored {stored:?}");
    assert!(
        stored.is_some(),
        "the retried refresh must have stored a value"
    );
    assert_ne!(
        stored.as_deref(),
        Some(cold.as_str()),
        "the retry must have replaced the value the panicking refresh failed to"
    );

    // The other way a spawned task can stop running without reaching its own end: the
    // caller aborts or drops the `JoinHandle` instead of the body panicking. A cancelled
    // task's future is dropped just like any other, so the guard releases here too.
    let stuck = spawn_stuck_refresh("h").expect("the first caller must claim the refresh");
    // Let the task actually start running (and park inside its sleep) before aborting
    // it - aborting a task the runtime never polled would prove nothing about a claim
    // held mid-execution.
    tokio::time::sleep(Duration::from_millis(50)).await;
    stuck.abort();
    let cancelled = stuck.await;
    assert!(
        cancelled.as_ref().is_err_and(|e| e.is_cancelled()),
        "the task must have been cancelled, not merely finished on its own: {cancelled:?}"
    );
    let retry_claim = claim_refresh("h");
    println!(
        "  an aborted refresh released its claim: retry claim succeeded {}",
        retry_claim.is_some()
    );
    assert!(
        retry_claim.is_some(),
        "an aborted refresh must release its claim, not wedge the key forever"
    );
    drop(retry_claim);

    // --- 5. single-flight revalidation --------------------------------------
    println!("\nsingle-flight revalidation");

    // Cold: 8 callers, nothing cached. They deduplicate and wait together, so the
    // body runs once and every caller gets the same value.
    let before = compute_count();
    let started = std::time::Instant::now();
    let mut handles = Vec::new();
    for _ in 0..8 {
        handles.push(tokio::spawn(async { single_flight_swr("d").await }));
    }
    let mut cold = Vec::new();
    for handle in handles {
        cold.push(handle.await.expect("cold task panicked"));
    }
    println!(
        "  cold : 8 callers -> {} compute(s) in {}ms, all callers agree: {}",
        compute_count() - before,
        started.elapsed().as_millis(),
        cold.iter().all(|v| *v == cold[0])
    );
    assert_eq!(
        compute_count() - before,
        1,
        "`by_key` must dedupe the cold path"
    );
    assert!(cold.iter().all(|v| *v == cold[0]));

    // Stale: the same 8 callers, now with an expired entry. Nobody waits.
    tokio::time::sleep(Duration::from_millis(1_100)).await;
    let before = compute_count();
    let started = std::time::Instant::now();
    let mut handles = Vec::new();
    for _ in 0..8 {
        handles.push(tokio::spawn(async { single_flight_swr("d").await }));
    }
    let mut served = Vec::new();
    for handle in handles {
        served.push(handle.await.expect("stale task panicked"));
    }
    let stale_ms = started.elapsed().as_millis();
    println!(
        "  stale: 8 callers -> served {} in {}ms without waiting",
        served[0], stale_ms
    );
    assert!(stale_ms < 100, "a stale read must not wait for the refresh");
    assert_eq!(
        served[0], cold[0],
        "the stale reads must serve the old value"
    );

    tokio::time::sleep(Duration::from_millis(400)).await;
    println!(
        "  after: {} refresh(es), next read -> {}",
        compute_count() - before,
        single_flight_swr("d").await
    );
    assert_eq!(
        compute_count() - before,
        1,
        "exactly one caller may revalidate"
    );

    // Section 4 proved the panicking-refresh case under plain `#[cached]`; the fix
    // touched the `sync_writes = "by_key"` call site too, and nothing above drives a
    // panic through it, so repeat the same proof here.
    let cold_flaky = flaky_single_flight("i").await;
    PANIC_NEXT_SINGLE_FLIGHT_REFRESH.store(true, Ordering::SeqCst);
    let outcome = {
        let _silence = SilencedPanic::new(); // the panic below is deliberate
        let failing = spawn_flaky_single_flight_refresh("i")
            .expect("the first caller must claim the refresh");
        failing.await
    };
    assert!(
        outcome.is_err(),
        "the single-flight refresh task must have panicked"
    );

    let retry = spawn_flaky_single_flight_refresh("i")
        .expect("a panicking single-flight refresh must not wedge the key");
    retry.await.expect("the retry must not panic");
    let stored = {
        let cache = FLAKY_SINGLE_FLIGHT.read().await;
        cache.cache_peek_with_expiry_status(&"i".to_string()).0
    };
    println!("  a panicking single-flight refresh released its claim, retry stored {stored:?}");
    assert!(
        stored.is_some(),
        "the retried single-flight refresh must have stored a value"
    );
    assert_ne!(
        stored.as_deref(),
        Some(cold_flaky.as_str()),
        "the retry must have replaced the value the panicking refresh failed to"
    );

    // Poison recovery: `release_refresh` runs from `Drop`, including mid-unwind, where
    // a second panic would abort the process, so it recovers a poisoned `REFRESHING`
    // instead of unwrapping it. Nothing above ever poisons the lock, so poison it here
    // directly. This must run last - poisoning is permanent for the mutex's lifetime,
    // and `claim_refresh` (unlike `release_refresh`) still `.expect`s a healthy lock, so
    // any assertion after this point would panic there instead of testing what it means to.
    let claim = claim_refresh("poison-me").expect("the first claim on a fresh key must succeed");
    let poisoned = {
        let _silence = SilencedPanic::new(); // the panic below is deliberate
        std::panic::catch_unwind(|| {
            let _guard = REFRESHING.lock().expect("not yet poisoned");
            panic!("deliberately poison REFRESHING while a guard is held");
        })
    };
    assert!(
        poisoned.is_err(),
        "the deliberate panic must have propagated"
    );
    assert!(REFRESHING.is_poisoned(), "REFRESHING must now be poisoned");

    // Runs `RefreshClaim::drop` -> `release_refresh` against the now-poisoned mutex.
    // Were `release_refresh` still `.expect`-ing a healthy lock, this drop would panic
    // during a `Drop`, and if it happened during an unwind, that second panic would
    // abort the process rather than merely failing this assertion.
    drop(claim);
    let recovered = REFRESHING.lock().unwrap_or_else(PoisonError::into_inner);
    let still_claimed = recovered
        .as_ref()
        .is_some_and(|set| set.contains("poison-me"));
    println!("  release_refresh recovered a poisoned lock: still claimed = {still_claimed}");
    assert!(
        !still_claimed,
        "release_refresh must remove the key even from a poisoned set"
    );

    println!("\nstale-while-revalidate ok");
}