tidaldb/tidal/src/signals/ledger/hot_top_k.rs

//! Cached per-signal-type top-K (m12p2): bound `SignalRanked` candidate
//! generation to O(K) per query instead of an O(N) ledger scan on every call.
//!
//! # Why a cache is correct here
//!
//! A `SignalRanked` strategy ranks candidates by a signal's decayed score. The
//! score of an entity at query time is `accumulated * exp(-lambda * dt)`, and the
//! SAME `lambda` applies to every entity of a given signal type, so time decay
//! multiplies every entity's score by the same factor — it **never changes the
//! relative order**. The membership of the top-K therefore only changes when a
//! NEW signal write bumps some entity's accumulated score. That makes a cached
//! top-K valid until the next write, not until the next clock tick.
//!
//! # Freshness policy
//!
//! - **Small ledgers** (`< SMALL_LEDGER_ENTRIES`): the O(N) rebuild is
//!   sub-millisecond, so we rebuild whenever a write has landed since the cache
//!   was built — always fresh, which is what in-process tests and small apps
//!   expect.
//! - **Large ledgers** (`>= SMALL_LEDGER_ENTRIES`): a continuous-write workload
//!   would otherwise land an O(N) scan on the read hot path and blow the p99, so
//!   rebuilds are throttled to at most one per [`REFRESH_INTERVAL_NS`]. Trending
//!   tolerates a second of staleness; the p99 SLO does not tolerate a 50ms scan.

use std::sync::atomic::{AtomicU64, Ordering};

use dashmap::DashMap;

use super::super::SignalTypeId;
use super::types::EntitySignalEntry;
use crate::schema::EntityId;

/// Ledgers smaller than this rebuild the cache on every post-write query (always
/// fresh; the scan is cheap at this size). At or above it, rebuilds throttle to
/// [`REFRESH_INTERVAL_NS`] to keep the O(N) scan off the read hot path.
const SMALL_LEDGER_ENTRIES: usize = 50_000;

/// Minimum wall-clock between rebuilds for a large ledger (1 second).
const REFRESH_INTERVAL_NS: u64 = 1_000_000_000;

/// How many entities to materialize per signal type. Bounds cache memory (a few
/// thousand `EntityId`s per signal type) and caps the candidate pool any single
/// `SignalRanked` query draws from. Comfortably above `MAX_LIMIT × 4`.
const REBUILD_K: usize = 4096;

/// One signal type's materialized top-K plus the bookkeeping that decides
/// freshness.
struct CachedTopK {
    /// Top entities by decayed score, best-first (length ≤ [`REBUILD_K`]).
    items: Vec<EntityId>,
    /// Wall-clock (ns) the cache was built — drives the large-ledger throttle.
    built_at_ns: u64,
    /// The global write counter at build time — if it still matches, no signal
    /// has been written since, so the cache is exact regardless of elapsed time.
    built_dirty: u64,
}

/// Per-signal-type cached top-K, with a global write counter for invalidation.
pub struct HotTopKCache {
    per_type: DashMap<SignalTypeId, CachedTopK>,
    /// Bumped once per signal apply (one relaxed atomic add — negligible on the
    /// write path). A cache built at an earlier value knows writes have landed.
    dirty: AtomicU64,
}

impl HotTopKCache {
    pub fn new() -> Self {
        Self {
            per_type: DashMap::new(),
            dirty: AtomicU64::new(0),
        }
    }

    /// Record that a signal was written (invalidates caches on their next read).
    pub fn note_write(&self) {
        self.dirty.fetch_add(1, Ordering::Relaxed);
    }

    /// Top entities for `type_id` by decayed score, best-first, capped to the
    /// caller's `needed` count. Serves a fresh-enough cache in O(needed); rebuilds
    /// (one bounded O(N) scan) only when stale per the freshness policy.
    pub fn candidates(
        &self,
        entries: &DashMap<(EntityId, SignalTypeId), EntitySignalEntry>,
        type_id: SignalTypeId,
        needed: usize,
        now_ns: u64,
    ) -> Vec<EntityId> {
        let needed = needed.min(REBUILD_K);
        let dirty_now = self.dirty.load(Ordering::Relaxed);

        if let Some(cached) = self.per_type.get(&type_id) {
            let fresh = if cached.built_dirty == dirty_now {
                // No writes at all since the build ⇒ exact, regardless of clock.
                true
            } else if entries.len() < SMALL_LEDGER_ENTRIES {
                // Small ledger + new writes ⇒ rebuild (cheap, always fresh).
                false
            } else {
                // Large ledger ⇒ throttle the O(N) rebuild off the hot path.
                now_ns.saturating_sub(cached.built_at_ns) < REFRESH_INTERVAL_NS
            };
            if fresh {
                return cached.items.iter().take(needed).copied().collect();
            }
        }

        let items = rebuild(entries, type_id, now_ns);
        let out = items.iter().take(needed).copied().collect();
        self.per_type.insert(
            type_id,
            CachedTopK {
                items,
                built_at_ns: now_ns,
                built_dirty: dirty_now,
            },
        );
        out
    }
}

/// The bounded O(N) ledger scan: the top [`REBUILD_K`] entities of `type_id` by
/// decayed score, best-first. Memory stays O(K): a buffer capped at `2·K` is
/// partition-truncated to K whenever it fills, so peak allocation is bounded
/// regardless of how many entities carry the signal.
fn rebuild(
    entries: &DashMap<(EntityId, SignalTypeId), EntitySignalEntry>,
    type_id: SignalTypeId,
    now_ns: u64,
) -> Vec<EntityId> {
    let cap = REBUILD_K;
    // Descending by score; NaN sinks via a deterministic id tiebreak.
    let cmp = |a: &(EntityId, f64), b: &(EntityId, f64)| {
        b.1.partial_cmp(&a.1)
            .unwrap_or_else(|| b.0.as_u64().cmp(&a.0.as_u64()))
    };
    let buffer_cap = cap.saturating_mul(2).max(cap + 1);
    let mut scored: Vec<(EntityId, f64)> = Vec::with_capacity(buffer_cap);

    for entry in entries {
        let (entity_id, signal_type_id) = entry.key();
        if *signal_type_id == type_id {
            // Decay score with lambda=0 (no extra decay beyond what was applied at
            // write time) — the same simplified candidate-gen ranking the prior
            // O(N) scan used. Stage 3 does the full scoring.
            let score = entry.value().hot.current_score(0, now_ns, 0.0);
            scored.push((*entity_id, score));
            if scored.len() >= buffer_cap {
                scored.select_nth_unstable_by(cap - 1, cmp);
                scored.truncate(cap);
            }
        }
    }
    if scored.len() > cap {
        scored.select_nth_unstable_by(cap - 1, cmp);
        scored.truncate(cap);
    }
    scored.sort_unstable_by(cmp);
    scored.into_iter().map(|(id, _)| id).collect()
}