dd_sds/scanner/regex_rule/
regex_cache_store.rs

1use crate::SharedPool;
2use crate::scanner::regex_rule::regex_store::{RegexCacheKey, SharedRegex, gc_regex_store};
3use lazy_static::lazy_static;
4use regex_automata::meta::Regex as MetaRegex;
5use slotmap::SecondaryMap;
6use std::sync::{Arc, RwLock};
7extern crate num_cpus;
8
9fn new_regex_cache_pool() -> Arc<SharedPool<Box<RegexCaches>>> {
10    Arc::new(SharedPool::new(
11        Box::new(|| Box::new(RegexCaches::new())),
12        num_cpus::get(),
13    ))
14}
15
16lazy_static! {
17    // `RwLock` lets `reset_regex_caches` swap in a fresh pool, dropping every thread's caches.
18    static ref REGEX_CACHE_STORE: RwLock<Arc<SharedPool<Box<RegexCaches>>>> =
19        RwLock::new(new_regex_cache_pool());
20}
21
22pub fn access_regex_caches<T>(func: impl FnOnce(&mut RegexCaches) -> T) -> T {
23    // Clone the `Arc` under a short read lock so scanning holds no lock; a concurrent
24    // `reset_regex_caches` keeps this in-flight pool alive via the clone.
25    let pool = REGEX_CACHE_STORE.read().unwrap().clone();
26    let mut caches = pool.get();
27    func(caches.get_ref())
28}
29
30/// Swaps in a fresh pool, dropping every thread's cached scratch. In-flight scans keep the
31/// old pool alive via their `Arc` clone; caches are recreated on the next scan.
32pub fn reset_regex_caches() {
33    *REGEX_CACHE_STORE.write().unwrap() = new_regex_cache_pool();
34}
35
36/// Clears both scanning caches: per-thread scratch (`reset_regex_caches`) and unreferenced
37/// compiled regexes (`gc_regex_store`). Call when idle to drop dd-sds to a near-zero footprint.
38pub fn clear_all_caches() {
39    reset_regex_caches();
40    gc_regex_store();
41}
42
43pub struct RegexCaches {
44    map: SecondaryMap<RegexCacheKey, RegexCacheValue>,
45}
46
47pub struct RegexCacheValue {
48    pub cache: regex_automata::meta::Cache,
49    pub captures: regex_automata::util::captures::Captures,
50}
51
52impl RegexCaches {
53    pub fn new() -> Self {
54        Self {
55            map: SecondaryMap::new(),
56        }
57    }
58
59    pub fn get(&mut self, shared_regex: &SharedRegex) -> &mut RegexCacheValue {
60        self.raw_get(shared_regex.cache_key, &shared_regex.regex)
61    }
62
63    pub(super) fn raw_get(
64        &mut self,
65        key: RegexCacheKey,
66        regex: &MetaRegex,
67    ) -> &mut RegexCacheValue {
68        self.map
69            .entry(key)
70            .unwrap()
71            .or_insert_with(|| RegexCacheValue {
72                cache: regex.create_cache(),
73                captures: regex.create_captures(),
74            })
75    }
76}
77
78#[cfg(test)]
79mod test {
80    use super::*;
81    use crate::scanner::regex_rule::regex_store::get_memoized_regex;
82    use regex_automata::meta::Regex;
83
84    #[test]
85    fn reset_swaps_pool_and_still_serves_caches() {
86        let shared = get_memoized_regex("unique-reset-pattern", Regex::new).unwrap();
87        access_regex_caches(|caches| {
88            let _ = caches.get(&shared);
89        });
90
91        // Hold the old pool so its allocation can't be reused (flaky ptr compare), then
92        // confirm reset installs a different pool.
93        let old_pool = REGEX_CACHE_STORE.read().unwrap().clone();
94        reset_regex_caches();
95        let new_pool = REGEX_CACHE_STORE.read().unwrap().clone();
96        assert!(!Arc::ptr_eq(&old_pool, &new_pool));
97
98        // Fresh pool recreates caches on next access.
99        access_regex_caches(|caches| {
100            let _ = caches.get(&shared);
101        });
102    }
103}