benilla-octo/crates/benilla-app/src/query_cache.rs

//! `QueryCache<K, V>`, the ask-once cache: the reference's `DBCache<T>` (vtable `0x80912c`). A
//! miss sends the query once, a second lookup of a pending key does not re-send, the response
//! lands the record, and eviction is explicit only.
//!
//! The read that asks takes `&self` (the in-flight set is behind a lock), so read-only systems
//! share the owning resource and its change detection means an answer landed.
//!
//! A negative answer is cached as `None` and never re-asked: an integer-keyed store's high-bit miss
//! (`0x556e67`), or a player's empty name.

use std::collections::HashMap;
use std::hash::Hash;
use std::sync::Mutex;

use benilla_assets::LockRecover;
use bevy::prelude::*;

use crate::net::EnteredWorldMessage;

pub(crate) struct QueryCache<K, V> {
    answered: HashMap<K, Option<V>>,
    /// The keys in flight, locked so the read that asks is `&self`; uncontended on the main
    /// thread.
    pending: Mutex<HashMap<K, Ask>>,
    /// Bumped by every landing or eviction, the reference's `DBCACHECALLBACK` redisplay.
    generation: u64,
}

/// One in-flight ask, for [`QueryCache::sweep_pending`].
#[derive(Default)]
struct Ask {
    /// When the ask was last sent, as the sweep's own clock reads it; `None` until the first
    /// sweep sees it, since `get_or_ask` has no clock to stamp it with.
    last: Option<f64>,
    /// How many times it has been sent, the first included.
    tries: u8,
    /// Whether giving up on it has been reported, so the warning is said once.
    abandoned: bool,
}

/// What a [`QueryCache::sweep_pending`] pass found.
pub(crate) struct AskSweep<K> {
    /// Asks past their deadline with tries left: re-send these.
    pub(crate) retry: Vec<K>,
    /// Asks past their deadline with no tries left, reported once each.
    pub(crate) abandoned: Vec<K>,
}

impl<K: Copy + Eq + Hash, V> Default for QueryCache<K, V> {
    fn default() -> Self {
        Self {
            answered: HashMap::new(),
            pending: Mutex::new(HashMap::new()),
            generation: 0,
        }
    }
}

impl<K: Copy + Eq + Hash, V> QueryCache<K, V> {
    /// The answer for `key`; a miss runs `ask` once per key per connection. `None` for a miss
    /// and a cached negative alike.
    pub(crate) fn get_or_ask(&self, key: K, ask: impl FnOnce()) -> Option<&V> {
        match self.answered.get(&key) {
            Some(answer) => answer.as_ref(),
            None => {
                let mut pending = self.pending.lock_recover();
                if let std::collections::hash_map::Entry::Vacant(slot) = pending.entry(key) {
                    slot.insert(Ask {
                        last: None,
                        tries: 1,
                        abandoned: false,
                    });
                    drop(pending);
                    ask();
                }
                None
            }
        }
    }

    /// The answer if it is already here, never asking.
    pub(crate) fn get(&self, key: K) -> Option<&V> {
        self.answered.get(&key)?.as_ref()
    }

    /// The answer, mutably, for a record the wire patches in place; never asking.
    pub(crate) fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.answered.get_mut(&key)?.as_mut()
    }

    /// Whether the server answered `key`, with a record or "unknown".
    pub(crate) fn answered(&self, key: K) -> bool {
        self.answered.contains_key(&key)
    }

    /// Whether the server answered "unknown", as opposed to a still-pending ask.
    pub(crate) fn answered_unknown(&self, key: K) -> bool {
        self.answered.get(&key).is_some_and(Option::is_none)
    }

    #[cfg(test)]
    /// Whether an ask for `key` is in flight.
    pub(crate) fn is_pending(&self, key: K) -> bool {
        self.pending.lock_recover().contains_key(&key)
    }

    /// Age the in-flight asks against `now` and report the ones that have waited longer than
    /// `after` seconds: those with tries left (re-send them), and those without, once each.
    ///
    /// **Why an ask-once cache needs this at all.** The reference re-asks nothing either, and on
    /// a client that never loses a response it never has to. Here a query that goes out and is
    /// never answered leaves its key answered-never: `GetItemInfo` keeps returning nil, so the
    /// item has no name, no icon and no tooltip for the rest of the session, and nothing in the
    /// UI ever asks again. grom, 2026-11-08, on the quest item in bag 0 slot 2: *"on the quest
    /// Deeprun Rat Roundup, I'm not recieving the flute item needed for the quest"* — he had it,
    /// entry 17117, and it was blank.
    ///
    /// The first sweep to see an ask only stamps it, because `get_or_ask` takes `&self` and has
    /// no clock; the deadline therefore runs from that stamp, which is at most one sweep period
    /// after the send.
    pub(crate) fn sweep_pending(&self, now: f64, after: f64, max_tries: u8) -> AskSweep<K> {
        let mut sweep = AskSweep {
            retry: Vec::new(),
            abandoned: Vec::new(),
        };
        for (key, ask) in self.pending.lock_recover().iter_mut() {
            let Some(last) = ask.last else {
                ask.last = Some(now);
                continue;
            };
            if now - last < after {
                continue;
            }
            if ask.tries < max_tries {
                ask.tries += 1;
                ask.last = Some(now);
                sweep.retry.push(*key);
            } else if !ask.abandoned {
                ask.abandoned = true;
                sweep.abandoned.push(*key);
            }
        }
        sweep
    }

    /// Lands an answer; `None` is the server's negative.
    pub(crate) fn insert(&mut self, key: K, answer: Option<V>) {
        self.pending
            .get_mut()
            .unwrap_or_else(|p| p.into_inner())
            .remove(&key);
        self.answered.insert(key, answer);
        self.generation = self.generation.wrapping_add(1);
    }

    /// The reference's explicit eviction (`SMSG_INVALIDATE_PLAYER` for a player's name): the next
    /// read re-asks. Returns whether there was anything to evict.
    pub(crate) fn evict(&mut self, key: K) -> bool {
        self.pending
            .get_mut()
            .unwrap_or_else(|p| p.into_inner())
            .remove(&key);
        let was = self.answered.remove(&key).is_some();
        if was {
            self.generation = self.generation.wrapping_add(1);
        }
        was
    }

    /// Forgets the in-flight asks, which a disconnect or a send before the writer existed may
    /// have dropped; [`register`] runs it on every world entry.
    pub(crate) fn clear_pending(&mut self) {
        self.pending
            .get_mut()
            .unwrap_or_else(|p| p.into_inner())
            .clear();
    }

    /// Drops everything, answers included (a session-scoped cache on disconnect).
    pub(crate) fn clear(&mut self) {
        self.answered.clear();
        self.clear_pending();
    }

    pub(crate) fn generation(&self) -> u64 {
        self.generation
    }

    /// Every answered key with its answer, for a persisted cache's save.
    pub(crate) fn iter(&self) -> impl Iterator<Item = (&K, &Option<V>)> {
        self.answered.iter()
    }

    pub(crate) fn len(&self) -> usize {
        self.answered.len()
    }

    #[cfg(test)]
    pub(crate) fn is_empty(&self) -> bool {
        self.answered.is_empty()
    }
}

/// A resource that owns one or more [`QueryCache`]s: what to forget on world entry.
pub(crate) trait AskOnce: Resource {
    fn clear_pending(&mut self);
}

/// Registers a cache owner to clear its in-flight asks on world entry, so a key asked before the
/// writer existed, or across a disconnect, is asked again.
pub(crate) fn register<T: AskOnce>(app: &mut App) {
    app.add_systems(
        Update,
        release_on_enter::<T>
            .in_set(benilla_world::schedule::WorldStage::Net)
            .after(crate::net::apply_net_updates),
    );
}

pub(crate) fn release_on_enter<T: AskOnce>(
    mut entered: MessageReader<EnteredWorldMessage>,
    mut cache: ResMut<T>,
) {
    if entered.read().count() == 0 {
        return;
    }
    cache.clear_pending();
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The sweep's whole job: a key asked once and never answered gets asked again, a bounded
    /// number of times, and is then reported once rather than re-asked forever.
    #[test]
    fn an_unanswered_ask_is_re_sent_until_the_tries_run_out() {
        let cache: QueryCache<u32, &str> = QueryCache::default();
        let mut asks = 0;
        cache.get_or_ask(17117, || asks += 1);

        // The first sweep only stamps it: `get_or_ask` has no clock of its own.
        assert!(cache.sweep_pending(0.0, 5.0, 3).retry.is_empty());
        assert!(
            cache.sweep_pending(4.9, 5.0, 3).retry.is_empty(),
            "and nothing is re-asked before the deadline"
        );

        assert_eq!(cache.sweep_pending(5.0, 5.0, 3).retry, vec![17117]);
        assert_eq!(cache.sweep_pending(10.0, 5.0, 3).retry, vec![17117]);
        let third = cache.sweep_pending(15.0, 5.0, 3);
        assert!(
            third.retry.is_empty(),
            "three sends in all, the original included"
        );
        assert_eq!(third.abandoned, vec![17117]);
        assert!(
            cache.sweep_pending(20.0, 5.0, 3).abandoned.is_empty(),
            "and giving up is said once, not once a sweep"
        );
    }

    /// An answer takes the key out of the sweep, whichever answer it is: a server that knows
    /// nothing about an entry says so with the high-bit miss, and that is an answer.
    #[test]
    fn an_answered_key_is_never_swept_again() {
        for answer in [Some("flute"), None] {
            let mut cache: QueryCache<u32, &str> = QueryCache::default();
            cache.get_or_ask(17117, || ());
            cache.sweep_pending(0.0, 5.0, 3);
            cache.insert(17117, answer);
            let sweep = cache.sweep_pending(60.0, 5.0, 3);
            assert!(sweep.retry.is_empty() && sweep.abandoned.is_empty());
        }
    }

    #[test]
    fn a_miss_asks_once_and_a_landing_answers_it() {
        let mut cache: QueryCache<u32, &str> = QueryCache::default();
        let mut asks = 0;
        assert!(cache.get_or_ask(7, || asks += 1).is_none());
        assert!(cache.get_or_ask(7, || asks += 1).is_none());
        assert_eq!(
            asks, 1,
            "the second lookup of a pending key does not re-send"
        );
        assert!(cache.is_pending(7));
        cache.insert(7, Some("seven"));
        assert_eq!(cache.get_or_ask(7, || asks += 1), Some(&"seven"));
        assert_eq!(asks, 1);
        assert!(!cache.is_pending(7));
        assert_eq!(cache.generation(), 1);
    }

    #[test]
    fn a_negative_is_cached_and_never_re_asked() {
        let mut cache: QueryCache<u32, &str> = QueryCache::default();
        let mut asks = 0;
        cache.get_or_ask(1, || asks += 1);
        cache.insert(1, None);
        assert!(cache.get_or_ask(1, || asks += 1).is_none());
        assert_eq!(asks, 1);
        assert!(cache.answered_unknown(1));
        assert!(!cache.answered_unknown(2));
    }

    #[test]
    fn clearing_pending_lets_a_dropped_ask_go_again_and_keeps_the_answers() {
        let mut cache: QueryCache<u32, &str> = QueryCache::default();
        let mut asks = 0;
        cache.get_or_ask(1, || asks += 1);
        cache.insert(2, Some("two"));
        cache.clear_pending();
        cache.get_or_ask(1, || asks += 1);
        assert_eq!(asks, 2);
        assert_eq!(cache.get(2), Some(&"two"));
    }

    #[test]
    fn eviction_re_asks() {
        let mut cache: QueryCache<u32, &str> = QueryCache::default();
        cache.insert(3, Some("three"));
        assert!(cache.evict(3));
        assert!(!cache.evict(3));
        let mut asks = 0;
        assert!(cache.get_or_ask(3, || asks += 1).is_none());
        assert_eq!(asks, 1);
    }
}