cached 4.0.0

/*
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 with `cached::claim::ClaimRegistry`; 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::sync::LazyLock;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;

use cached::claim::ClaimRegistry;
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.
//
// `cached::claim::ClaimRegistry` is that set of keys. `claim` hands the first
// caller a `Claim` and every later caller `None`, until that claim is dropped.
//
// The release comes from the claim's `Drop` rather than from a `release(id)`
// statement at the end of the spawned task, because such a statement 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 panicking and the cancelled path.
//
// `ClaimRegistry::new` is not `const`, so a `static` registry needs `LazyLock`.
// The registry is also cheap to clone, so it can live in a struct field instead.
//
// One registry serves every refresh below only because the ids they claim are
// disjoint ("b", "g", "h", "d", "i"). A registry is a single key space: an id
// shared by two of these caches would let a refresh in one suppress a legitimate
// refresh in the other. Give each key space its own registry when the ids can
// overlap.
// ============================================================================

static REFRESHING: LazyLock<ClaimRegistry<String>> = LazyLock::new(ClaimRegistry::new);

/// Upper bound on a single background refresh. `Claim` has no expiry of its own, and the
/// peek deliberately never removes an expired entry, so a refresh body that hangs (a
/// stuck origin call, say) would otherwise hold its claim forever: every later caller
/// would keep being served the same stale value with nobody able to recompute it. Wrap
/// every spawned refresh in `tokio::time::timeout(REFRESH_TIMEOUT, ..)` so that instead
/// the claim drops - and the key releases - on elapse, the same as it does on a normal
/// completion, a panic, or a cancellation. Named so copy-paste call sites share one
/// tunable instead of each hardcoding their own.
const REFRESH_TIMEOUT: Duration = Duration::from_secs(5);

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) = REFRESHING.claim(id.to_string()) {
                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.
                    // Borrowing the key out of the guard is what keeps it alive for
                    // the duration of the call the borrow is passed into: the guard
                    // cannot be dropped while that borrow is live. It is dropped, and
                    // the key released, on completion, on an unwind out of the
                    // refresh, and on cancellation alike - including here, where the
                    // timeout elapsing drops the future (and with it the claim)
                    // instead of ever finishing the call.
                    if tokio::time::timeout(REFRESH_TIMEOUT, async_lookup_prime_cache(claim.key()))
                        .await
                        .is_err()
                    {
                        eprintln!("refresh for {:?} timed out; claim released", claim.key());
                    }
                });
            }
            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 = REFRESHING.claim(id.to_string())?;
    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
/// 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 the `Claim`'s `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 = REFRESHING.claim(id.to_string())?;
    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 registry as section 4: released by drop, so a refresh that
            // panics or is cancelled cannot wedge the key.
            if let Some(claim) = REFRESHING.claim(id.to_string()) {
                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 claim releases under `#[cached]`; this is the same proof
/// under `sync_writes = "by_key"`, the other call site that takes a claim, 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 = REFRESHING.claim(id.to_string())?;
    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 releases from `Drop`, 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 = REFRESHING.claim("h".to_string());
    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);
    assert!(
        !REFRESHING.is_claimed("h"),
        "the registry must drain the key once the retry claim is dropped too"
    );

    // --- 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"
    );

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