use std::fs::{File, OpenOptions};
use std::path::{Path, PathBuf};
use fjall::{Database, Keyspace, KeyspaceCreateOptions};
use fs2::FileExt;
use thiserror::Error;
use super::COMMS_SCHEMA_VER;
use super::ids::{AgentId, ThreadId};
use super::keys;
use super::model::{
AgentRecord, MESSAGE_REFERENCE_PREFIX, Membership, MessageBody, MessageMeta, Thread, message_reference, now_micros,
};
mod idempotency;
mod retention;
pub use idempotency::{IDEMPOTENCY_TTL, MAX_IDEMPOTENCY_KEY_BYTES, PostOutcome, valid_idempotency_key};
const META_SCHEMA_VER: &[u8] = b"schema_ver";
const META_LAST_MAINTENANCE: &[u8] = b"last_maintenance_micros";
const STORE_DIR: &str = "store.fjall";
const LOCK_FILE: &str = ".lock";
pub const MESSAGE_TTL: std::time::Duration = std::time::Duration::from_secs(3 * 24 * 60 * 60);
pub const MAX_MESSAGES_PER_THREAD: u64 = 1_000;
pub const MAX_ACTIVE_THREADS: usize = 256;
pub const THREAD_IDLE_TTL: std::time::Duration = std::time::Duration::from_secs(7 * 24 * 60 * 60);
pub const THREAD_RETENTION_TTL: std::time::Duration = std::time::Duration::from_secs(14 * 24 * 60 * 60);
pub const EPHEMERAL_AGENT_TTL: std::time::Duration = std::time::Duration::from_secs(3 * 24 * 60 * 60);
pub const AGENT_TOUCH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
const LOCK_ATTEMPTS: u32 = 25;
const LOCK_BACKOFF: std::time::Duration = std::time::Duration::from_millis(20);
#[derive(Debug, Error)]
pub enum CommsStoreError {
#[error("fjall error: {0}")]
Fjall(#[from] fjall::Error),
#[error("io error on {path}: {source}")]
Io {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("msgpack encode error: {0}")]
Encode(#[from] rmp_serde::encode::Error),
#[error("msgpack decode error: {0}")]
Decode(#[from] rmp_serde::decode::Error),
#[error("comms retention limit reached: {0}")]
Limit(&'static str),
#[error("another basemind comms daemon holds the lock on {0}")]
Locked(PathBuf),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum MessageReferenceResolution {
Found {
message_id: String,
thread: ThreadId,
seq: u64,
},
Malformed,
Missing,
Ambiguous,
}
pub struct CommsStore {
db: Database,
meta: Keyspace,
threads: Keyspace,
thread_members: Keyspace,
messages_by_thread: Keyspace,
message_body: Keyspace,
thread_subs: Keyspace,
cursors: Keyspace,
agents: Keyspace,
post_lock: std::sync::Mutex<()>,
_lock: File,
}
impl CommsStore {
pub fn record_maintenance(&self) -> Result<(), CommsStoreError> {
self.meta.insert(META_LAST_MAINTENANCE, now_micros().to_be_bytes())?;
Ok(())
}
pub fn last_maintenance(&self) -> Result<Option<i64>, CommsStoreError> {
Ok(self
.meta
.get(META_LAST_MAINTENANCE)?
.and_then(|value| <[u8; 8]>::try_from(value.as_ref()).ok().map(i64::from_be_bytes)))
}
pub fn open(comms_dir: &Path) -> Result<Self, CommsStoreError> {
std::fs::create_dir_all(comms_dir).map_err(|source| CommsStoreError::Io {
path: comms_dir.to_path_buf(),
source,
})?;
let lock = acquire_lock(comms_dir)?;
let dir = comms_dir.join(STORE_DIR);
std::fs::create_dir_all(&dir).map_err(|source| CommsStoreError::Io {
path: dir.clone(),
source,
})?;
let mut db = Database::builder(&dir).worker_threads(1).open()?;
let mut meta = db.keyspace("meta", KeyspaceCreateOptions::default)?;
let on_disk_ver = meta
.get(META_SCHEMA_VER)?
.and_then(|bytes| <[u8; 4]>::try_from(&bytes[..]).ok())
.map(u32::from_be_bytes);
if matches!(on_disk_ver, Some(ver) if ver != COMMS_SCHEMA_VER) {
drop(meta);
drop(db);
std::fs::remove_dir_all(&dir).map_err(|source| CommsStoreError::Io {
path: dir.clone(),
source,
})?;
std::fs::create_dir_all(&dir).map_err(|source| CommsStoreError::Io {
path: dir.clone(),
source,
})?;
db = Database::builder(&dir).worker_threads(1).open()?;
meta = db.keyspace("meta", KeyspaceCreateOptions::default)?;
}
let threads = db.keyspace("threads", KeyspaceCreateOptions::default)?;
let thread_members = db.keyspace("thread_members", KeyspaceCreateOptions::default)?;
let messages_by_thread = db.keyspace("messages_by_thread", KeyspaceCreateOptions::default)?;
let message_body = db.keyspace("message_body", KeyspaceCreateOptions::default)?;
let thread_subs = db.keyspace("thread_subs", KeyspaceCreateOptions::default)?;
let cursors = db.keyspace("cursors", KeyspaceCreateOptions::default)?;
let agents = db.keyspace("agents", KeyspaceCreateOptions::default)?;
meta.insert(META_SCHEMA_VER, COMMS_SCHEMA_VER.to_be_bytes())?;
Ok(Self {
db,
meta,
threads,
thread_members,
messages_by_thread,
message_body,
thread_subs,
cursors,
agents,
post_lock: std::sync::Mutex::new(()),
_lock: lock,
})
}
pub fn put_thread(&self, thread: &Thread) -> Result<(), CommsStoreError> {
if thread.active
&& self.get_thread(&thread.id)?.is_none()
&& self.list_threads()?.iter().filter(|existing| existing.active).count() >= MAX_ACTIVE_THREADS
{
return Err(CommsStoreError::Limit("maximum active thread count"));
}
let bytes = rmp_serde::to_vec_named(thread)?;
self.threads.insert(keys::thread_key(thread.id.as_str()), bytes)?;
Ok(())
}
pub fn get_thread(&self, thread: &ThreadId) -> Result<Option<Thread>, CommsStoreError> {
match self.threads.get(keys::thread_key(thread.as_str()))? {
Some(v) => Ok(Some(rmp_serde::from_slice(&v)?)),
None => Ok(None),
}
}
pub fn list_threads(&self) -> Result<Vec<Thread>, CommsStoreError> {
let mut out = Vec::new();
for guard in self.threads.iter() {
let (_, v) = guard.into_inner()?;
out.push(rmp_serde::from_slice(&v)?);
}
Ok(out)
}
pub fn put_agent(&self, agent: &AgentRecord) -> Result<(), CommsStoreError> {
let bytes = rmp_serde::to_vec_named(agent)?;
self.agents.insert(keys::agent_key(agent.agent_id.as_str()), bytes)?;
Ok(())
}
pub fn get_agent(&self, agent: &AgentId) -> Result<Option<AgentRecord>, CommsStoreError> {
match self.agents.get(keys::agent_key(agent.as_str()))? {
Some(v) => Ok(Some(rmp_serde::from_slice(&v)?)),
None => Ok(None),
}
}
pub fn list_agents(&self) -> Result<Vec<AgentRecord>, CommsStoreError> {
let mut out = Vec::new();
for guard in self.agents.iter() {
let (_, v) = guard.into_inner()?;
out.push(rmp_serde::from_slice(&v)?);
}
Ok(out)
}
pub fn add_member(&self, membership: &Membership) -> Result<(), CommsStoreError> {
let key = keys::thread_agent(membership.thread.as_str(), membership.agent_id.as_str());
let bytes = rmp_serde::to_vec_named(membership)?;
self.thread_members.insert(&key, &bytes)?;
self.thread_subs.insert(&key, &bytes)?;
Ok(())
}
pub fn remove_member(&self, thread: &ThreadId, agent: &AgentId) -> Result<(), CommsStoreError> {
let key = keys::thread_agent(thread.as_str(), agent.as_str());
self.thread_members.remove(&key)?;
self.thread_subs.remove(&key)?;
Ok(())
}
pub fn members(&self, thread: &ThreadId) -> Result<Vec<AgentId>, CommsStoreError> {
let prefix = keys::thread_agent_prefix(thread.as_str());
let mut out = Vec::new();
for guard in self.thread_members.prefix(prefix) {
let (k, _) = guard.into_inner()?;
if let Some((_, agent)) = keys::parse_thread_agent(&k)
&& let Ok(id) = AgentId::parse(agent)
{
out.push(id);
}
}
Ok(out)
}
pub fn is_member(&self, thread: &ThreadId, agent: &AgentId) -> Result<bool, CommsStoreError> {
let key = keys::thread_agent(thread.as_str(), agent.as_str());
Ok(self.thread_members.get(key)?.is_some())
}
pub fn threads_for_agent(&self, agent: &AgentId) -> Result<Vec<ThreadId>, CommsStoreError> {
let mut out = Vec::new();
for guard in self.thread_members.iter() {
let (k, _) = guard.into_inner()?;
if let Some((thread, a)) = keys::parse_thread_agent(&k)
&& a == agent.as_str()
&& let Ok(id) = ThreadId::parse(thread)
{
out.push(id);
}
}
Ok(out)
}
fn current_seq(&self, thread: &ThreadId) -> Result<u64, CommsStoreError> {
let key = keys::thread_seq_meta_key(thread.as_str());
Ok(match self.meta.get(&key)? {
Some(v) if v.len() == 8 => u64::from_be_bytes([v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7]]),
_ => 0,
})
}
pub fn post(
&self,
thread: &ThreadId,
meta: MessageMeta,
body: MessageBody,
) -> Result<(u64, MessageMeta), CommsStoreError> {
self.post_with_row(thread, meta, body, None)
}
fn post_with_row(
&self,
thread: &ThreadId,
meta: MessageMeta,
body: MessageBody,
extra_meta_row: Option<(Vec<u8>, Vec<u8>)>,
) -> Result<(u64, MessageMeta), CommsStoreError> {
let seq = self.current_seq(thread)?.saturating_add(1);
let mut batch = self.db.batch();
if let Some((key, value)) = extra_meta_row {
batch.insert(&self.meta, key, value);
}
batch.insert(
&self.meta,
keys::thread_seq_meta_key(thread.as_str()),
seq.to_be_bytes(),
);
let meta_key = keys::message_by_thread(thread.as_str(), seq);
let meta_bytes = rmp_serde::to_vec_named(&meta)?;
batch.insert(&self.messages_by_thread, meta_key, meta_bytes);
let body_bytes = rmp_serde::to_vec_named(&body)?;
batch.insert(&self.message_body, meta.id.as_bytes().to_vec(), body_bytes);
if let Some(expired_seq) = seq.checked_sub(MAX_MESSAGES_PER_THREAD) {
let expired_key = keys::message_by_thread(thread.as_str(), expired_seq);
if let Some(bytes) = self.messages_by_thread.get(&expired_key)? {
let expired: MessageMeta = rmp_serde::from_slice(&bytes)?;
batch.remove(&self.messages_by_thread, expired_key);
batch.remove(&self.message_body, expired.id.as_bytes().to_vec());
}
}
batch.commit()?;
Ok((seq, meta))
}
pub fn history(&self, thread: &ThreadId, after_seq: u64, limit: usize) -> Result<HistoryPage, CommsStoreError> {
self.history_since(thread, after_seq, limit, None)
}
pub fn history_since(
&self,
thread: &ThreadId,
after_seq: u64,
limit: usize,
since_micros: Option<i64>,
) -> Result<HistoryPage, CommsStoreError> {
let prefix = keys::messages_by_thread_prefix(thread.as_str());
let mut messages = Vec::new();
let mut last_seq = after_seq;
let mut more = false;
for guard in self.messages_by_thread.prefix(&prefix) {
let (k, v) = guard.into_inner()?;
let Some((_, seq)) = keys::parse_message_by_thread(&k) else {
continue;
};
if seq <= after_seq {
continue;
}
let meta: MessageMeta = rmp_serde::from_slice(&v)?;
if since_micros.is_some_and(|cutoff| meta.ts_micros < cutoff) {
last_seq = seq;
continue;
}
if messages.len() >= limit {
more = true;
break;
}
messages.push((seq, meta));
last_seq = seq;
}
Ok(HistoryPage {
messages,
last_seq,
more,
})
}
pub fn history_with_seq(
&self,
thread: &ThreadId,
after_seq: u64,
limit: usize,
) -> Result<Vec<(u64, MessageMeta)>, CommsStoreError> {
let start = keys::message_by_thread(thread.as_str(), after_seq.saturating_add(1));
let end = keys::message_by_thread(thread.as_str(), u64::MAX);
let mut out = Vec::new();
for guard in self.messages_by_thread.range(start..=end) {
let (k, v) = guard.into_inner()?;
let Some((_, seq)) = keys::parse_message_by_thread(&k) else {
continue;
};
if seq <= after_seq {
continue;
}
if out.len() >= limit {
break;
}
out.push((seq, rmp_serde::from_slice(&v)?));
}
Ok(out)
}
pub fn get_body(&self, message_id: &str) -> Result<Option<Vec<u8>>, CommsStoreError> {
match self.message_body.get(message_id.as_bytes())? {
Some(v) => {
let body: MessageBody = rmp_serde::from_slice(&v)?;
Ok(Some(body.0))
}
None => Ok(None),
}
}
pub fn resolve_ids(&self, message_ids: &[String]) -> Result<Vec<(String, ThreadId, u64)>, CommsStoreError> {
if message_ids.is_empty() {
return Ok(Vec::new());
}
let wanted: ahash::AHashSet<&str> = message_ids.iter().map(String::as_str).collect();
let mut out = Vec::with_capacity(message_ids.len());
for guard in self.messages_by_thread.iter() {
let (k, v) = guard.into_inner()?;
let Some((_, seq)) = keys::parse_message_by_thread(&k) else {
continue;
};
let meta: MessageMeta = rmp_serde::from_slice(&v)?;
if wanted.contains(meta.id.as_str()) {
out.push((meta.id.clone(), meta.thread, seq));
if out.len() == wanted.len() {
break;
}
}
}
Ok(out)
}
pub fn resolve_message_reference(&self, reference: &str) -> Result<MessageReferenceResolution, CommsStoreError> {
let compact = reference.strip_prefix(MESSAGE_REFERENCE_PREFIX);
if let Some(hash_prefix) = compact
&& (!(4..=16).contains(&hash_prefix.len()) || !hash_prefix.bytes().all(|byte| byte.is_ascii_hexdigit()))
{
return Ok(MessageReferenceResolution::Malformed);
}
let mut found = None;
for guard in self.messages_by_thread.iter() {
let (key, value) = guard.into_inner()?;
let Some((_, seq)) = keys::parse_message_by_thread(&key) else {
continue;
};
let meta: MessageMeta = rmp_serde::from_slice(&value)?;
let matches = match compact {
Some(hash_prefix) => message_reference(&meta.id)[MESSAGE_REFERENCE_PREFIX.len()..]
.starts_with(&hash_prefix.to_ascii_lowercase()),
None => meta.id == reference,
};
if !matches {
continue;
}
if found.is_some() {
return Ok(MessageReferenceResolution::Ambiguous);
}
found = Some(MessageReferenceResolution::Found {
message_id: meta.id,
thread: meta.thread,
seq,
});
}
Ok(found.unwrap_or(MessageReferenceResolution::Missing))
}
pub fn read_cursor(&self, agent: &AgentId, thread: &ThreadId) -> Result<u64, CommsStoreError> {
let key = keys::cursor_key(agent.as_str(), thread.as_str());
match self.cursors.get(key)? {
Some(v) if v.len() == 8 => Ok(u64::from_be_bytes([v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7]])),
_ => Ok(0),
}
}
pub fn set_read_cursor(&self, agent: &AgentId, thread: &ThreadId, seq: u64) -> Result<(), CommsStoreError> {
let current = self.read_cursor(agent, thread)?;
if seq <= current {
return Ok(());
}
let key = keys::cursor_key(agent.as_str(), thread.as_str());
self.cursors.insert(key, seq.to_be_bytes())?;
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HistoryPage {
pub messages: Vec<(u64, MessageMeta)>,
pub last_seq: u64,
pub more: bool,
}
pub fn body_hash_hex(body: &[u8]) -> String {
crate::hashing::hex(&crate::hashing::hash_bytes(body))
}
pub fn build_meta(
id: String,
thread: ThreadId,
from: AgentId,
subject: String,
tags: Vec<String>,
reply_to: Option<String>,
body: &[u8],
) -> MessageMeta {
MessageMeta {
id,
thread,
from,
ts_micros: now_micros(),
subject,
tags,
reply_to,
body_len: u32::try_from(body.len()).unwrap_or(u32::MAX),
body_sha: body_hash_hex(body),
}
}
fn acquire_lock(comms_dir: &Path) -> Result<File, CommsStoreError> {
let path = comms_dir.join(LOCK_FILE);
let file = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(false)
.open(&path)
.map_err(|source| CommsStoreError::Io {
path: path.clone(),
source,
})?;
for attempt in 0..LOCK_ATTEMPTS {
match file.try_lock_exclusive() {
Ok(()) => return Ok(file),
Err(_) if attempt + 1 < LOCK_ATTEMPTS => std::thread::sleep(LOCK_BACKOFF),
Err(_) => return Err(CommsStoreError::Locked(path)),
}
}
unreachable!("loop returns on the final attempt")
}
#[cfg(test)]
#[path = "store_tests.rs"]
mod tests;