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; const GLOBAL_CACHE_NAME: &str = "global";
19const ANON_CACHE_NAME: &str = "anon";
20
21pub 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 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 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 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 pub fn del_global(key: &str) -> EgResult<()> {
190 Cache::del_from(GLOBAL_CACHE_NAME, key)
191 }
192
193 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 pub fn get_global(key: &str) -> EgResult<Option<EgValue>> {
207 Cache::get(GLOBAL_CACHE_NAME, key)
208 }
209
210 pub fn get_anon(key: &str) -> EgResult<Option<EgValue>> {
212 Cache::get(ANON_CACHE_NAME, key)
213 }
214
215 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 pub fn set_global(key: &str, value: EgValue) -> EgResult<()> {
228 Cache::set(GLOBAL_CACHE_NAME, key, value, 0)
229 }
230
231 pub fn set_global_for(key: &str, value: EgValue, timeout: u32) -> EgResult<()> {
234 Cache::set(GLOBAL_CACHE_NAME, key, value, timeout)
235 }
236
237 pub fn set_anon(key: &str, value: EgValue) -> EgResult<()> {
240 Cache::set(ANON_CACHE_NAME, key, value, 0)
241 }
242
243 pub fn set_anon_for(key: &str, value: EgValue, timeout: u32) -> EgResult<()> {
246 Cache::set(ANON_CACHE_NAME, key, value, timeout)
247 }
248}