1use std::num::NonZeroUsize;
34use std::time::{Duration, SystemTime, UNIX_EPOCH};
35
36use crate::cache::LruTtlStore;
37use crate::pow_profile::profile::{PowCapabilityProfile, PowCapabilitySample};
38use crate::types::TargetClass;
39use crate::vendor_classifier::VendorId;
40
41pub const DEFAULT_POW_TTL: Duration = Duration::from_hours(1);
49
50#[allow(clippy::unwrap_used)]
53pub const DEFAULT_POW_CAPACITY: NonZeroUsize = match NonZeroUsize::new(128) {
54 Some(value) => value,
55 None => NonZeroUsize::MIN,
56};
57
58const ZERO_FALLBACK_UNIX_SECS: u64 = 0;
63
64#[must_use]
83pub fn pow_profile_key(domain: &str, target_class: TargetClass, vendor: VendorId) -> String {
84 format!(
85 "charon:pow:{}:{}:{}",
86 domain.to_ascii_lowercase(),
87 target_class_label(target_class),
88 vendor.label()
89 )
90}
91
92const fn target_class_label(c: TargetClass) -> &'static str {
93 match c {
94 TargetClass::Api => "api",
95 TargetClass::ContentSite => "content_site",
96 TargetClass::HighSecurity => "high_security",
97 TargetClass::Unknown => "unknown",
98 }
99}
100
101pub struct PowCapabilityStore {
131 store: LruTtlStore<PowCapabilityProfile>,
132}
133
134impl std::fmt::Debug for PowCapabilityStore {
135 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136 f.debug_struct("PowCapabilityStore")
137 .field("ttl", &self.store.ttl())
138 .field("len", &self.store.len())
139 .finish()
140 }
141}
142
143impl PowCapabilityStore {
144 #[must_use]
146 pub fn new(capacity: NonZeroUsize, ttl: Duration) -> Self {
147 Self {
148 store: LruTtlStore::new(capacity, ttl),
149 }
150 }
151
152 #[must_use]
155 pub fn with_default_ttl(capacity: NonZeroUsize) -> Self {
156 Self::new(capacity, DEFAULT_POW_TTL)
157 }
158
159 #[must_use]
162 pub fn with_defaults() -> Self {
163 Self::new(DEFAULT_POW_CAPACITY, DEFAULT_POW_TTL)
164 }
165
166 #[must_use]
168 pub const fn ttl(&self) -> Duration {
169 self.store.ttl()
170 }
171
172 pub fn record_sample(
183 &self,
184 domain: &str,
185 target_class: TargetClass,
186 vendor: VendorId,
187 sample: &PowCapabilitySample,
188 ) {
189 let key = pow_profile_key(domain, target_class, vendor);
190 self.store.mutate(key, |existing| {
194 let mut profile =
195 existing.unwrap_or_else(|| PowCapabilityProfile::new(domain, target_class, vendor));
196 profile.merge(sample);
197 profile.recorded_at_unix_secs = current_unix_secs();
203 profile
204 });
205 }
206
207 #[must_use]
211 pub fn lookup(
212 &self,
213 domain: &str,
214 target_class: TargetClass,
215 vendor: VendorId,
216 ) -> Option<PowCapabilityProfile> {
217 self.store
218 .get(&pow_profile_key(domain, target_class, vendor))
219 }
220
221 #[must_use]
223 pub fn len(&self) -> usize {
224 self.store.len()
225 }
226
227 #[must_use]
229 pub fn is_empty(&self) -> bool {
230 self.store.is_empty()
231 }
232
233 pub fn clear(&self) {
235 self.store.clear();
236 }
237
238 pub fn invalidate(&self, domain: &str, target_class: TargetClass, vendor: VendorId) {
241 self.store
242 .invalidate(&pow_profile_key(domain, target_class, vendor));
243 }
244}
245
246fn current_unix_secs() -> u64 {
247 SystemTime::now()
248 .duration_since(UNIX_EPOCH)
249 .map_or(ZERO_FALLBACK_UNIX_SECS, |duration| duration.as_secs())
250}
251
252#[cfg(test)]
253#[allow(
254 clippy::unwrap_used,
255 clippy::expect_used,
256 clippy::panic,
257 clippy::indexing_slicing
258)]
259mod tests {
260 use super::*;
261 use std::thread;
262
263 #[test]
264 fn record_sample_creates_new_profile_on_first_call() {
265 let store = PowCapabilityStore::new(NonZeroUsize::new(4).unwrap(), DEFAULT_POW_TTL);
266 store.record_sample(
267 "example.com",
268 TargetClass::ContentSite,
269 VendorId::Cloudflare,
270 &PowCapabilitySample::solved(1_000, 0),
271 );
272 let profile = store
273 .lookup(
274 "example.com",
275 TargetClass::ContentSite,
276 VendorId::Cloudflare,
277 )
278 .expect("profile");
279 assert_eq!(profile.domain, "example.com");
280 assert_eq!(profile.solved_count, 1);
281 assert_eq!(profile.failed_count, 0);
282 assert_eq!(profile.vendor_family, VendorId::Cloudflare);
283 }
284
285 #[test]
286 fn record_sample_merges_into_existing_profile() {
287 let store = PowCapabilityStore::new(NonZeroUsize::new(4).unwrap(), DEFAULT_POW_TTL);
288 let key = (
289 "example.com",
290 TargetClass::ContentSite,
291 VendorId::Cloudflare,
292 );
293 store.record_sample(key.0, key.1, key.2, &PowCapabilitySample::solved(1_000, 0));
294 store.record_sample(key.0, key.1, key.2, &PowCapabilitySample::solved(1_500, 1));
295 store.record_sample(
296 key.0,
297 key.1,
298 key.2,
299 &PowCapabilitySample::failed(
300 2_000,
301 1,
302 crate::pow_profile::profile::PowFailureMode::Timeout,
303 ),
304 );
305 let profile = store.lookup(key.0, key.1, key.2).expect("profile");
306 assert_eq!(profile.solved_count, 2);
307 assert_eq!(profile.failed_count, 1);
308 assert_eq!(profile.retry_count, 2);
309 assert_eq!(
310 profile
311 .failure_modes
312 .get(&crate::pow_profile::profile::PowFailureMode::Timeout),
313 Some(&1)
314 );
315 }
316
317 #[test]
318 fn distinct_keys_keep_distinct_profiles() {
319 let store = PowCapabilityStore::new(NonZeroUsize::new(8).unwrap(), DEFAULT_POW_TTL);
320 store.record_sample(
321 "example.com",
322 TargetClass::ContentSite,
323 VendorId::Cloudflare,
324 &PowCapabilitySample::solved(1_000, 0),
325 );
326 store.record_sample(
327 "example.com",
328 TargetClass::Api,
329 VendorId::Cloudflare,
330 &PowCapabilitySample::solved(2_000, 0),
331 );
332 store.record_sample(
333 "example.com",
334 TargetClass::ContentSite,
335 VendorId::Akamai,
336 &PowCapabilitySample::solved(3_000, 0),
337 );
338 let cs_cf = store
339 .lookup(
340 "example.com",
341 TargetClass::ContentSite,
342 VendorId::Cloudflare,
343 )
344 .unwrap();
345 let api_cf = store
346 .lookup("example.com", TargetClass::Api, VendorId::Cloudflare)
347 .unwrap();
348 let cs_ak = store
349 .lookup("example.com", TargetClass::ContentSite, VendorId::Akamai)
350 .unwrap();
351 assert_eq!(cs_cf.solve_latency_ms_p50, Some(1_000));
352 assert_eq!(api_cf.solve_latency_ms_p50, Some(2_000));
353 assert_eq!(cs_ak.solve_latency_ms_p50, Some(3_000));
354 }
355
356 #[test]
357 fn domain_is_normalised_to_lower_case() {
358 let store = PowCapabilityStore::new(NonZeroUsize::new(4).unwrap(), DEFAULT_POW_TTL);
359 store.record_sample(
360 "Example.COM",
361 TargetClass::Api,
362 VendorId::Cloudflare,
363 &PowCapabilitySample::solved(1_000, 0),
364 );
365 let profile = store
366 .lookup("EXAMPLE.com", TargetClass::Api, VendorId::Cloudflare)
367 .expect("profile");
368 assert_eq!(profile.domain, "example.com");
369 }
370
371 #[test]
372 fn entries_decay_after_ttl() {
373 let store =
374 PowCapabilityStore::new(NonZeroUsize::new(4).unwrap(), Duration::from_millis(1));
375 store.record_sample(
376 "example.com",
377 TargetClass::Api,
378 VendorId::Cloudflare,
379 &PowCapabilitySample::solved(1_000, 0),
380 );
381 thread::sleep(Duration::from_millis(5));
382 assert!(
383 store
384 .lookup("example.com", TargetClass::Api, VendorId::Cloudflare)
385 .is_none()
386 );
387 }
388
389 #[test]
390 fn clear_drops_everything() {
391 let store = PowCapabilityStore::new(NonZeroUsize::new(4).unwrap(), DEFAULT_POW_TTL);
392 store.record_sample(
393 "a.example",
394 TargetClass::Api,
395 VendorId::Cloudflare,
396 &PowCapabilitySample::solved(1_000, 0),
397 );
398 store.record_sample(
399 "b.example",
400 TargetClass::Api,
401 VendorId::Cloudflare,
402 &PowCapabilitySample::solved(1_000, 0),
403 );
404 assert_eq!(store.len(), 2);
405 store.clear();
406 assert!(store.is_empty());
407 }
408
409 #[test]
410 fn invalidate_drops_single_key() {
411 let store = PowCapabilityStore::new(NonZeroUsize::new(4).unwrap(), DEFAULT_POW_TTL);
412 store.record_sample(
413 "a.example",
414 TargetClass::Api,
415 VendorId::Cloudflare,
416 &PowCapabilitySample::solved(1_000, 0),
417 );
418 store.record_sample(
419 "b.example",
420 TargetClass::Api,
421 VendorId::Cloudflare,
422 &PowCapabilitySample::solved(1_000, 0),
423 );
424 store.invalidate("a.example", TargetClass::Api, VendorId::Cloudflare);
425 assert!(
426 store
427 .lookup("a.example", TargetClass::Api, VendorId::Cloudflare)
428 .is_none()
429 );
430 assert!(
431 store
432 .lookup("b.example", TargetClass::Api, VendorId::Cloudflare)
433 .is_some()
434 );
435 }
436
437 #[test]
438 fn lru_capacity_is_respected() {
439 let store = PowCapabilityStore::new(NonZeroUsize::new(2).unwrap(), DEFAULT_POW_TTL);
440 store.record_sample(
441 "a.example",
442 TargetClass::Api,
443 VendorId::Cloudflare,
444 &PowCapabilitySample::solved(1_000, 0),
445 );
446 store.record_sample(
447 "b.example",
448 TargetClass::Api,
449 VendorId::Cloudflare,
450 &PowCapabilitySample::solved(1_000, 0),
451 );
452 store.record_sample(
453 "c.example",
454 TargetClass::Api,
455 VendorId::Cloudflare,
456 &PowCapabilitySample::solved(1_000, 0),
457 );
458 assert!(store.len() <= 2);
459 }
460
461 #[test]
462 fn key_namespace_is_pow_prefixed() {
463 let key = pow_profile_key("Example.COM", TargetClass::Api, VendorId::Akamai);
464 assert_eq!(key, "charon:pow:example.com:api:akamai");
465 }
466
467 #[test]
468 fn default_ttl_matches_default_sample_window() {
469 assert_eq!(
470 DEFAULT_POW_TTL.as_secs(),
471 crate::pow_profile::DEFAULT_SAMPLE_WINDOW_SECS
472 );
473 }
474}