use std::collections::BTreeMap;
use std::path::Path;
use std::sync::Arc;
use web_time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use log::{debug, error, info};
use persy::{Persy, PersyId};
use serde::{Deserialize, Serialize};
use crate::actor::{Actor, ActorContext, Addr};
use crate::message::{BatchPut, Get, Message, Put};
use crate::types::*;
pub(crate) const BEAM_NODES: &str = "beam_nodes_v1";
#[derive(Serialize, Deserialize, Default, Debug, Clone)]
pub(crate) struct NodeRecord {
pub(crate) node_id: String,
pub(crate) children: Children,
}
pub struct PersyStorage {
db: Arc<Persy>,
path: String,
}
impl Clone for PersyStorage {
fn clone(&self) -> Self {
Self {
db: Arc::clone(&self.db),
path: self.path.clone(),
}
}
}
impl Default for PersyStorage {
fn default() -> Self {
Self::new()
}
}
impl PersyStorage {
pub fn new() -> Self {
Self::new_with_path("beam.persy")
}
pub fn new_with_path<P: AsRef<Path>>(path: P) -> Self {
let path = path.as_ref().to_string_lossy().into_owned();
let db = Persy::open_or_create_with(&path, persy::Config::new(), |persy| {
let mut tx = persy
.begin()
.map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
tx.create_segment(BEAM_NODES)
.map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
tx.prepare()
.map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?
.commit()
.map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
Ok(())
})
.unwrap_or_else(|e| {
panic!("Failed to create/open persy at {}: {:?}", path, e);
});
Self {
db: Arc::new(db),
path,
}
}
fn handle_get(&self, get: Get, ctx: &ActorContext) {
let segment_id = match self.db.solve_segment_id(BEAM_NODES) {
Ok(id) => id,
Err(e) => {
error!("persy solve_segment_id failed: {:?}", e);
return;
}
};
let scan_iter = match self.db.scan(segment_id) {
Ok(it) => it,
Err(e) => {
error!("persy scan failed: {:?}", e);
return;
}
};
let mut reply_children = BTreeMap::new();
for (_id, bytes) in scan_iter {
let record: NodeRecord = match postcard::from_bytes(&bytes) {
Ok(r) => r,
Err(e) => {
error!("persy get: deserialize failed: {:?}", e);
continue;
}
};
if record.node_id == get.node_id {
reply_children = record.children;
break; }
}
let final_children = match &get.child_key {
Some(target) => reply_children
.into_iter()
.filter(|(k, _)| k == target)
.collect::<BTreeMap<_, _>>(),
None => reply_children,
};
let mut reply_nodes = BTreeMap::new();
reply_nodes.insert(get.node_id.clone(), final_children);
let mut put = Put::new(reply_nodes, Some(get.id.clone()), ctx.addr.clone());
put.to_string();
let is_ack = put.in_response_to.is_some();
if is_ack || put.checksum != get.checksum {
let _ = get.from.send(Message::Put(put));
} else {
debug!("persy get: checksum match, not replying");
}
}
fn apply_put_to_tx(
&self,
tx: &mut persy::Transaction,
segment_id: persy::SegmentId,
put: Put,
) -> Result<(), String> {
for (node_id, update_data) in put.updated_nodes.iter().rev() {
if !node_id.is_empty() && node_id.starts_with('_') {
continue;
}
let mut existing_children: BTreeMap<String, NodeData> = BTreeMap::new();
let mut stale_ids: Vec<PersyId> = Vec::new();
let scan_iter = tx.scan(segment_id).map_err(|e| format!("scan: {:?}", e))?;
for (id, bytes) in scan_iter {
let record: NodeRecord = match postcard::from_bytes(&bytes) {
Ok(r) => r,
Err(e) => {
error!("persy put: deserialize failed: {:?}", e);
continue;
}
};
if record.node_id == node_id {
for (k, v) in record.children {
match existing_children.get(&k) {
Some(existing) if existing.updated_at >= v.updated_at => {}
_ => {
existing_children.insert(k, v);
}
}
}
stale_ids.push(id);
}
}
for (k, v) in update_data {
match existing_children.get(k) {
Some(existing) if existing.updated_at >= v.updated_at => {}
_ => {
existing_children.insert(k.clone(), v.clone());
}
}
}
for id in &stale_ids {
tx.delete(segment_id, id)
.map_err(|e| format!("delete stale {:?}: {:?}", id, e))?;
}
if !existing_children.is_empty() {
let record = NodeRecord {
node_id: node_id.clone(),
children: existing_children,
};
let bytes = postcard::to_allocvec(&record)
.map_err(|e| format!("postcard serialize: {:?}", e))?;
tx.insert(segment_id, &bytes)
.map_err(|e| format!("insert: {:?}", e))?;
}
}
Ok(())
}
fn handle_put_internal(&self, put: Put) -> Result<(), String> {
let mut tx = self.db.begin().map_err(|e| format!("begin tx: {:?}", e))?;
let segment_id = self
.db
.solve_segment_id(BEAM_NODES)
.map_err(|e| format!("solve_segment_id: {:?}", e))?;
self.apply_put_to_tx(&mut tx, segment_id, put)?;
let prepared = tx.prepare().map_err(|e| format!("prepare: {:?}", e))?;
prepared.commit().map_err(|e| format!("commit: {:?}", e))?;
Ok(())
}
fn handle_batch_put(&self, batch: BatchPut) -> Result<(), String> {
let mut tx = self.db.begin().map_err(|e| format!("begin tx: {:?}", e))?;
let segment_id = self
.db
.solve_segment_id(BEAM_NODES)
.map_err(|e| format!("solve_segment_id: {:?}", e))?;
for put in batch.puts {
self.apply_put_to_tx(&mut tx, segment_id, put)?;
}
let prepared = tx.prepare().map_err(|e| format!("prepare: {:?}", e))?;
prepared.commit().map_err(|e| format!("commit: {:?}", e))?;
Ok(())
}
}
fn build_ack_children(
result: &Result<Result<(), String>, tokio::task::JoinError>,
) -> (BTreeMap<String, NodeData>, Option<String>) {
let now_millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as f64;
match result {
Ok(Ok(())) => (
vec![(
"_ack".to_string(),
NodeData {
value: Value::Text("ok".to_string()),
updated_at: now_millis,
},
)]
.into_iter()
.collect(),
None,
),
Ok(Err(e)) => {
error!("persy put commit failed: {}", e);
(
vec![(
"_err".to_string(),
NodeData {
value: Value::Text(e.clone()),
updated_at: now_millis,
},
)]
.into_iter()
.collect(),
Some(e.clone()),
)
}
Err(e) => {
let msg = format!("task panicked: {:?}", e);
error!("persy put task panicked: {:?}", e);
(
vec![(
"_err".to_string(),
NodeData {
value: Value::Text(msg.clone()),
updated_at: now_millis,
},
)]
.into_iter()
.collect(),
Some(msg),
)
}
}
}
#[async_trait]
impl Actor for PersyStorage {
async fn pre_start(&mut self, _ctx: &ActorContext) {
debug!("PersyStorage started at {}", self.path);
}
async fn stopping(&mut self, _ctx: &ActorContext) {
info!(
"PersyStorage stopping at {} — all writes committed",
self.path
);
}
async fn handle(&mut self, message: Arc<Message>, ctx: &ActorContext) {
match &*message {
Message::Get(get) => self.handle_get(get.clone(), ctx),
Message::Put(put) => {
let put_id = put.id.clone();
let put_from = put.from.clone();
let storage = self.clone();
let result =
tokio::task::spawn_blocking(move || storage.handle_put_internal(put.clone()))
.await;
self.send_put_ack_after_commit(&put_id, &put_from, &result, ctx);
}
Message::BatchPut(batch) => {
let batch_id = batch.id.clone();
let batch_from = batch.from.clone();
let storage = self.clone();
let result =
tokio::task::spawn_blocking(move || storage.handle_batch_put(batch.clone()))
.await;
self.send_put_ack_after_commit(&batch_id, &batch_from, &result, ctx);
}
Message::Flush(flush) => {
let flush_id = flush.id.clone();
let from_addr = flush.from.clone();
let ctx_addr = ctx.addr.clone();
let mut ack_children = BTreeMap::new();
ack_children.insert(
"_flushed".to_string(),
NodeData {
value: Value::Text("true".to_string()),
updated_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as f64,
},
);
let mut ack_nodes = BTreeMap::new();
ack_nodes.insert("_ack".to_string(), ack_children);
let mut put = Put::new(ack_nodes, Some(flush_id), ctx_addr);
put.to_string();
let _ = from_addr.send(Message::Put(put));
}
_ => {}
}
}
fn try_clone_storage(&self) -> Option<Box<dyn Actor>> {
Some(Box::new(self.clone()))
}
}
impl PersyStorage {
fn send_put_ack_after_commit(
&self,
put_id: &str,
put_from: &Addr,
result: &Result<Result<(), String>, tokio::task::JoinError>,
ctx: &ActorContext,
) {
let (ack_children, err_msg) = build_ack_children(result);
let mut nodes = BTreeMap::new();
nodes.insert("_ack".to_string(), ack_children);
let ack = Put::new(nodes, Some(put_id.to_string()), ctx.addr.clone());
let _ = put_from.send(Message::Put(ack));
if err_msg.is_some() {
debug!("persy put ack sent with _err for {}", put_id);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::actor::Addr;
use crate::types::{NodeData, Value};
fn make_node_data(value: &str, ts: f64) -> NodeData {
NodeData {
value: Value::Text(value.to_string()),
updated_at: ts,
}
}
fn fresh_path(name: &str) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
p.push(format!("beam_persy_test_{}_{}", name, std::process::id()));
let _ = std::fs::remove_file(&p);
p
}
#[test]
fn new_with_path_creates_db_and_segment() {
let path = fresh_path("open");
let storage = PersyStorage::new_with_path(&path);
assert!(storage.db.exists_segment(BEAM_NODES).unwrap_or(false));
assert!(path.exists());
let _ = std::fs::remove_file(&path);
}
#[test]
fn new_with_path_is_idempotent() {
let path = fresh_path("idem");
let _ = PersyStorage::new_with_path(&path);
let _ = PersyStorage::new_with_path(&path);
assert!(path.exists());
let _ = std::fs::remove_file(&path);
}
#[test]
fn put_then_get_roundtrips_children() {
let path = fresh_path("rt");
let storage = PersyStorage::new_with_path(&path);
let mut children = BTreeMap::new();
children.insert("b".to_string(), make_node_data("hello", 100.0));
let mut updated_nodes = BTreeMap::new();
updated_nodes.insert("a".to_string(), children);
let put = Put::new(updated_nodes, Some("put-1".to_string()), Addr::noop());
storage
.handle_put_internal(put)
.expect("put commit should succeed");
let segment_id = storage.db.solve_segment_id(BEAM_NODES).unwrap();
let scan_iter = storage.db.scan(segment_id).unwrap();
let mut found = false;
for (_id, bytes) in scan_iter {
let record: NodeRecord = postcard::from_bytes(&bytes).unwrap();
if record.node_id == "a" {
found = true;
assert_eq!(record.children.len(), 1);
let child = record.children.get("b").unwrap();
match &child.value {
Value::Text(s) => assert_eq!(s, "hello"),
_ => panic!("expected text value"),
}
assert_eq!(child.updated_at, 100.0);
}
}
assert!(found, "node 'a' must be present in segment after Put");
let _ = std::fs::remove_file(&path);
}
#[test]
fn lww_merge_prefers_newer_updated_at() {
let path = fresh_path("lww");
let storage = PersyStorage::new_with_path(&path);
let mut children = BTreeMap::new();
children.insert("x".to_string(), make_node_data("old", 100.0));
let mut nodes = BTreeMap::new();
nodes.insert("n1".to_string(), children);
storage
.handle_put_internal(Put::new(nodes, Some("p1".to_string()), Addr::noop()))
.unwrap();
let mut children2 = BTreeMap::new();
children2.insert("x".to_string(), make_node_data("older", 50.0));
let mut nodes2 = BTreeMap::new();
nodes2.insert("n1".to_string(), children2);
storage
.handle_put_internal(Put::new(nodes2, Some("p2".to_string()), Addr::noop()))
.unwrap();
let mut children3 = BTreeMap::new();
children3.insert("x".to_string(), make_node_data("newest", 200.0));
let mut nodes3 = BTreeMap::new();
nodes3.insert("n1".to_string(), children3);
storage
.handle_put_internal(Put::new(nodes3, Some("p3".to_string()), Addr::noop()))
.unwrap();
let segment_id = storage.db.solve_segment_id(BEAM_NODES).unwrap();
let scan_iter = storage.db.scan(segment_id).unwrap();
let mut record_count = 0;
let mut found_value = None;
for (_id, bytes) in scan_iter {
let record: NodeRecord = postcard::from_bytes(&bytes).unwrap();
if record.node_id == "n1" {
record_count += 1;
let child = record.children.get("x").unwrap();
if let Value::Text(s) = &child.value {
found_value = Some(s.clone());
}
}
}
assert_eq!(
record_count, 1,
"stale records must be deleted, leaving 1 fresh record"
);
assert_eq!(found_value.as_deref(), Some("newest"));
let _ = std::fs::remove_file(&path);
}
#[test]
fn handle_get_missing_node_returns_empty_reply() {
let path = fresh_path("missing");
let storage = PersyStorage::new_with_path(&path);
let segment_id = storage.db.solve_segment_id(BEAM_NODES).unwrap();
let scan_iter = storage.db.scan(segment_id).unwrap();
let mut count = 0;
for _ in scan_iter {
count += 1;
}
assert_eq!(count, 0, "fresh db has no records");
let _ = std::fs::remove_file(&path);
}
}