Skip to main content

kevy_store/
accounting.rs

1//! Internal accounting helpers on [`Store`]: per-entry weight bookkeeping,
2//! LRU/LFU clock advance, prefetch, and the lazy-expire `live_entry` /
3//! `live_entry_mut` lookups used by every typed accessor.
4//!
5//! Split out of [`crate`] for file-size hygiene. Nothing here is part of
6//! the public surface — all methods are `pub(crate)` and called by sibling
7//! modules (string/hash/list/set/zset/evict/expire/keyspace).
8
9use kevy_hash::KevyHash;
10
11use crate::value::ENTRY_OVERHEAD;
12use crate::{Entry, SmallBytes, Store, apply_delta, evict, key_heap_bytes_for};
13
14impl Store {
15    /// Insert a fresh entry, replacing any prior. Stamps `entry.weight` from
16    /// the live value and key, then updates `used_memory` for either the
17    /// new-key (charges [`ENTRY_OVERHEAD`]) or overwrite (weight swap) case.
18    pub(crate) fn insert_entry(&mut self, key: SmallBytes, mut entry: Entry) -> Option<Entry> {
19        // New-key event capture: the owned key copy is only paid when
20        // the capture flag is on (server with `n` notifications).
21        let new_key_copy = (self.notify_capture & crate::notify::CAPTURE_NEW != 0)
22            .then(|| key.to_vec());
23        // A wholesale value replacement (type change / RESTORE)
24        // discards any per-field hash TTLs; a fresh create is a no-op.
25        self.clear_hash_key_ttls(key.as_slice());
26        let key_heap = key.heap_bytes() as u64;
27        entry.set_weight(key_heap + entry.value.weight());
28        if self.clock_on() {
29            self.tick_clock();
30            entry.set_lru_clock(self.clock_counter as u32);
31        }
32        let new_w = entry.weight();
33        let new_has_ttl = entry.expire_at_ns.is_some();
34        let prev = self.map.insert(key, entry);
35        match &prev {
36            Some(old) => {
37                // A displaced cold stub's vlog record dies with it.
38                self.tier_note_dead(key_heap, &old.value);
39                self.used_memory = self
40                    .used_memory
41                    .saturating_sub(old.weight())
42                    .saturating_add(new_w);
43            }
44            None => {
45                self.used_memory = self.used_memory.saturating_add(new_w + ENTRY_OVERHEAD);
46            }
47        }
48        let old_has_ttl = prev.as_ref().is_some_and(|o| o.expire_at_ns.is_some());
49        self.adjust_expires(i64::from(new_has_ttl) - i64::from(old_has_ttl));
50        self.update_peak();
51        if prev.is_none()
52            && let Some(k) = new_key_copy
53        {
54            self.notify_events.push((crate::notify::KeyspaceEvent::New, k));
55        }
56        prev
57    }
58
59    /// Remove a key, returning the displaced entry (`None` if absent).
60    /// Frees the entry's cached weight + [`ENTRY_OVERHEAD`]. This is the
61    /// DISCARD form: a cold stub's vlog record is credited dead. A
62    /// caller re-homing the entry intact (RENAME) uses
63    /// [`Self::take_entry_keepalive`] instead.
64    pub(crate) fn remove_entry(&mut self, key: &[u8]) -> Option<Entry> {
65        let old = self.take_entry_keepalive(key)?;
66        self.tier_note_dead(key_heap_bytes_for(key), &old.value);
67        Some(old)
68    }
69
70    /// [`Self::remove_entry`] minus the cold-record dead-credit — for
71    /// moves that keep the entry (and any [`crate::value::ColdRef`] in
72    /// it) alive under another key. Same accounting/hfttl behaviour.
73    pub(crate) fn take_entry_keepalive(&mut self, key: &[u8]) -> Option<Entry> {
74        self.clear_hash_key_ttls(key);
75        let old = self.map.remove(key)?;
76        self.used_memory = self
77            .used_memory
78            .saturating_sub(old.weight() + ENTRY_OVERHEAD);
79        if old.expire_at_ns.is_some() {
80            self.adjust_expires(-1);
81        }
82        Some(old)
83    }
84
85    /// Apply a signed weight delta to `key`'s cached `Entry::weight` AND to
86    /// the shard-wide `used_memory`. Used by in-place collection mutators
87    /// (HSET adding a field, LPUSH adding an item, …) so we account in O(1)
88    /// without re-walking the container.
89    pub(crate) fn account_delta(&mut self, key: &[u8], delta: i64) {
90        if delta == 0 {
91            return;
92        }
93        if let Some(e) = self.map.get_mut(key) {
94            e.add_to_weight(delta);
95        }
96        apply_delta(&mut self.used_memory, delta);
97        if delta > 0 {
98            self.update_peak();
99        }
100    }
101
102    /// Recompute `weight` for the entry at `key` from its current value +
103    /// key, then propagate the delta to `used_memory`. Use after a wholesale
104    /// in-place value swap (SET / APPEND / INCRBYFLOAT) where the prior
105    /// `Value`'s weight was already cached on the entry.
106    pub(crate) fn reweigh_entry(&mut self, key: &[u8]) {
107        let key_heap = key_heap_bytes_for(key);
108        let Some(e) = self.map.get_mut(key) else {
109            return;
110        };
111        let new_w = key_heap + e.value.weight();
112        let delta = new_w as i64 - e.weight() as i64;
113        e.set_weight(new_w);
114        apply_delta(&mut self.used_memory, delta);
115        if delta > 0 {
116            self.update_peak();
117        }
118    }
119
120    /// Advance the global access ordinal by one tick. Only invoked under
121    /// `maxmemory > 0` so the wrapping_add cost stays out of the unlimited
122    /// fast path.
123    #[inline]
124    pub(crate) fn tick_clock(&mut self) {
125        self.clock_counter = self.clock_counter.wrapping_add(1);
126    }
127
128    #[inline]
129    fn update_peak(&mut self) {
130        if self.used_memory > self.used_memory_peak {
131            self.used_memory_peak = self.used_memory;
132        }
133    }
134
135    /// Apply a weight delta computed in-place by a caller that already held
136    /// `&mut Entry` (overwrite-SET fast path) — same arithmetic as
137    /// [`Self::reweigh_entry`] but WITHOUT the second hash + map probe that
138    /// `reweigh_entry(key)` pays to re-find the entry it just mutated.
139    #[inline]
140    pub(crate) fn apply_weight_delta(&mut self, delta: i64) {
141        apply_delta(&mut self.used_memory, delta);
142        if delta > 0 {
143            self.update_peak();
144        }
145    }
146
147    /// Hint the CPU to fetch the bucket cache line for `key` into L1. Called
148    /// by the reactor's parse loop on command N+1 while command N is still
149    /// being dispatched — by the time N+1 actually probes the table, the
150    /// metadata line is hot. No-op when the table is empty. Cheap when not.
151    #[inline]
152    pub fn prefetch_for_key(&self, key: &[u8]) {
153        let hash = key.kevy_hash();
154        self.map.prefetch_for_hash(hash);
155    }
156
157    pub(crate) fn expired(&self, key: &[u8], now: u64) -> bool {
158        match self.map.get(key) {
159            Some(e) => e.is_expired_at(now),
160            None => false,
161        }
162    }
163
164    /// Drop `key` if expired; returns whether it is live afterwards. `now` is
165    /// monotonic ns since epoch (from [`crate::now_ns`]).
166    pub(crate) fn reap(&mut self, key: &[u8], now: u64) -> bool {
167        if self.expired(key, now) {
168            self.note_expired(key);
169            self.remove_entry(key);
170            self.expired_keys_total = self.expired_keys_total.saturating_add(1);
171            false
172        } else {
173            self.map.contains_key(key)
174        }
175    }
176
177    /// Single-lookup lazy-expiring read: the live `Entry` for `key`, or `None` if
178    /// absent or expired (expired keys are dropped here, as `reap` would).
179    ///
180    /// Two wins over the old `reap(now)`-then-`get` read path: (1) the clock is
181    /// read **only when the entry actually carries a TTL** — most keys don't, so
182    /// the common hit skips `Instant::now()` (~20–40 ns); (2) one fewer keyspace
183    /// lookup on hits (was peek-expiry + `contains_key` + `get` = 3; now peek +
184    /// `get` = 2). The two-phase shape (decide, then mutate/fetch) keeps the
185    /// borrow checker happy without an owning key clone.
186    pub(crate) fn live_entry(&mut self, key: &[u8]) -> Option<&Entry> {
187        // TTL-free fast path. Read cached clock fields
188        // ONLY when the entry actually carries a TTL — most keys don't,
189        // and a prior implementation paid two field reads + a pass
190        // through `is_expired` (which itself short-circuits on None)
191        // unconditionally. Saves ~5 ns / hot lookup across every
192        // collection / string read path.
193        let needs_check = self.map.get(key)?.expire_at_ns.is_some();
194        if needs_check {
195            let (uc, cn) = (self.cached_clock, self.cached_ns);
196            let expired = self.map.get(key).is_some_and(|e| e.is_expired(uc, cn));
197            if expired {
198                self.note_expired(key);
199                self.remove_entry(key);
200                self.expired_keys_total = self.expired_keys_total.saturating_add(1);
201                return None;
202            }
203        }
204        if self.clock_on() {
205            self.tick_clock();
206            let c = self.clock_counter as u32;
207            let policy = self.touch_policy();
208            let e = self.map.get_mut(key)?;
209            evict::touch_on_access(e, policy, c);
210            return Some(&*e);
211        }
212        self.map.get(key)
213    }
214
215    /// Mutable [`live_entry`](Self::live_entry): the live `Entry` for `key` by
216    /// `&mut`, or `None` if absent/expired (expired dropped). Same wins — clock
217    /// read only on TTL'd keys, one fewer lookup than `reap`-then-`get_mut`.
218    /// Read-modify commands (INCR/APPEND/…) get the entry once and mutate in
219    /// place, preserving any TTL on it.
220    pub(crate) fn live_entry_mut(&mut self, key: &[u8]) -> Option<&mut Entry> {
221        // See `live_entry` doc — TTL-free fast path.
222        let needs_check = self.map.get(key)?.expire_at_ns.is_some();
223        if needs_check {
224            let (uc, cn) = (self.cached_clock, self.cached_ns);
225            let expired = self.map.get(key).is_some_and(|e| e.is_expired(uc, cn));
226            if expired {
227                self.note_expired(key);
228                self.remove_entry(key);
229                self.expired_keys_total = self.expired_keys_total.saturating_add(1);
230                return None;
231            }
232        }
233        if self.clock_on() {
234            self.tick_clock();
235            let c = self.clock_counter as u32;
236            let policy = self.touch_policy();
237            let e = self.map.get_mut(key)?;
238            evict::touch_on_access(e, policy, c);
239            return Some(e);
240        }
241        self.map.get_mut(key)
242    }
243}