evergreen/osrf/
cache.rs

1use crate::osrf::sclient::HostSettings;
2use crate::EgResult;
3use crate::EgValue;
4use memcache;
5use std::cell::RefCell;
6use std::collections::HashMap;
7use std::fmt;
8use std::sync::OnceLock;
9
10static GLOBAL_MEMCACHE_CLIENT: OnceLock<memcache::Client> = OnceLock::new();
11
12thread_local! {
13    static CACHE_CONNECTIONS: RefCell<HashMap<String, CacheConnection>> = RefCell::new(HashMap::new());
14}
15
16const DEFAULT_MAX_CACHE_TIME: u32 = 86400;
17const DEFAULT_MAX_CACHE_SIZE: u32 = 10_000_000; // ~10M
18const GLOBAL_CACHE_NAME: &str = "global";
19const ANON_CACHE_NAME: &str = "anon";
20
21/*
22<cache>
23  <global>
24    <servers>
25      <server>127.0.0.1:11211</server>
26    </servers>
27    <max_cache_time>86400</max_cache_time>
28  </global>
29  <anon>
30    <servers>
31      <server>127.0.0.1:11211</server>
32    </servers>
33    <max_cache_time>1800</max_cache_time>
34    <max_cache_size>102400</max_cache_size>
35  </anon>
36</cache>
37*/
38
39pub struct CacheConnection {
40    name: String,
41    memcache: memcache::Client,
42    max_cache_time: u32,
43    max_cache_size: u32,
44}
45
46impl fmt::Display for CacheConnection {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        write!(f, "Cache name={}", self.name)
49    }
50}
51
52impl CacheConnection {
53    /// Store a value in the cache with the provided timeout.
54    ///
55    /// If the timeout is 0, the default timeout for the connection type is used.
56    fn set(&self, key: &str, value: EgValue, mut timeout: u32) -> EgResult<()> {
57        let value = value.into_json_value().dump();
58        let byte_count = value.len();
59
60        log::debug!("{self} caching {byte_count} bytes at key={key}");
61
62        if byte_count > self.max_cache_size as usize {
63            return Err(format!(
64                "{self} key={key} exceeds the max size of {}",
65                self.max_cache_size
66            )
67            .into());
68        }
69
70        if timeout == 0 {
71            timeout = self.max_cache_time;
72        }
73
74        self.memcache
75            .set(key, &value, timeout)
76            .map_err(|e| format!("{self} set key={key} failed: {e}").into())
77    }
78
79    fn get(&self, key: &str) -> EgResult<Option<EgValue>> {
80        let result: Option<String> = match self.memcache.get(key) {
81            Ok(r) => r,
82            Err(e) => return Err(format!("{self} get key={key} failed: {e}").into()),
83        };
84
85        if let Some(value) = result {
86            let obj = json::parse(&value)
87                .map_err(|e| format!("Cached JSON parse failure on key {key}: {e} [{value}]"))?;
88
89            let v = EgValue::try_from(obj)?;
90
91            return Ok(Some(v));
92        }
93
94        Ok(None)
95    }
96
97    fn del(&self, key: &str) -> EgResult<()> {
98        self.memcache
99            .delete(key)
100            .map(|_| ())
101            .map_err(|e| format!("{self} del key={key} failed: {e}").into())
102    }
103}
104
105pub struct Cache;
106
107impl Cache {
108    /// Returns OK if the specified cache type has been initialized, Err otherwise.
109    fn verify_cache(cache_name: &str) -> EgResult<()> {
110        let mut has = false;
111        CACHE_CONNECTIONS.with(|cc| has = cc.borrow().contains_key(cache_name));
112        if has {
113            Ok(())
114        } else {
115            Err(format!("No such cache initialized: {cache_name}").into())
116        }
117    }
118
119    pub fn init_cache(cache_name: &str) -> EgResult<()> {
120        if Cache::verify_cache(cache_name).is_ok() {
121            log::warn!("Cache {cache_name} is already connected; ignoring");
122            return Ok(());
123        }
124
125        let conf_key = format!("cache/{}", cache_name);
126        let config = HostSettings::get(&conf_key)?;
127
128        let mut servers = Vec::new();
129        if let Some(server) = config["servers"]["server"].as_str() {
130            servers.push(format!("memcache://{server}"));
131        } else {
132            for server in config["servers"]["server"].members() {
133                servers.push(format!("memcache://{server}"));
134            }
135        }
136
137        let cache_time = config["max_cache_time"]
138            .as_int()
139            .map(|n| n as u32)
140            .unwrap_or(DEFAULT_MAX_CACHE_TIME);
141
142        let cache_size = config["max_cache_size"]
143            .as_int()
144            .map(|n| n as u32)
145            .unwrap_or(DEFAULT_MAX_CACHE_SIZE);
146
147        log::info!("Connecting to cache servers: {servers:?}");
148
149        let mc = if let Some(client) = GLOBAL_MEMCACHE_CLIENT.get() {
150            client.clone()
151        } else {
152            let mc = match memcache::connect(servers) {
153                Ok(mc) => mc,
154                Err(e) => {
155                    return Err(format!(
156                        "Cannot connect to memcache with config: {} : {e}",
157                        config.clone().into_json_value().dump()
158                    )
159                    .into());
160                }
161            };
162
163            GLOBAL_MEMCACHE_CLIENT.set(mc.clone()).ok();
164            mc
165        };
166
167        let cache = CacheConnection {
168            name: GLOBAL_CACHE_NAME.to_string(),
169            memcache: mc,
170            max_cache_time: cache_time,
171            max_cache_size: cache_size,
172        };
173
174        CACHE_CONNECTIONS.with(|c| c.borrow_mut().insert(GLOBAL_CACHE_NAME.to_string(), cache));
175
176        Ok(())
177    }
178
179    /// Remove a thing from the cache.
180    pub fn del_from(cache_name: &str, key: &str) -> EgResult<()> {
181        Cache::verify_cache(cache_name)?;
182
183        let mut result = Ok(());
184        CACHE_CONNECTIONS.with(|c| result = c.borrow().get(cache_name).unwrap().del(key));
185        result
186    }
187
188    /// Shortcut to remove a value from the "global" cache
189    pub fn del_global(key: &str) -> EgResult<()> {
190        Cache::del_from(GLOBAL_CACHE_NAME, key)
191    }
192
193    /// Shortcut to remove a value from the "anon" cache
194    pub fn del_anon(key: &str) -> EgResult<()> {
195        Cache::del_from(ANON_CACHE_NAME, key)
196    }
197
198    pub fn get(cache_name: &str, key: &str) -> EgResult<Option<EgValue>> {
199        Cache::verify_cache(cache_name)?;
200        let mut result = Ok(None);
201        CACHE_CONNECTIONS.with(|c| result = c.borrow().get(cache_name).unwrap().get(key));
202        result
203    }
204
205    /// Shortcut to return a value from the "global" cache
206    pub fn get_global(key: &str) -> EgResult<Option<EgValue>> {
207        Cache::get(GLOBAL_CACHE_NAME, key)
208    }
209
210    /// Shortcut to return a value from the "anon" cache
211    pub fn get_anon(key: &str) -> EgResult<Option<EgValue>> {
212        Cache::get(ANON_CACHE_NAME, key)
213    }
214
215    /// Store a value using the specified cache.
216    pub fn set(cache_name: &str, key: &str, value: EgValue, timeout: u32) -> EgResult<()> {
217        Cache::verify_cache(cache_name)?;
218
219        let mut result = Ok(());
220        CACHE_CONNECTIONS
221            .with(|c| result = c.borrow().get(cache_name).unwrap().set(key, value, timeout));
222        result
223    }
224
225    /// Shortcut for storing a value in the "global" cache with the
226    /// default timeout.
227    pub fn set_global(key: &str, value: EgValue) -> EgResult<()> {
228        Cache::set(GLOBAL_CACHE_NAME, key, value, 0)
229    }
230
231    /// Shortcut for storing a value in the "global" cache with the
232    /// provided timeout.
233    pub fn set_global_for(key: &str, value: EgValue, timeout: u32) -> EgResult<()> {
234        Cache::set(GLOBAL_CACHE_NAME, key, value, timeout)
235    }
236
237    /// Shortcut for storing a value in the "anon" cache with the
238    /// default timeout.
239    pub fn set_anon(key: &str, value: EgValue) -> EgResult<()> {
240        Cache::set(ANON_CACHE_NAME, key, value, 0)
241    }
242
243    /// Shortcut for storing a value in the "anon" cache with the
244    /// provided timeout.
245    pub fn set_anon_for(key: &str, value: EgValue, timeout: u32) -> EgResult<()> {
246        Cache::set(ANON_CACHE_NAME, key, value, timeout)
247    }
248}