decdn_cache/policy/tinylfu.rs
1//! W-TinyLFU estimator + policies (filled in Stage B, Tasks 5-7).
2use super::EvictionPolicy;
3use super::FrequencyEstimator;
4use super::sketch::ShardedCountMinSketch;
5use super::{AdmissionContext, AdmissionDecision, AdmissionPolicy, Segment};
6use super::{EvictionContext, EvictionPlan};
7use crate::Hash;
8use std::collections::HashSet;
9use std::sync::Arc;
10
11/// Shared W-TinyLFU frequency evidence: one sighting per served request from
12/// the serve path, one estimate read per fill-path admission decision and per
13/// probation candidate at sweep time.
14///
15/// The counters live in a [`ShardedCountMinSketch`], so a sighting or an
16/// estimate locks the one shard its hash routes to. Serve completions and
17/// admission reads for different hashes therefore proceed in parallel, and the
18/// periodic aging pass touches a single shard rather than the whole array.
19///
20/// Aging runs on one node-wide clock across all shards. Both readers below
21/// compare estimates on a shared scale — [`ProbationAdmission`] against a
22/// node-wide threshold, [`TinyLfuEviction`] by ranking candidates from
23/// different shards against each other — so counters are only meaningful if
24/// each is caught up to that one clock when it is read. A clock per shard would
25/// age at a rate set by key skew and put the two readers on scales that cannot
26/// be compared at all.
27#[derive(Debug)]
28pub struct TinyLfuEstimator {
29 inner: ShardedCountMinSketch,
30}
31
32impl TinyLfuEstimator {
33 /// Size the sketch to roughly `sketch_bytes` of counters: one `u8` per
34 /// counter over `ROWS` rows, so `cols = bytes / 4`.
35 ///
36 /// The floor of 64 columns is a backstop for callers outside config
37 /// resolution, which rejects a `cache.tinylfu.sketch_bytes` small enough to
38 /// reach it (`decdn_common::config::MIN_TINYLFU_SKETCH_BYTES`). It keeps the
39 /// width above [`super::sketch::SHARDS`], below which a sharded sketch
40 /// degenerates to one column per shard (see [`ShardedCountMinSketch::new`]).
41 /// It is not a width the estimates read usefully at: 64 columns over
42 /// [`super::sketch::SHARDS`] shards is four columns per shard, and the
43 /// collision rate that matters is the per-shard one.
44 #[must_use]
45 pub fn new(sketch_bytes: usize) -> Self {
46 // one u8 per counter, ROWS(=4) rows: cols = bytes / 4.
47 let cols = (sketch_bytes / 4).max(64);
48 Self {
49 inner: ShardedCountMinSketch::new(cols),
50 }
51 }
52}
53
54impl FrequencyEstimator for TinyLfuEstimator {
55 fn observe(&self, hash: Hash) {
56 self.inner.increment(&hash);
57 }
58 fn estimate(&self, hash: Hash) -> u32 {
59 u32::from(self.inner.estimate(&hash))
60 }
61}
62
63/// Admits a first-ever miss into `Probation`; once the shared estimator has
64/// seen `promotion_threshold` prior sightings of a hash, admits straight to
65/// `Main`. Reads `estimate` only — never calls `observe` (see the ordering
66/// invariant in the module docs: a request must never count as evidence for
67/// its own promotion).
68///
69/// The estimate read takes one sketch shard, so a fill decision and a
70/// concurrent serve completion collide only when their hashes share a shard.
71#[derive(Debug)]
72pub struct ProbationAdmission {
73 /// The shared frequency estimator; read, never written, by this policy.
74 pub freq: Arc<dyn FrequencyEstimator>,
75 /// Prior sightings a hash needs before a miss is admitted straight to
76 /// [`Segment::Main`] instead of `Probation`.
77 pub promotion_threshold: u32,
78}
79
80impl AdmissionPolicy for ProbationAdmission {
81 fn admit(&self, ctx: &AdmissionContext) -> AdmissionDecision {
82 let segment = if self.freq.estimate(ctx.hash) < self.promotion_threshold {
83 Segment::Probation
84 } else {
85 Segment::Main
86 };
87 AdmissionDecision::Store { segment }
88 }
89}
90
91/// Ranks eviction candidates least-frequent-first, reading frequency from a
92/// shared estimator; ties break oldest-access-first (LRFU), then largest
93/// first, so blobs sharing the open-time recency seed release in the order that
94/// reaches the target in the fewest releases. Also owns the
95/// probation lifecycle at sweep time: promotes probation members whose
96/// buffered frequency has reached `promotion_threshold`, and caps probation's
97/// footprint to `probation_target_pct` of `cache_bytes` by evicting the
98/// least-frequent non-promoted probation members first.
99#[derive(Debug)]
100pub struct TinyLfuEviction {
101 freq: Arc<dyn FrequencyEstimator>,
102 promotion_threshold: u32,
103 probation_target_pct: u64,
104}
105
106impl TinyLfuEviction {
107 /// Rank against `freq`, promoting probation members at
108 /// `promotion_threshold` sightings and holding probation to
109 /// `probation_target_pct` of the cache.
110 #[must_use]
111 pub fn new(
112 freq: Arc<dyn FrequencyEstimator>,
113 promotion_threshold: u32,
114 probation_target_pct: u64,
115 ) -> Self {
116 Self {
117 freq,
118 promotion_threshold,
119 probation_target_pct,
120 }
121 }
122}
123
124impl EvictionPolicy for TinyLfuEviction {
125 fn plan(&self, ctx: &EvictionContext<'_>) -> EvictionPlan {
126 let size = |h: &Hash| ctx.sizes.get(h).copied().unwrap_or(0);
127 // 1. Promote: probation members whose buffered frequency now clears
128 // the threshold graduate to Main. They're excluded from the
129 // probation-cap eviction below (and from the global loop, since a
130 // hash can't need eviction and promotion in the same sweep).
131 let mut promote = Vec::new();
132 let mut promoted: HashSet<Hash> = HashSet::new();
133 for (h, seg) in ctx.segments {
134 if *seg == Segment::Probation && self.freq.estimate(*h) >= self.promotion_threshold {
135 promote.push((*h, Segment::Main));
136 promoted.insert(*h);
137 }
138 }
139
140 let mut evict = Vec::new();
141 let mut evicted: HashSet<Hash> = HashSet::new();
142 let mut freed = 0u64;
143
144 // 2. Cap: bound probation's footprint (excluding promoted members) to
145 // probation_target_pct of the configured cache size.
146 let probation_limit = ctx
147 .cache_bytes
148 .saturating_mul(self.probation_target_pct)
149 .saturating_div(100);
150 let mut probation_members: Vec<(Hash, u32, std::time::Instant)> = ctx
151 .candidates
152 .iter()
153 .filter(|(h, _)| {
154 ctx.segments.get(*h).copied() == Some(Segment::Probation) && !promoted.contains(*h)
155 })
156 .map(|(h, t)| (*h, self.freq.estimate(*h), *t))
157 .collect();
158 let probation_footprint: u64 = probation_members
159 .iter()
160 .map(|(h, _, _)| ctx.sizes.get(h).copied().unwrap_or(0))
161 .sum();
162 if probation_footprint > probation_limit {
163 // Least frequent first; tie-break oldest access first, then largest.
164 probation_members.sort_by(|a, b| {
165 a.1.cmp(&b.1)
166 .then(a.2.cmp(&b.2))
167 .then_with(|| size(&b.0).cmp(&size(&a.0)))
168 });
169 let mut remaining = probation_footprint;
170 for (h, _, _) in probation_members {
171 if remaining <= probation_limit {
172 break;
173 }
174 if evict.len() as u64 >= ctx.budget {
175 break;
176 }
177 let size = ctx.sizes.get(&h).copied().unwrap_or(0);
178 remaining = remaining.saturating_sub(size);
179 freed = freed.saturating_add(size);
180 evicted.insert(h);
181 evict.push(h);
182 }
183 }
184
185 // 3. Global target: continue least-frequent-first eviction toward
186 // target_bytes, skipping anything already evicted or promoted this
187 // sweep.
188 let mut scored: Vec<(Hash, u32, std::time::Instant)> = ctx
189 .candidates
190 .iter()
191 .filter(|(h, _)| !evicted.contains(*h) && !promoted.contains(*h))
192 .map(|(h, t)| (*h, self.freq.estimate(*h), *t))
193 .collect();
194 scored.sort_by(|a, b| {
195 a.1.cmp(&b.1)
196 .then(a.2.cmp(&b.2))
197 .then_with(|| size(&b.0).cmp(&size(&a.0)))
198 });
199 for (h, _, _) in scored {
200 if ctx.total_bytes.saturating_sub(freed) <= ctx.target_bytes {
201 break;
202 }
203 if evict.len() as u64 >= ctx.budget {
204 break;
205 }
206 freed = freed.saturating_add(ctx.sizes.get(&h).copied().unwrap_or(0));
207 evict.push(h);
208 }
209
210 EvictionPlan { evict, promote }
211 }
212}
213
214#[cfg(test)]
215mod tests;