1use std::collections::HashMap;
19use std::sync::Mutex;
20
21use pagesift_stat::beta::{BetaCounts, Prior};
22use pagesift_stat::rng::Rng;
23use pagesift_stat::{Reader, Stat, Writer};
24
25const TABLE_VERSION: u32 = 1;
27
28#[derive(Debug, Clone, Copy, PartialEq)]
30pub struct RenderWhy {
31 pub learned: bool,
33 pub p: f64,
34 pub value: f64,
37 pub charge: f64,
38}
39
40pub type PlanObserver = std::sync::Arc<dyn Fn(u64, bool) + Send + Sync>;
42
43pub struct RenderPlanner {
44 shards: Box<[Mutex<HashMap<u64, BetaCounts>>]>,
45 observer: std::sync::OnceLock<PlanObserver>,
46}
47
48impl Default for RenderPlanner {
49 fn default() -> Self {
50 let n = std::thread::available_parallelism()
52 .map_or(1, |n| n.get())
53 .next_power_of_two();
54 Self {
55 shards: (0..n).map(|_| Mutex::new(HashMap::new())).collect(),
56 observer: std::sync::OnceLock::new(),
57 }
58 }
59}
60
61thread_local! {
62 static RNG: std::cell::RefCell<Rng> = std::cell::RefCell::new(Rng::from_entropy());
63}
64
65impl RenderPlanner {
66 fn with<R>(&self, family: &str, f: impl FnOnce(&mut HashMap<u64, BetaCounts>, u64) -> R) -> R {
67 let key = family_key(family);
68 let shard = &self.shards[(key as usize) & (self.shards.len() - 1)];
69 f(&mut shard.lock().unwrap_or_else(|e| e.into_inner()), key)
70 }
71
72 pub fn decide(
80 &self,
81 family: &str,
82 heuristic: bool,
83 render_secs: f64,
84 remaining_secs: f64,
85 browser_price: f64,
86 ) -> (bool, RenderWhy) {
87 let counts = self.with(family, |m, k| m.get(&k).copied());
88 let value = 1.0;
89 let charge = render_secs * (browser_price + 1.0 / remaining_secs.max(f64::MIN_POSITIVE));
90 match counts.filter(|c| c.trials() > 0.0) {
91 None => (
92 heuristic,
93 RenderWhy {
94 learned: false,
95 p: f64::from(u8::from(heuristic)),
96 value,
97 charge,
98 },
99 ),
100 Some(c) => {
101 let p = RNG.with(|r| c.posterior(Prior::UNIFORM).sample(&mut r.borrow_mut()));
102 (
103 p * value > charge,
104 RenderWhy {
105 learned: true,
106 p,
107 value,
108 charge,
109 },
110 )
111 }
112 }
113 }
114
115 pub fn record(&self, family: &str, gain: bool) {
117 let key = self.with(family, |m, k| {
118 m.entry(k).or_default().observe(gain);
119 k
120 });
121 if let Some(o) = self.observer.get() {
122 o(key, gain);
123 }
124 }
125
126 pub fn observe_with(&self, o: PlanObserver) -> bool {
128 self.observer.set(o).is_ok()
129 }
130
131 pub fn import(&self, key: u64, counts: BetaCounts) {
134 let shard = &self.shards[(key as usize) & (self.shards.len() - 1)];
135 let mut m = shard.lock().unwrap_or_else(|e| e.into_inner());
136 let e = m.entry(key).or_default();
137 if e.trials() == 0.0 {
138 *e = counts;
139 }
140 }
141
142 pub fn mean(&self, family: &str) -> Option<f64> {
144 self.with(family, |m, k| m.get(&k).copied())
145 .filter(|c| c.trials() > 0.0)
146 .map(|c| c.posterior(Prior::UNIFORM).mean())
147 }
148
149 pub fn encode(&self) -> Vec<u8> {
151 let mut w = Writer::new();
152 w.u32(TABLE_VERSION);
153 let entries: Vec<(u64, BetaCounts)> = self
154 .shards
155 .iter()
156 .flat_map(|s| {
157 let m = s.lock().unwrap_or_else(|e| e.into_inner());
158 m.iter().map(|(k, v)| (*k, *v)).collect::<Vec<_>>()
159 })
160 .collect();
161 w.u64(entries.len() as u64);
162 for (k, v) in &entries {
163 w.u64(*k);
164 v.encode(&mut w);
165 }
166 w.into_inner()
167 }
168
169 pub fn decode_merge(&self, buf: &[u8]) {
171 let mut r = Reader::new(buf);
172 let read = |r: &mut Reader<'_>| -> Option<Vec<(u64, BetaCounts)>> {
173 if r.u32().ok()? != TABLE_VERSION {
174 return None;
175 }
176 let n = r.u64().ok()?;
177 (0..n)
178 .map(|_| Some((r.u64().ok()?, BetaCounts::decode(r).ok()?)))
179 .collect()
180 };
181 for (k, v) in read(&mut r).unwrap_or_default() {
182 let shard = &self.shards[(k as usize) & (self.shards.len() - 1)];
183 shard
184 .lock()
185 .unwrap_or_else(|e| e.into_inner())
186 .entry(k)
187 .or_default()
188 .merge(&v);
189 }
190 }
191}
192
193pub fn family_key(family: &str) -> u64 {
194 let h = blake3::hash(family.as_bytes());
195 let (head, _) = h
197 .as_bytes()
198 .split_first_chunk::<8>()
199 .expect("BLAKE3 digest is 32 bytes");
200 u64::from_le_bytes(*head)
201}
202
203pub fn render_family(url: &str, client_code: bool) -> String {
206 format!(
207 "{}#{}",
208 family(url),
209 if client_code { "js" } else { "static" }
210 )
211}
212
213pub fn fallback_family(url: &str, engine: usize) -> String {
215 format!("{}#after{engine}", family(url))
216}
217
218pub fn family(raw: &str) -> String {
220 let Ok(url) = url::Url::parse(raw) else {
221 return raw.to_string();
222 };
223 let section = url
224 .path()
225 .trim_start_matches('/')
226 .split('/')
227 .next()
228 .unwrap_or("");
229 format!("{}/{}", url.origin().ascii_serialization(), section)
230}
231
232pub fn quality(doc: &pagesift_content::Document, html_bytes: usize) -> Option<f64> {
234 let model = crate::quality::QualityModel::builtin()?;
235 let weights = crate::script::ScriptWeights::builtin();
236 Some(model.predict(&crate::quality::features(doc, html_bytes, weights)))
237}