unifier/
store.rs

1//! In-memory hot store with dirty tracking; flush writes only changed data to disk.
2
3use std::collections::{BTreeMap, BTreeSet, VecDeque};
4use std::fs;
5use std::path::{Path, PathBuf};
6
7use uuid::Uuid;
8
9use crate::constants::{CRON, KEYS, MAILBOX};
10use crate::cron::CronSchedule;
11use crate::error::{Error, Result};
12use crate::fs_text::{read_text, write_text, write_text_atomic};
13use crate::home::UnifierHome;
14use crate::paths::{cron_dir, key_path, mailbox_dir, message_path, parse_message_id};
15use crate::postbox::Message;
16use crate::scope::resolve_under_root;
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub(crate) enum KeyState {
20    Present { value: String, dirty: bool },
21    Deleted { dirty: bool },
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25struct MsgState {
26    body: String,
27    dirty: bool,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
31enum MessageKind {
32    Mailbox,
33    Cron,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub(crate) struct EventState {
38    pub(crate) body: String,
39    pub(crate) created_at: String,
40    /// RFC3339 expiry; `None` means the event never expires.
41    pub(crate) expires_at: Option<String>,
42    pub(crate) dirty: bool,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub(crate) enum StagingValue {
47    Present(String),
48    Deleted,
49}
50
51/// Active tick: reads frozen snapshot, writes go to staging only.
52#[derive(Debug, Clone)]
53pub(crate) struct ActiveTick {
54    pub number: u64,
55    pub(crate) label: String,
56    /// Current lifecycle phase (`start`, or a name set via `tick phase`).
57    pub(crate) phase: String,
58    pub(crate) read_snapshot: BTreeMap<String, String>,
59    pub(crate) staging: BTreeMap<String, StagingValue>,
60    pub(crate) locks: BTreeSet<String>,
61}
62
63/// In-memory mirror of the postbox; disk is touched only on [`HotStore::flush`].
64#[derive(Debug, Default)]
65pub struct HotStore {
66    pub(crate) keys: BTreeMap<String, KeyState>,
67    mailboxes: BTreeMap<String, BTreeMap<Uuid, MsgState>>,
68    cron: BTreeMap<String, BTreeMap<Uuid, MsgState>>,
69    pub(crate) events: BTreeMap<Uuid, EventState>,
70    removed_messages: BTreeSet<(MessageKind, String, Uuid)>,
71    pub(crate) committed_tick: u64,
72    pub(crate) active_tick: Option<ActiveTick>,
73    pub(crate) tick_queue: VecDeque<String>,
74}
75
76impl HotStore {
77    pub fn load(home: &UnifierHome) -> Result<Self> {
78        let mut store = Self::default();
79        store.load_keys(home)?;
80        store.mailboxes = load_message_buckets(home, MAILBOX)?;
81        store.cron = load_message_buckets(home, CRON)?;
82        store.events = crate::tick::load_events(home)?;
83        store.committed_tick = crate::tick::load_committed_tick(home)?;
84        Ok(store)
85    }
86
87    pub fn is_dirty(&self) -> bool {
88        self.keys.values().any(|k| match k {
89            KeyState::Present { dirty, .. } | KeyState::Deleted { dirty } => *dirty,
90        }) || self.has_dirty_messages(&self.mailboxes)
91            || self.has_dirty_messages(&self.cron)
92            || self.events.values().any(|e| e.dirty)
93            || !self.removed_messages.is_empty()
94            || self.active_tick.is_some()
95    }
96
97    pub fn put_key(&mut self, key: &str, value: &str) -> Result<()> {
98        validate_key(key)?;
99        if let Some(tick) = &self.active_tick {
100            if tick.locks.contains(key) {
101                return Err(Error::msg(format!(
102                    "key locked during tick {}: {key}",
103                    tick.number
104                )));
105            }
106        }
107        if let Some(tick) = &mut self.active_tick {
108            tick.staging
109                .insert(key.to_string(), StagingValue::Present(value.to_string()));
110            return Ok(());
111        }
112        self.keys.insert(
113            key.to_string(),
114            KeyState::Present {
115                value: value.to_string(),
116                dirty: true,
117            },
118        );
119        Ok(())
120    }
121
122    pub fn get_key(&self, key: &str) -> Result<Option<String>> {
123        validate_key(key)?;
124        if let Some(tick) = &self.active_tick {
125            if let Some(staged) = tick.staging.get(key) {
126                return Ok(match staged {
127                    StagingValue::Present(v) => Some(v.clone()),
128                    StagingValue::Deleted => None,
129                });
130            }
131            return Ok(tick.read_snapshot.get(key).cloned());
132        }
133        Ok(match self.keys.get(key) {
134            Some(KeyState::Present { value, .. }) => Some(value.clone()),
135            Some(KeyState::Deleted { .. }) | None => None,
136        })
137    }
138
139    /// Present key names, optionally restricted to those under `prefix/` (or exact `prefix`).
140    pub fn list_keys(&self, prefix: Option<&str>) -> Result<Vec<String>> {
141        if let Some(p) = prefix {
142            if !p.is_empty() {
143                validate_key(p)?;
144            }
145        }
146        let mut out: Vec<String> = self
147            .keys
148            .iter()
149            .filter_map(|(k, state)| match state {
150                KeyState::Present { .. } => Some(k.clone()),
151                KeyState::Deleted { .. } => None,
152            })
153            .filter(|k| match prefix {
154                None | Some("") => true,
155                Some(p) => k == p || k.starts_with(&format!("{p}/")),
156            })
157            .collect();
158        out.sort();
159        Ok(out)
160    }
161
162    pub fn delete_key(&mut self, key: &str) -> Result<bool> {
163        validate_key(key)?;
164        if let Some(tick) = &self.active_tick {
165            if tick.locks.contains(key) {
166                return Err(Error::msg(format!(
167                    "key locked during tick {}: {key}",
168                    tick.number
169                )));
170            }
171        }
172        if let Some(tick) = &mut self.active_tick {
173            let existed = tick.staging.contains_key(key) || tick.read_snapshot.contains_key(key);
174            if !existed {
175                return Ok(false);
176            }
177            tick.staging.insert(key.to_string(), StagingValue::Deleted);
178            return Ok(true);
179        }
180        let existed = matches!(self.keys.get(key), Some(KeyState::Present { .. }));
181        if !existed {
182            return Ok(false);
183        }
184        self.keys
185            .insert(key.to_string(), KeyState::Deleted { dirty: true });
186        Ok(true)
187    }
188
189    pub fn send(&mut self, recipient: &str, body: &str) -> Result<Uuid> {
190        self.send_from(crate::envelope::DEFAULT_SENDER, recipient, body)
191    }
192
193    pub fn send_from(&mut self, from: &str, recipient: &str, body: &str) -> Result<Uuid> {
194        validate_segment(from, "from")?;
195        validate_segment(recipient, "recipient")?;
196        let env = crate::envelope::Envelope::new(
197            from,
198            recipient,
199            crate::envelope::Envelope::parse_payload(body),
200        );
201        let id = env.id;
202        self.mailboxes
203            .entry(recipient.to_string())
204            .or_default()
205            .insert(
206                id,
207                MsgState {
208                    body: env.to_json()?,
209                    dirty: true,
210                },
211            );
212        Ok(id)
213    }
214
215    pub fn post_cron(&mut self, schedule: &str, body: &str) -> Result<Uuid> {
216        CronSchedule::parse(schedule)?;
217        let id = Uuid::new_v4();
218        self.cron.entry(schedule.to_string()).or_default().insert(
219            id,
220            MsgState {
221                body: body.to_string(),
222                dirty: true,
223            },
224        );
225        Ok(id)
226    }
227
228    pub fn poll_mailbox(&self, recipient: &str) -> Result<Vec<Message>> {
229        validate_segment(recipient, "recipient")?;
230        Ok(self.collect_messages(
231            MessageKind::Mailbox,
232            recipient,
233            self.mailboxes.get(recipient),
234        ))
235    }
236
237    pub fn poll_cron(&self) -> Result<Vec<Message>> {
238        let mut out = Vec::new();
239        for (schedule, msgs) in &self.cron {
240            let parsed = CronSchedule::parse(schedule)?;
241            if !parsed.matches_now() {
242                continue;
243            }
244            out.extend(self.collect_messages(MessageKind::Cron, schedule, Some(msgs)));
245        }
246        out.sort_by(|a, b| a.path.cmp(&b.path));
247        Ok(out)
248    }
249
250    pub fn list_dir(&self, home: &UnifierHome, subpath: &str) -> Result<Vec<Message>> {
251        let dir = resolve_under_root(home.path(), subpath)?;
252        let rel = dir
253            .strip_prefix(home.path())
254            .map_err(|_| Error::msg("path escapes store root"))?;
255        let parts: Vec<_> = rel.iter().collect();
256        match parts.as_slice() {
257            [p1, p2] if p1.to_string_lossy() == MAILBOX => self.poll_mailbox(&p2.to_string_lossy()),
258            [p1, p2] if p1.to_string_lossy() == CRON => {
259                let schedule = p2.to_string_lossy();
260                Ok(self.collect_messages(
261                    MessageKind::Cron,
262                    &schedule,
263                    self.cron.get(schedule.as_ref()),
264                ))
265            }
266            _ => Ok(Vec::new()),
267        }
268    }
269
270    pub fn ack(&mut self, home: &UnifierHome, id_or_path: &str) -> Result<bool> {
271        if id_or_path.contains('/') || Path::new(id_or_path).is_absolute() {
272            let path = if Path::new(id_or_path).is_absolute() {
273                PathBuf::from(id_or_path)
274            } else {
275                resolve_under_root(home.path(), id_or_path)?
276            };
277            return self.ack_path(home, &path);
278        }
279        let id = Uuid::parse_str(id_or_path)
280            .map_err(|_| Error::msg(format!("invalid message id or path: {id_or_path}")))?;
281        self.ack_id(home, &id)
282    }
283
284    pub fn flush(&mut self, home: &UnifierHome) -> Result<()> {
285        if self.active_tick.is_some() {
286            return Err(Error::msg(
287                "cannot flush while a tick is active; run tick end first",
288            ));
289        }
290        home.ensure()?;
291        self.flush_keys(home)?;
292        self.flush_mailboxes(home)?;
293        self.flush_cron(home)?;
294        self.flush_events(home)?;
295        self.removed_messages.clear();
296        Ok(())
297    }
298
299    fn load_keys(&mut self, home: &UnifierHome) -> Result<()> {
300        let root = home.path().join(KEYS);
301        if !root.is_dir() {
302            return Ok(());
303        }
304        let mut prefix = Vec::new();
305        self.walk_key_files(&root, &mut prefix)?;
306        Ok(())
307    }
308
309    fn walk_key_files(&mut self, dir: &Path, prefix: &mut Vec<String>) -> Result<()> {
310        for entry in fs::read_dir(dir)? {
311            let entry = entry?;
312            let name = entry.file_name().to_string_lossy().into_owned();
313            if entry.file_type()?.is_dir() {
314                prefix.push(name);
315                self.walk_key_files(&entry.path(), prefix)?;
316                prefix.pop();
317            } else if entry.file_type()?.is_file() {
318                prefix.push(name);
319                let key = prefix.join("/");
320                prefix.pop();
321                let value = read_text(&entry.path())?;
322                self.keys.insert(
323                    key,
324                    KeyState::Present {
325                        value,
326                        dirty: false,
327                    },
328                );
329            }
330        }
331        Ok(())
332    }
333
334    fn flush_keys(&mut self, home: &UnifierHome) -> Result<()> {
335        let mut to_remove = Vec::new();
336        for (key, state) in &mut self.keys {
337            match state {
338                KeyState::Present { value, dirty: true } => {
339                    write_text_atomic(&key_path(home, key), value)?;
340                    *state = KeyState::Present {
341                        value: value.clone(),
342                        dirty: false,
343                    };
344                }
345                KeyState::Deleted { dirty: true } => {
346                    let path = key_path(home, key);
347                    if path.is_file() {
348                        fs::remove_file(&path)?;
349                    }
350                    to_remove.push(key.clone());
351                }
352                _ => {}
353            }
354        }
355        for key in to_remove {
356            self.keys.remove(&key);
357        }
358        Ok(())
359    }
360
361    fn flush_mailboxes(&mut self, home: &UnifierHome) -> Result<()> {
362        flush_message_buckets(
363            home,
364            MessageKind::Mailbox,
365            &mut self.mailboxes,
366            &mut self.removed_messages,
367            mailbox_dir,
368        )
369    }
370
371    fn flush_cron(&mut self, home: &UnifierHome) -> Result<()> {
372        flush_message_buckets(
373            home,
374            MessageKind::Cron,
375            &mut self.cron,
376            &mut self.removed_messages,
377            cron_dir,
378        )
379    }
380
381    fn ack_path(&mut self, home: &UnifierHome, path: &Path) -> Result<bool> {
382        let rel = path.strip_prefix(home.path()).ok();
383        let Some(rel) = rel else {
384            return Ok(false);
385        };
386        let parts: Vec<_> = rel
387            .iter()
388            .map(|p| p.to_string_lossy().into_owned())
389            .collect();
390        if parts.len() != 3 {
391            return Ok(false);
392        }
393        let Some(id) = parse_message_id(&parts[2]) else {
394            return Ok(false);
395        };
396        let existed = match parts[0].as_str() {
397            MAILBOX => {
398                let had = self
399                    .mailboxes
400                    .get(&parts[1])
401                    .is_some_and(|m| m.contains_key(&id))
402                    || path.is_file();
403                if had {
404                    self.removed_messages
405                        .insert((MessageKind::Mailbox, parts[1].clone(), id));
406                    if let Some(msgs) = self.mailboxes.get_mut(&parts[1]) {
407                        msgs.remove(&id);
408                    }
409                }
410                had
411            }
412            CRON => {
413                let had = self
414                    .cron
415                    .get(&parts[1])
416                    .is_some_and(|m| m.contains_key(&id))
417                    || path.is_file();
418                if had {
419                    self.removed_messages
420                        .insert((MessageKind::Cron, parts[1].clone(), id));
421                    if let Some(msgs) = self.cron.get_mut(&parts[1]) {
422                        msgs.remove(&id);
423                    }
424                }
425                had
426            }
427            _ => false,
428        };
429        Ok(existed)
430    }
431
432    fn ack_id(&mut self, home: &UnifierHome, id: &Uuid) -> Result<bool> {
433        if let Some(recipient) = self
434            .mailboxes
435            .iter()
436            .find_map(|(name, msgs)| msgs.contains_key(id).then(|| name.clone()))
437        {
438            self.removed_messages
439                .insert((MessageKind::Mailbox, recipient.clone(), *id));
440            if let Some(m) = self.mailboxes.get_mut(&recipient) {
441                m.remove(id);
442            }
443            return Ok(true);
444        }
445        if let Some(schedule) = self
446            .cron
447            .iter()
448            .find_map(|(name, msgs)| msgs.contains_key(id).then(|| name.clone()))
449        {
450            self.removed_messages
451                .insert((MessageKind::Cron, schedule.clone(), *id));
452            if let Some(m) = self.cron.get_mut(&schedule) {
453                m.remove(id);
454            }
455            return Ok(true);
456        }
457        // Message may exist on disk but not loaded into a bucket (empty dir walk).
458        let filename = format!("{}.txt", id.hyphenated());
459        for sub in [MAILBOX, CRON] {
460            let base = home.path().join(sub);
461            if !base.is_dir() {
462                continue;
463            }
464            for entry in fs::read_dir(&base)? {
465                let entry = entry?;
466                if entry.file_type()?.is_dir() {
467                    let candidate = entry.path().join(&filename);
468                    if candidate.is_file() {
469                        let bucket = entry.file_name().to_string_lossy().into_owned();
470                        let kind = if sub == MAILBOX {
471                            MessageKind::Mailbox
472                        } else {
473                            MessageKind::Cron
474                        };
475                        self.removed_messages.insert((kind, bucket, *id));
476                        return Ok(true);
477                    }
478                }
479            }
480        }
481        Ok(false)
482    }
483
484    fn collect_messages(
485        &self,
486        kind: MessageKind,
487        bucket: &str,
488        msgs: Option<&BTreeMap<Uuid, MsgState>>,
489    ) -> Vec<Message> {
490        let Some(msgs) = msgs else {
491            return Vec::new();
492        };
493        let base = match kind {
494            MessageKind::Mailbox => PathBuf::from(MAILBOX).join(bucket),
495            MessageKind::Cron => PathBuf::from(CRON).join(bucket),
496        };
497        let mut out = Vec::new();
498        for (id, msg) in msgs {
499            out.push(Message {
500                id: *id,
501                path: base.join(format!("{}.txt", id.hyphenated())),
502                body: msg.body.clone(),
503            });
504        }
505        out.sort_by_key(|m| m.id);
506        out
507    }
508
509    fn has_dirty_messages(&self, buckets: &BTreeMap<String, BTreeMap<Uuid, MsgState>>) -> bool {
510        buckets.values().any(|msgs| msgs.values().any(|m| m.dirty))
511    }
512}
513
514fn load_message_buckets(
515    home: &UnifierHome,
516    top: &str,
517) -> Result<BTreeMap<String, BTreeMap<Uuid, MsgState>>> {
518    let mut target = BTreeMap::new();
519    let base = home.path().join(top);
520    if !base.is_dir() {
521        return Ok(target);
522    }
523    for entry in fs::read_dir(&base)? {
524        let entry = entry?;
525        if !entry.file_type()?.is_dir() {
526            continue;
527        }
528        let bucket = entry.file_name().to_string_lossy().into_owned();
529        let mut msgs = BTreeMap::new();
530        for msg_entry in fs::read_dir(entry.path())? {
531            let msg_entry = msg_entry?;
532            if !msg_entry.file_type()?.is_file() {
533                continue;
534            }
535            let name = msg_entry.file_name().to_string_lossy().into_owned();
536            let Some(id) = parse_message_id(&name) else {
537                continue;
538            };
539            msgs.insert(
540                id,
541                MsgState {
542                    body: read_text(&msg_entry.path())?,
543                    dirty: false,
544                },
545            );
546        }
547        if !msgs.is_empty() {
548            target.insert(bucket, msgs);
549        }
550    }
551    Ok(target)
552}
553
554fn flush_message_buckets<F>(
555    home: &UnifierHome,
556    kind: MessageKind,
557    buckets: &mut BTreeMap<String, BTreeMap<Uuid, MsgState>>,
558    removed_messages: &mut BTreeSet<(MessageKind, String, Uuid)>,
559    dir_for: F,
560) -> Result<()>
561where
562    F: Fn(&UnifierHome, &str) -> PathBuf,
563{
564    for (bucket, msgs) in buckets.iter_mut() {
565        for (id, msg) in msgs.iter_mut() {
566            if msg.dirty {
567                write_text(&message_path(&dir_for(home, bucket), id), &msg.body)?;
568                msg.dirty = false;
569            }
570        }
571    }
572    let removed: Vec<_> = removed_messages
573        .iter()
574        .filter(|(k, _, _)| *k == kind)
575        .cloned()
576        .collect();
577    for (_, bucket, id) in removed {
578        let path = message_path(&dir_for(home, &bucket), &id);
579        if path.is_file() {
580            fs::remove_file(path)?;
581        }
582        if let Some(msgs) = buckets.get_mut(&bucket) {
583            msgs.remove(&id);
584        }
585        removed_messages.remove(&(kind, bucket.clone(), id));
586    }
587    buckets.retain(|_, msgs| !msgs.is_empty());
588    Ok(())
589}
590
591pub(crate) fn validate_key(key: &str) -> Result<()> {
592    if key.is_empty() {
593        return Err(Error::msg("key must not be empty"));
594    }
595    if key.contains("..") {
596        return Err(Error::msg("key must not contain '..'"));
597    }
598    Ok(())
599}
600
601pub(crate) fn validate_segment(segment: &str, label: &str) -> Result<()> {
602    if segment.is_empty() || segment.contains('/') || segment.contains("..") {
603        return Err(Error::msg(format!("invalid {label}: {segment}")));
604    }
605    Ok(())
606}
607
608#[cfg(test)]
609mod tests {
610    use super::*;
611    use crate::home::UnifierHome;
612    use tempfile::tempdir;
613
614    #[test]
615    fn put_get_without_flush_leaves_disk_clean() {
616        let tmp = tempdir().unwrap();
617        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
618        let mut store = HotStore::load(&home).unwrap();
619
620        store.put_key("app/theme", "dark").unwrap();
621        assert_eq!(store.get_key("app/theme").unwrap(), Some("dark".into()));
622        assert!(store.is_dirty());
623        assert!(!home.path().join("keys/app/theme").exists());
624
625        store.flush(&home).unwrap();
626        assert!(!store.is_dirty());
627        assert!(home.path().join("keys/app/theme").is_file());
628    }
629
630    #[test]
631    fn list_keys_filters_by_prefix() {
632        let tmp = tempdir().unwrap();
633        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
634        let mut store = HotStore::load(&home).unwrap();
635        store.put_key("gifts/ideas/a", "1").unwrap();
636        store.put_key("gifts/plans/a", "2").unwrap();
637        store.put_key("other/x", "3").unwrap();
638        let under = store.list_keys(Some("gifts")).unwrap();
639        assert_eq!(
640            under,
641            vec!["gifts/ideas/a".to_string(), "gifts/plans/a".to_string()]
642        );
643        assert_eq!(store.list_keys(Some("gifts/ideas")).unwrap().len(), 1);
644    }
645
646    #[test]
647    fn send_without_flush_then_flush_persists() {
648        let tmp = tempdir().unwrap();
649        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
650        let mut store = HotStore::load(&home).unwrap();
651
652        store.send("worker", "hello").unwrap();
653        assert!(store.poll_mailbox("worker").unwrap().len() == 1);
654        assert!(fs::read_dir(home.path().join("mailbox")).is_err());
655
656        store.flush(&home).unwrap();
657        assert!(home.path().join("mailbox/worker").is_dir());
658    }
659
660    #[test]
661    fn load_existing_keys_from_disk() {
662        let tmp = tempdir().unwrap();
663        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
664        put_key_fs(&home, "k", "v").unwrap();
665
666        let store = HotStore::load(&home).unwrap();
667        assert_eq!(store.get_key("k").unwrap(), Some("v".into()));
668    }
669
670    fn put_key_fs(home: &UnifierHome, key: &str, value: &str) -> Result<()> {
671        write_text_atomic(&key_path(home, key), value)
672    }
673}
unifier/store.rs

unifier/
store.rs
