1#![forbid(unsafe_code)]
21
22use std::collections::BTreeMap;
23use std::path::Path;
24use std::sync::Mutex;
25use std::time::{Duration, SystemTime, UNIX_EPOCH};
26
27use rusqlite::{Connection, OptionalExtension, params};
28use zerodds_durability_store::{
29 Contract, Cursor, DurabilitySample, DurabilityStore, Page, Result, Selector, StoreError,
30 StoreStats,
31};
32use zerodds_qos::policies::history::HistoryKind;
33
34const DEFAULT_PAGE: usize = 1024;
35
36const SCHEMA: &str = "
37CREATE TABLE IF NOT EXISTS samples (
38 topic TEXT NOT NULL,
39 instance BLOB NOT NULL,
40 sequence INTEGER NOT NULL,
41 created_nanos INTEGER NOT NULL,
42 payload BLOB NOT NULL,
43 representation INTEGER NOT NULL DEFAULT 1,
44 big_endian INTEGER NOT NULL DEFAULT 0,
45 PRIMARY KEY (topic, instance, sequence)
46) WITHOUT ROWID;
47CREATE INDEX IF NOT EXISTS idx_samples_topic ON samples(topic, instance, sequence);
48CREATE TABLE IF NOT EXISTS unregistered (
49 topic TEXT NOT NULL,
50 instance BLOB NOT NULL,
51 at_nanos INTEGER NOT NULL,
52 PRIMARY KEY (topic, instance)
53);
54";
55
56pub struct SqliteStore {
58 conn: Mutex<Connection>,
59 contracts: Mutex<BTreeMap<String, Contract>>,
60 default_contract: Contract,
61}
62
63fn backend<E: core::fmt::Display>(ctx: &str, e: E) -> StoreError {
64 StoreError::Backend(format!("sqlite store: {ctx}: {e}"))
65}
66
67fn nanos_of(t: SystemTime) -> i64 {
68 t.duration_since(UNIX_EPOCH)
69 .map(|d| i64::try_from(d.as_nanos()).unwrap_or(i64::MAX))
70 .unwrap_or(0)
71}
72
73fn time_of(nanos: i64) -> SystemTime {
74 UNIX_EPOCH + Duration::from_nanos(nanos.max(0) as u64)
75}
76
77impl SqliteStore {
78 pub fn open<P: AsRef<Path>>(path: P, default_contract: Contract) -> Result<Self> {
83 let conn = Connection::open(path.as_ref()).map_err(|e| backend("open", e))?;
84 Self::init(conn, default_contract)
85 }
86
87 pub fn open_in_memory(default_contract: Contract) -> Result<Self> {
92 let conn = Connection::open_in_memory().map_err(|e| backend("open mem", e))?;
93 Self::init(conn, default_contract)
94 }
95
96 fn init(conn: Connection, default_contract: Contract) -> Result<Self> {
97 conn.pragma_update(None, "journal_mode", "WAL")
98 .map_err(|e| backend("pragma wal", e))?;
99 conn.pragma_update(None, "synchronous", "NORMAL")
100 .map_err(|e| backend("pragma synchronous", e))?;
101 conn.execute_batch(SCHEMA)
102 .map_err(|e| backend("schema", e))?;
103 Ok(Self {
104 conn: Mutex::new(conn),
105 contracts: Mutex::new(BTreeMap::new()),
106 default_contract,
107 })
108 }
109
110 fn lock_conn(&self) -> Result<std::sync::MutexGuard<'_, Connection>> {
111 self.conn
112 .lock()
113 .map_err(|_| StoreError::Poisoned("sqlite conn"))
114 }
115
116 fn contract_for(&self, topic: &str) -> Result<Contract> {
117 Ok(self
118 .contracts
119 .lock()
120 .map_err(|_| StoreError::Poisoned("sqlite contracts"))?
121 .get(topic)
122 .copied()
123 .unwrap_or(self.default_contract))
124 }
125
126 fn count(conn: &Connection, sql: &str, p: &[&dyn rusqlite::ToSql]) -> Result<i64> {
127 conn.query_row(sql, p, |r| r.get::<_, i64>(0))
128 .map_err(|e| backend("count", e))
129 }
130}
131
132impl DurabilityStore for SqliteStore {
133 fn set_contract(&self, topic: &str, contract: Contract) -> Result<()> {
134 self.contracts
135 .lock()
136 .map_err(|_| StoreError::Poisoned("sqlite contracts"))?
137 .insert(topic.to_string(), contract);
138 Ok(())
139 }
140
141 fn store(&self, sample: DurabilitySample) -> Result<()> {
142 let contract = self.contract_for(&sample.topic)?;
143 let conn = self.lock_conn()?;
144 let inst = &sample.instance_key[..];
145
146 let is_resend = Self::count(
150 &conn,
151 "SELECT COUNT(*) FROM samples WHERE topic=?1 AND instance=?2 AND sequence=?3",
152 params![sample.topic, inst, sample.sequence as i64],
153 )? > 0;
154
155 if !is_resend
157 && contract.samples_bounded()
158 && matches!(contract.history_kind, HistoryKind::KeepAll)
159 {
160 let n = Self::count(
161 &conn,
162 "SELECT COUNT(*) FROM samples WHERE topic=?1",
163 params![sample.topic],
164 )?;
165 if n >= i64::from(contract.max_samples) {
166 return Err(StoreError::OutOfResources("max_samples"));
167 }
168 }
169 if contract.instances_bounded() {
170 let exists = Self::count(
171 &conn,
172 "SELECT COUNT(*) FROM samples WHERE topic=?1 AND instance=?2",
173 params![sample.topic, inst],
174 )? > 0;
175 if !exists {
176 let insts = Self::count(
177 &conn,
178 "SELECT COUNT(DISTINCT instance) FROM samples WHERE topic=?1",
179 params![sample.topic],
180 )?;
181 if insts >= i64::from(contract.max_instances) {
182 return Err(StoreError::OutOfResources("max_instances"));
183 }
184 }
185 }
186 if !is_resend
187 && contract.per_instance_bounded()
188 && matches!(contract.history_kind, HistoryKind::KeepAll)
189 {
190 let n = Self::count(
191 &conn,
192 "SELECT COUNT(*) FROM samples WHERE topic=?1 AND instance=?2",
193 params![sample.topic, inst],
194 )?;
195 if n >= i64::from(contract.max_samples_per_instance) {
196 return Err(StoreError::OutOfResources("max_samples_per_instance"));
197 }
198 }
199
200 conn.execute(
201 "INSERT OR REPLACE INTO samples(topic,instance,sequence,created_nanos,payload,representation,big_endian) \
202 VALUES (?1,?2,?3,?4,?5,?6,?7)",
203 params![
204 sample.topic,
205 inst,
206 sample.sequence as i64,
207 nanos_of(sample.created_at),
208 sample.payload,
209 i64::from(sample.representation),
210 i64::from(sample.big_endian)
211 ],
212 )
213 .map_err(|e| backend("insert", e))?;
214
215 if matches!(contract.history_kind, HistoryKind::KeepLast) {
217 conn.execute(
218 "DELETE FROM samples WHERE topic=?1 AND instance=?2 AND sequence NOT IN \
219 (SELECT sequence FROM samples WHERE topic=?1 AND instance=?2 \
220 ORDER BY sequence DESC LIMIT ?3)",
221 params![sample.topic, inst, contract.effective_depth() as i64],
222 )
223 .map_err(|e| backend("keep_last trim", e))?;
224 }
225 Ok(())
226 }
227
228 fn query(&self, topic: &str, selector: &Selector) -> Result<Page> {
229 let conn = self.lock_conn()?;
230 let limit = selector.limit.unwrap_or(DEFAULT_PAGE);
231 let mut sql = String::from(
233 "SELECT instance,sequence,created_nanos,payload,representation,big_endian \
234 FROM samples WHERE topic=?1",
235 );
236 let mut binds: Vec<Box<dyn rusqlite::ToSql>> = vec![Box::new(topic.to_string())];
237 if let Some(k) = selector.instance_key {
238 sql.push_str(" AND instance=?");
239 sql.push_str(&(binds.len() + 1).to_string());
240 binds.push(Box::new(k.to_vec()));
241 }
242 if let Some(lo) = selector.seq_from {
243 sql.push_str(&format!(" AND sequence>=?{}", binds.len() + 1));
244 binds.push(Box::new(lo as i64));
245 }
246 if let Some(hi) = selector.seq_to {
247 sql.push_str(&format!(" AND sequence<=?{}", binds.len() + 1));
248 binds.push(Box::new(hi as i64));
249 }
250 if let Some(t0) = selector.time_from {
251 sql.push_str(&format!(" AND created_nanos>=?{}", binds.len() + 1));
252 binds.push(Box::new(nanos_of(t0)));
253 }
254 if let Some(t1) = selector.time_to {
255 sql.push_str(&format!(" AND created_nanos<=?{}", binds.len() + 1));
256 binds.push(Box::new(nanos_of(t1)));
257 }
258 if let Some((ck, cs)) = selector.after {
259 let i1 = binds.len() + 1;
261 let i2 = binds.len() + 2;
262 let i3 = binds.len() + 3;
263 sql.push_str(&format!(
264 " AND (instance>?{i1} OR (instance=?{i2} AND sequence>?{i3}))"
265 ));
266 binds.push(Box::new(ck.to_vec()));
267 binds.push(Box::new(ck.to_vec()));
268 binds.push(Box::new(cs as i64));
269 }
270 sql.push_str(&format!(" ORDER BY instance,sequence LIMIT {}", limit + 1));
271
272 let mut stmt = conn.prepare(&sql).map_err(|e| backend("prepare", e))?;
273 let param_refs: Vec<&dyn rusqlite::ToSql> = binds.iter().map(|b| b.as_ref()).collect();
274 let topic_owned = topic.to_string();
275 let rows = stmt
276 .query_map(param_refs.as_slice(), |r| {
277 let inst: Vec<u8> = r.get(0)?;
278 let mut key = [0u8; 16];
279 if inst.len() == 16 {
280 key.copy_from_slice(&inst);
281 }
282 Ok(DurabilitySample {
283 topic: topic_owned.clone(),
284 instance_key: key,
285 sequence: r.get::<_, i64>(1)? as u64,
286 created_at: time_of(r.get::<_, i64>(2)?),
287 payload: r.get(3)?,
288 representation: r.get::<_, i64>(4)? as u8,
289 big_endian: r.get::<_, i64>(5)? != 0,
290 })
291 })
292 .map_err(|e| backend("query_map", e))?;
293 let mut samples = Vec::new();
294 for row in rows {
295 samples.push(row.map_err(|e| backend("row", e))?);
296 }
297 let exhausted = samples.len() <= limit;
298 samples.truncate(limit);
299 let next: Option<Cursor> = if exhausted {
300 None
301 } else {
302 samples.last().map(|s| (s.instance_key, s.sequence))
303 };
304 Ok(Page { samples, next })
305 }
306
307 fn unregister(&self, topic: &str, instance_key: &[u8; 16], now: SystemTime) -> Result<()> {
308 let conn = self.lock_conn()?;
309 conn.execute(
310 "INSERT OR REPLACE INTO unregistered(topic,instance,at_nanos) VALUES (?1,?2,?3)",
311 params![topic, &instance_key[..], nanos_of(now)],
312 )
313 .map_err(|e| backend("unregister", e))?;
314 Ok(())
315 }
316
317 fn cleanup(&self, now: SystemTime) -> Result<usize> {
318 let conn = self.lock_conn()?;
319 let mut stmt = conn
321 .prepare("SELECT topic,instance,at_nanos FROM unregistered")
322 .map_err(|e| backend("prepare cleanup", e))?;
323 let due: Vec<(String, Vec<u8>)> = stmt
324 .query_map([], |r| {
325 Ok((
326 r.get::<_, String>(0)?,
327 r.get::<_, Vec<u8>>(1)?,
328 r.get::<_, i64>(2)?,
329 ))
330 })
331 .map_err(|e| backend("cleanup scan", e))?
332 .filter_map(|row| row.ok())
333 .filter_map(|(topic, inst, at)| {
334 let delay = self.contract_for(&topic).ok()?.cleanup_delay;
335 let deadline = time_of(at).checked_add(delay)?;
336 if now >= deadline {
337 Some((topic, inst))
338 } else {
339 None
340 }
341 })
342 .collect();
343 drop(stmt);
344 let mut removed = 0usize;
345 for (topic, inst) in due {
346 conn.execute(
347 "DELETE FROM samples WHERE topic=?1 AND instance=?2",
348 params![topic, inst],
349 )
350 .map_err(|e| backend("cleanup delete samples", e))?;
351 conn.execute(
352 "DELETE FROM unregistered WHERE topic=?1 AND instance=?2",
353 params![topic, inst],
354 )
355 .map_err(|e| backend("cleanup delete marker", e))?;
356 removed += 1;
357 }
358 Ok(removed)
359 }
360
361 fn stats(&self, topic: &str) -> Result<StoreStats> {
362 let conn = self.lock_conn()?;
363 let samples = Self::count(
364 &conn,
365 "SELECT COUNT(*) FROM samples WHERE topic=?1",
366 params![topic],
367 )? as usize;
368 let instances = Self::count(
369 &conn,
370 "SELECT COUNT(DISTINCT instance) FROM samples WHERE topic=?1",
371 params![topic],
372 )? as usize;
373 let bytes = conn
374 .query_row(
375 "SELECT COALESCE(SUM(LENGTH(payload)),0) FROM samples WHERE topic=?1",
376 params![topic],
377 |r| r.get::<_, i64>(0),
378 )
379 .optional()
380 .map_err(|e| backend("stats bytes", e))?
381 .unwrap_or(0) as u64;
382 Ok(StoreStats {
383 samples,
384 instances,
385 bytes,
386 })
387 }
388}