use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::fs;
use std::path::{Path, PathBuf};
use uuid::Uuid;
use crate::constants::{CRON, KEYS, MAILBOX};
use crate::cron::CronSchedule;
use crate::error::{Error, Result};
use crate::fs_text::{read_text, write_text, write_text_atomic};
use crate::home::UnifierHome;
use crate::paths::{cron_dir, key_path, mailbox_dir, message_path, parse_message_id};
use crate::postbox::Message;
use crate::scope::resolve_under_root;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum KeyState {
Present { value: String, dirty: bool },
Deleted { dirty: bool },
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct MsgState {
body: String,
dirty: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum MessageKind {
Mailbox,
Cron,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct EventState {
pub(crate) body: String,
pub(crate) created_at: String,
pub(crate) expires_at: Option<String>,
pub(crate) dirty: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum StagingValue {
Present(String),
Deleted,
}
#[derive(Debug, Clone)]
pub(crate) struct ActiveTick {
pub number: u64,
pub(crate) read_snapshot: BTreeMap<String, String>,
pub(crate) staging: BTreeMap<String, StagingValue>,
pub(crate) locks: BTreeSet<String>,
}
#[derive(Debug, Default)]
pub struct HotStore {
pub(crate) keys: BTreeMap<String, KeyState>,
mailboxes: BTreeMap<String, BTreeMap<Uuid, MsgState>>,
cron: BTreeMap<String, BTreeMap<Uuid, MsgState>>,
pub(crate) events: BTreeMap<Uuid, EventState>,
removed_messages: BTreeSet<(MessageKind, String, Uuid)>,
pub(crate) committed_tick: u64,
pub(crate) active_tick: Option<ActiveTick>,
pub(crate) tick_queue: VecDeque<String>,
}
impl HotStore {
pub fn load(home: &UnifierHome) -> Result<Self> {
let mut store = Self::default();
store.load_keys(home)?;
store.mailboxes = load_message_buckets(home, MAILBOX)?;
store.cron = load_message_buckets(home, CRON)?;
store.events = crate::tick::load_events(home)?;
store.committed_tick = crate::tick::load_committed_tick(home)?;
Ok(store)
}
pub fn is_dirty(&self) -> bool {
self.keys.values().any(|k| match k {
KeyState::Present { dirty, .. } | KeyState::Deleted { dirty } => *dirty,
}) || self.has_dirty_messages(&self.mailboxes)
|| self.has_dirty_messages(&self.cron)
|| self.events.values().any(|e| e.dirty)
|| !self.removed_messages.is_empty()
|| self.active_tick.is_some()
}
pub fn put_key(&mut self, key: &str, value: &str) -> Result<()> {
validate_key(key)?;
if let Some(tick) = &self.active_tick {
if tick.locks.contains(key) {
return Err(Error::msg(format!(
"key locked during tick {}: {key}",
tick.number
)));
}
}
if let Some(tick) = &mut self.active_tick {
tick.staging
.insert(key.to_string(), StagingValue::Present(value.to_string()));
return Ok(());
}
self.keys.insert(
key.to_string(),
KeyState::Present {
value: value.to_string(),
dirty: true,
},
);
Ok(())
}
pub fn get_key(&self, key: &str) -> Result<Option<String>> {
validate_key(key)?;
if let Some(tick) = &self.active_tick {
if let Some(staged) = tick.staging.get(key) {
return Ok(match staged {
StagingValue::Present(v) => Some(v.clone()),
StagingValue::Deleted => None,
});
}
return Ok(tick.read_snapshot.get(key).cloned());
}
Ok(match self.keys.get(key) {
Some(KeyState::Present { value, .. }) => Some(value.clone()),
Some(KeyState::Deleted { .. }) | None => None,
})
}
pub fn delete_key(&mut self, key: &str) -> Result<bool> {
validate_key(key)?;
if let Some(tick) = &self.active_tick {
if tick.locks.contains(key) {
return Err(Error::msg(format!(
"key locked during tick {}: {key}",
tick.number
)));
}
}
if let Some(tick) = &mut self.active_tick {
let existed = tick.staging.contains_key(key) || tick.read_snapshot.contains_key(key);
if !existed {
return Ok(false);
}
tick.staging.insert(key.to_string(), StagingValue::Deleted);
return Ok(true);
}
let existed = matches!(self.keys.get(key), Some(KeyState::Present { .. }));
if !existed {
return Ok(false);
}
self.keys
.insert(key.to_string(), KeyState::Deleted { dirty: true });
Ok(true)
}
pub fn send(&mut self, recipient: &str, body: &str) -> Result<Uuid> {
self.send_from(crate::envelope::DEFAULT_SENDER, recipient, body)
}
pub fn send_from(&mut self, from: &str, recipient: &str, body: &str) -> Result<Uuid> {
validate_segment(from, "from")?;
validate_segment(recipient, "recipient")?;
let env = crate::envelope::Envelope::new(
from,
recipient,
crate::envelope::Envelope::parse_payload(body),
);
let id = env.id;
self.mailboxes
.entry(recipient.to_string())
.or_default()
.insert(
id,
MsgState {
body: env.to_json()?,
dirty: true,
},
);
Ok(id)
}
pub fn post_cron(&mut self, schedule: &str, body: &str) -> Result<Uuid> {
CronSchedule::parse(schedule)?;
let id = Uuid::new_v4();
self.cron.entry(schedule.to_string()).or_default().insert(
id,
MsgState {
body: body.to_string(),
dirty: true,
},
);
Ok(id)
}
pub fn poll_mailbox(&self, recipient: &str) -> Result<Vec<Message>> {
validate_segment(recipient, "recipient")?;
Ok(self.collect_messages(
MessageKind::Mailbox,
recipient,
self.mailboxes.get(recipient),
))
}
pub fn poll_cron(&self) -> Result<Vec<Message>> {
let mut out = Vec::new();
for (schedule, msgs) in &self.cron {
let parsed = CronSchedule::parse(schedule)?;
if !parsed.matches_now() {
continue;
}
out.extend(self.collect_messages(MessageKind::Cron, schedule, Some(msgs)));
}
out.sort_by(|a, b| a.path.cmp(&b.path));
Ok(out)
}
pub fn list_dir(&self, home: &UnifierHome, subpath: &str) -> Result<Vec<Message>> {
let dir = resolve_under_root(home.path(), subpath)?;
let rel = dir
.strip_prefix(home.path())
.map_err(|_| Error::msg("path escapes store root"))?;
let parts: Vec<_> = rel.iter().collect();
match parts.as_slice() {
[p1, p2] if p1.to_string_lossy() == MAILBOX => self.poll_mailbox(&p2.to_string_lossy()),
[p1, p2] if p1.to_string_lossy() == CRON => {
let schedule = p2.to_string_lossy();
Ok(self.collect_messages(
MessageKind::Cron,
&schedule,
self.cron.get(schedule.as_ref()),
))
}
_ => Ok(Vec::new()),
}
}
pub fn ack(&mut self, home: &UnifierHome, id_or_path: &str) -> Result<bool> {
if id_or_path.contains('/') || Path::new(id_or_path).is_absolute() {
let path = if Path::new(id_or_path).is_absolute() {
PathBuf::from(id_or_path)
} else {
resolve_under_root(home.path(), id_or_path)?
};
return self.ack_path(home, &path);
}
let id = Uuid::parse_str(id_or_path)
.map_err(|_| Error::msg(format!("invalid message id or path: {id_or_path}")))?;
self.ack_id(home, &id)
}
pub fn flush(&mut self, home: &UnifierHome) -> Result<()> {
if self.active_tick.is_some() {
return Err(Error::msg(
"cannot flush while a tick is active; run tick end first",
));
}
home.ensure()?;
self.flush_keys(home)?;
self.flush_mailboxes(home)?;
self.flush_cron(home)?;
self.flush_events(home)?;
self.removed_messages.clear();
Ok(())
}
fn load_keys(&mut self, home: &UnifierHome) -> Result<()> {
let root = home.path().join(KEYS);
if !root.is_dir() {
return Ok(());
}
let mut prefix = Vec::new();
self.walk_key_files(&root, &mut prefix)?;
Ok(())
}
fn walk_key_files(&mut self, dir: &Path, prefix: &mut Vec<String>) -> Result<()> {
for entry in fs::read_dir(dir)? {
let entry = entry?;
let name = entry.file_name().to_string_lossy().into_owned();
if entry.file_type()?.is_dir() {
prefix.push(name);
self.walk_key_files(&entry.path(), prefix)?;
prefix.pop();
} else if entry.file_type()?.is_file() {
prefix.push(name);
let key = prefix.join("/");
prefix.pop();
let value = read_text(&entry.path())?;
self.keys.insert(
key,
KeyState::Present {
value,
dirty: false,
},
);
}
}
Ok(())
}
fn flush_keys(&mut self, home: &UnifierHome) -> Result<()> {
let mut to_remove = Vec::new();
for (key, state) in &mut self.keys {
match state {
KeyState::Present { value, dirty: true } => {
write_text_atomic(&key_path(home, key), value)?;
*state = KeyState::Present {
value: value.clone(),
dirty: false,
};
}
KeyState::Deleted { dirty: true } => {
let path = key_path(home, key);
if path.is_file() {
fs::remove_file(&path)?;
}
to_remove.push(key.clone());
}
_ => {}
}
}
for key in to_remove {
self.keys.remove(&key);
}
Ok(())
}
fn flush_mailboxes(&mut self, home: &UnifierHome) -> Result<()> {
flush_message_buckets(
home,
MessageKind::Mailbox,
&mut self.mailboxes,
&mut self.removed_messages,
mailbox_dir,
)
}
fn flush_cron(&mut self, home: &UnifierHome) -> Result<()> {
flush_message_buckets(
home,
MessageKind::Cron,
&mut self.cron,
&mut self.removed_messages,
cron_dir,
)
}
fn ack_path(&mut self, home: &UnifierHome, path: &Path) -> Result<bool> {
let rel = path.strip_prefix(home.path()).ok();
let Some(rel) = rel else {
return Ok(false);
};
let parts: Vec<_> = rel
.iter()
.map(|p| p.to_string_lossy().into_owned())
.collect();
if parts.len() != 3 {
return Ok(false);
}
let Some(id) = parse_message_id(&parts[2]) else {
return Ok(false);
};
let existed = match parts[0].as_str() {
MAILBOX => {
let had = self
.mailboxes
.get(&parts[1])
.is_some_and(|m| m.contains_key(&id))
|| path.is_file();
if had {
self.removed_messages
.insert((MessageKind::Mailbox, parts[1].clone(), id));
if let Some(msgs) = self.mailboxes.get_mut(&parts[1]) {
msgs.remove(&id);
}
}
had
}
CRON => {
let had = self
.cron
.get(&parts[1])
.is_some_and(|m| m.contains_key(&id))
|| path.is_file();
if had {
self.removed_messages
.insert((MessageKind::Cron, parts[1].clone(), id));
if let Some(msgs) = self.cron.get_mut(&parts[1]) {
msgs.remove(&id);
}
}
had
}
_ => false,
};
Ok(existed)
}
fn ack_id(&mut self, home: &UnifierHome, id: &Uuid) -> Result<bool> {
if let Some(recipient) = self
.mailboxes
.iter()
.find_map(|(name, msgs)| msgs.contains_key(id).then(|| name.clone()))
{
self.removed_messages
.insert((MessageKind::Mailbox, recipient.clone(), *id));
if let Some(m) = self.mailboxes.get_mut(&recipient) {
m.remove(id);
}
return Ok(true);
}
if let Some(schedule) = self
.cron
.iter()
.find_map(|(name, msgs)| msgs.contains_key(id).then(|| name.clone()))
{
self.removed_messages
.insert((MessageKind::Cron, schedule.clone(), *id));
if let Some(m) = self.cron.get_mut(&schedule) {
m.remove(id);
}
return Ok(true);
}
let filename = format!("{}.txt", id.hyphenated());
for sub in [MAILBOX, CRON] {
let base = home.path().join(sub);
if !base.is_dir() {
continue;
}
for entry in fs::read_dir(&base)? {
let entry = entry?;
if entry.file_type()?.is_dir() {
let candidate = entry.path().join(&filename);
if candidate.is_file() {
let bucket = entry.file_name().to_string_lossy().into_owned();
let kind = if sub == MAILBOX {
MessageKind::Mailbox
} else {
MessageKind::Cron
};
self.removed_messages.insert((kind, bucket, *id));
return Ok(true);
}
}
}
}
Ok(false)
}
fn collect_messages(
&self,
kind: MessageKind,
bucket: &str,
msgs: Option<&BTreeMap<Uuid, MsgState>>,
) -> Vec<Message> {
let Some(msgs) = msgs else {
return Vec::new();
};
let base = match kind {
MessageKind::Mailbox => PathBuf::from(MAILBOX).join(bucket),
MessageKind::Cron => PathBuf::from(CRON).join(bucket),
};
let mut out = Vec::new();
for (id, msg) in msgs {
out.push(Message {
id: *id,
path: base.join(format!("{}.txt", id.hyphenated())),
body: msg.body.clone(),
});
}
out.sort_by_key(|m| m.id);
out
}
fn has_dirty_messages(&self, buckets: &BTreeMap<String, BTreeMap<Uuid, MsgState>>) -> bool {
buckets.values().any(|msgs| msgs.values().any(|m| m.dirty))
}
}
fn load_message_buckets(
home: &UnifierHome,
top: &str,
) -> Result<BTreeMap<String, BTreeMap<Uuid, MsgState>>> {
let mut target = BTreeMap::new();
let base = home.path().join(top);
if !base.is_dir() {
return Ok(target);
}
for entry in fs::read_dir(&base)? {
let entry = entry?;
if !entry.file_type()?.is_dir() {
continue;
}
let bucket = entry.file_name().to_string_lossy().into_owned();
let mut msgs = BTreeMap::new();
for msg_entry in fs::read_dir(entry.path())? {
let msg_entry = msg_entry?;
if !msg_entry.file_type()?.is_file() {
continue;
}
let name = msg_entry.file_name().to_string_lossy().into_owned();
let Some(id) = parse_message_id(&name) else {
continue;
};
msgs.insert(
id,
MsgState {
body: read_text(&msg_entry.path())?,
dirty: false,
},
);
}
if !msgs.is_empty() {
target.insert(bucket, msgs);
}
}
Ok(target)
}
fn flush_message_buckets<F>(
home: &UnifierHome,
kind: MessageKind,
buckets: &mut BTreeMap<String, BTreeMap<Uuid, MsgState>>,
removed_messages: &mut BTreeSet<(MessageKind, String, Uuid)>,
dir_for: F,
) -> Result<()>
where
F: Fn(&UnifierHome, &str) -> PathBuf,
{
for (bucket, msgs) in buckets.iter_mut() {
for (id, msg) in msgs.iter_mut() {
if msg.dirty {
write_text(&message_path(&dir_for(home, bucket), id), &msg.body)?;
msg.dirty = false;
}
}
}
let removed: Vec<_> = removed_messages
.iter()
.filter(|(k, _, _)| *k == kind)
.cloned()
.collect();
for (_, bucket, id) in removed {
let path = message_path(&dir_for(home, &bucket), &id);
if path.is_file() {
fs::remove_file(path)?;
}
if let Some(msgs) = buckets.get_mut(&bucket) {
msgs.remove(&id);
}
removed_messages.remove(&(kind, bucket.clone(), id));
}
buckets.retain(|_, msgs| !msgs.is_empty());
Ok(())
}
pub(crate) fn validate_key(key: &str) -> Result<()> {
if key.is_empty() {
return Err(Error::msg("key must not be empty"));
}
if key.contains("..") {
return Err(Error::msg("key must not contain '..'"));
}
Ok(())
}
pub(crate) fn validate_segment(segment: &str, label: &str) -> Result<()> {
if segment.is_empty() || segment.contains('/') || segment.contains("..") {
return Err(Error::msg(format!("invalid {label}: {segment}")));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::home::UnifierHome;
use tempfile::tempdir;
#[test]
fn put_get_without_flush_leaves_disk_clean() {
let tmp = tempdir().unwrap();
let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
let mut store = HotStore::load(&home).unwrap();
store.put_key("app/theme", "dark").unwrap();
assert_eq!(store.get_key("app/theme").unwrap(), Some("dark".into()));
assert!(store.is_dirty());
assert!(!home.path().join("keys/app/theme").exists());
store.flush(&home).unwrap();
assert!(!store.is_dirty());
assert!(home.path().join("keys/app/theme").is_file());
}
#[test]
fn send_without_flush_then_flush_persists() {
let tmp = tempdir().unwrap();
let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
let mut store = HotStore::load(&home).unwrap();
store.send("worker", "hello").unwrap();
assert!(store.poll_mailbox("worker").unwrap().len() == 1);
assert!(fs::read_dir(home.path().join("mailbox")).is_err());
store.flush(&home).unwrap();
assert!(home.path().join("mailbox/worker").is_dir());
}
#[test]
fn load_existing_keys_from_disk() {
let tmp = tempdir().unwrap();
let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
put_key_fs(&home, "k", "v").unwrap();
let store = HotStore::load(&home).unwrap();
assert_eq!(store.get_key("k").unwrap(), Some("v".into()));
}
fn put_key_fs(home: &UnifierHome, key: &str, value: &str) -> Result<()> {
write_text_atomic(&key_path(home, key), value)
}
}