use super::{CrdtError, CrdtResult};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::fs;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use yrs::updates::decoder::Decode;
use yrs::{Any, Doc, Map, MapPrelim, MapRef, Out, ReadTxn, Transact, TransactionMut, Update};
#[derive(Debug, Clone, Serialize, Deserialize)]
struct DocumentMetadata {
doc_id: String,
entity_type: String,
entity_id: String,
version: u32,
created_at: i64,
updated_at: i64,
}
#[derive(Debug, Clone)]
pub struct CompactionConfig {
pub tombstone_retention: Duration,
pub compaction_interval: Duration,
pub max_document_bytes: usize,
}
impl Default for CompactionConfig {
fn default() -> Self {
Self {
tombstone_retention: Duration::from_secs(7 * 24 * 60 * 60), compaction_interval: Duration::from_secs(60 * 60), max_document_bytes: 5 * 1024 * 1024, }
}
}
#[derive(Debug, Clone, Default)]
pub struct CompactionResult {
pub tombstones_removed: usize,
pub documents_compacted: usize,
pub bytes_saved: u64,
}
pub struct CrdtManager {
storage_dir: PathBuf,
config: CompactionConfig,
compaction_task: Arc<RwLock<Option<JoinHandle<()>>>>,
compaction_shutdown_tx: Arc<RwLock<Option<tokio::sync::mpsc::Sender<()>>>>,
}
impl CrdtManager {
pub async fn new<P: AsRef<Path>>(storage_dir: P) -> CrdtResult<Self> {
Self::new_with_config(storage_dir, CompactionConfig::default()).await
}
pub async fn new_with_config<P: AsRef<Path>>(
storage_dir: P,
config: CompactionConfig,
) -> CrdtResult<Self> {
let storage_dir = storage_dir.as_ref().to_path_buf();
let crdt_dir = storage_dir.join("crdt");
fs::create_dir_all(&crdt_dir).await.map_err(|e| {
CrdtError::FileSystem(format!("Failed to create CRDT directory: {}", e))
})?;
Ok(Self {
storage_dir,
config,
compaction_task: Arc::new(RwLock::new(None)),
compaction_shutdown_tx: Arc::new(RwLock::new(None)),
})
}
pub fn get_storage_dir(&self) -> &Path {
&self.storage_dir
}
fn entity_dir(&self, entity_type: &str) -> PathBuf {
self.storage_dir.join("crdt").join(entity_type)
}
fn sanitize_doc_id(doc_id: &str) -> String {
let hex_encoded = hex::encode(doc_id.as_bytes());
if hex_encoded.len() > 200 {
use sha2::{Digest, Sha256};
let hash = Sha256::digest(doc_id.as_bytes());
format!("h_{}", hex::encode(hash))
} else {
hex_encoded
}
}
fn doc_paths(&self, entity_type: &str, doc_id: &str) -> (PathBuf, PathBuf) {
let entity_dir = self.entity_dir(entity_type);
let safe_filename = Self::sanitize_doc_id(doc_id);
let yrs_path = entity_dir.join(format!("{}.yrs", safe_filename));
let meta_path = entity_dir.join(format!("{}.meta", safe_filename));
(yrs_path, meta_path)
}
pub async fn save_document(
&self,
doc_id: &str,
entity_type: &str,
entity_id: &str,
doc: &Doc,
) -> CrdtResult<()> {
if !doc_id.starts_with(entity_type) {
return Err(CrdtError::InvalidDocumentId(format!(
"doc_id '{}' must start with entity_type '{}'",
doc_id, entity_type
)));
}
let entity_dir = self.entity_dir(entity_type);
fs::create_dir_all(&entity_dir).await.map_err(|e| {
CrdtError::FileSystem(format!("Failed to create entity directory: {}", e))
})?;
let state = {
const MAX_TXN_RETRIES: usize = 8;
const BASE_DELAY_MS: u64 = 5;
let mut attempt = 0usize;
loop {
match doc.try_transact() {
Ok(txn) => {
break txn.encode_state_as_update_v1(&yrs::StateVector::default());
}
Err(err) => {
if attempt >= MAX_TXN_RETRIES {
return Err(CrdtError::Operation(format!(
"Failed to acquire read transaction for {} after {} attempts: {}",
doc_id,
attempt + 1,
err
)));
}
let delay = BASE_DELAY_MS.saturating_mul((attempt + 1) as u64);
tokio::time::sleep(Duration::from_millis(delay)).await;
attempt += 1;
}
}
}
};
const MAX_ENCODED_SIZE: usize = 10 * 1024 * 1024;
if state.len() > MAX_ENCODED_SIZE {
return Err(CrdtError::encoding_error(format!(
"Document too large: {} bytes (max: {})",
state.len(),
MAX_ENCODED_SIZE
)));
}
let (yrs_path, meta_path) = self.doc_paths(entity_type, doc_id);
let mut metadata = if meta_path.exists() {
let meta_json = fs::read_to_string(&meta_path)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read metadata: {}", e)))?;
let mut existing: DocumentMetadata = serde_json::from_str(&meta_json)
.map_err(|e| CrdtError::Serialization(format!("Invalid metadata JSON: {}", e)))?;
existing.version += 1; existing
} else {
let now = chrono::Utc::now().timestamp();
DocumentMetadata {
doc_id: doc_id.to_string(),
entity_type: entity_type.to_string(),
entity_id: entity_id.to_string(),
version: 1, created_at: now,
updated_at: now,
}
};
metadata.updated_at = chrono::Utc::now().timestamp();
let unique_suffix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or_default();
let temp_tag = format!("tmp.{}.{}", std::process::id(), unique_suffix);
let yrs_temp = yrs_path.with_extension(format!("yrs.{}", temp_tag));
let meta_temp = meta_path.with_extension(format!("meta.{}", temp_tag));
let meta_json = serde_json::to_string_pretty(&metadata).map_err(|e| {
CrdtError::Serialization(format!("Failed to serialize metadata: {}", e))
})?;
let yrs_write = fs::write(&yrs_temp, &state);
let meta_write = fs::write(&meta_temp, meta_json);
let (yrs_result, meta_result) = tokio::join!(yrs_write, meta_write);
yrs_result
.map_err(|e| CrdtError::FileSystem(format!("Failed to write Yrs state: {}", e)))?;
meta_result
.map_err(|e| CrdtError::FileSystem(format!("Failed to write metadata: {}", e)))?;
fs::rename(&yrs_temp, &yrs_path)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to rename Yrs file: {}", e)))?;
fs::rename(&meta_temp, &meta_path)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to rename metadata file: {}", e)))?;
Ok(())
}
pub async fn load_document(&self, doc_id: &str) -> CrdtResult<Doc> {
let parts: Vec<&str> = doc_id.split(':').collect();
if parts.len() < 2 {
return Err(CrdtError::InvalidDocumentId(format!(
"Invalid doc_id format (expected 'entity_type:entity_id:...'): {}",
doc_id
)));
}
let entity_type = parts[0];
let (yrs_path, _) = self.doc_paths(entity_type, doc_id);
if !yrs_path.exists() {
return Err(CrdtError::DocumentNotFound(doc_id.to_string()));
}
let state_bytes = fs::read(&yrs_path)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read Yrs state: {}", e)))?;
let update = Update::decode_v1(&state_bytes).map_err(|e| {
CrdtError::Deserialization(format!("Failed to decode Yrs state: {}", e))
})?;
let doc = Doc::new();
{
let mut txn = doc.transact_mut();
txn.apply_update(update);
}
Ok(doc)
}
pub async fn list_documents(&self, entity_type: &str) -> CrdtResult<Vec<String>> {
let entity_dir = self.entity_dir(entity_type);
if !entity_dir.exists() {
return Ok(Vec::new());
}
let mut doc_ids = Vec::new();
let mut read_dir = fs::read_dir(&entity_dir).await.map_err(|e| {
CrdtError::FileSystem(format!("Failed to read entity directory: {}", e))
})?;
while let Some(entry) = read_dir
.next_entry()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read directory entry: {}", e)))?
{
let file_name = entry.file_name();
let file_name_str = file_name.to_string_lossy();
if file_name_str.ends_with(".yrs") {
let hex_id = file_name_str.trim_end_matches(".yrs");
if let Ok(bytes) = hex::decode(hex_id)
&& let Ok(doc_id) = String::from_utf8(bytes)
{
doc_ids.push(doc_id);
}
}
}
Ok(doc_ids)
}
pub async fn apply_update(&self, doc_id: &str, update_bytes: &[u8]) -> CrdtResult<()> {
const MAX_ENCODED_SIZE: usize = 10 * 1024 * 1024;
if update_bytes.len() > MAX_ENCODED_SIZE {
return Err(CrdtError::encoding_error(format!(
"Update too large: {} bytes (max: {})",
update_bytes.len(),
MAX_ENCODED_SIZE
)));
}
let doc = match self.load_document(doc_id).await {
Ok(doc) => doc,
Err(CrdtError::DocumentNotFound(_)) => Doc::new(),
Err(e) => return Err(e),
};
let update = Update::decode_v1(update_bytes)
.map_err(|e| CrdtError::Deserialization(format!("Failed to decode update: {}", e)))?;
{
let mut txn = doc.transact_mut();
txn.apply_update(update);
}
let parts: Vec<&str> = doc_id.split(':').collect();
if parts.len() < 2 {
return Err(CrdtError::InvalidDocumentId(format!(
"Invalid doc_id format: {}",
doc_id
)));
}
let entity_type = parts[0];
let entity_id = parts[1];
self.save_document(doc_id, entity_type, entity_id, &doc)
.await?;
Ok(())
}
pub async fn merge_updates(&self, doc_id: &str, updates: Vec<Vec<u8>>) -> CrdtResult<Doc> {
const MAX_ENCODED_SIZE: usize = 10 * 1024 * 1024;
for (i, update_bytes) in updates.iter().enumerate() {
if update_bytes.len() > MAX_ENCODED_SIZE {
return Err(CrdtError::encoding_error(format!(
"Update {} too large: {} bytes (max: {})",
i,
update_bytes.len(),
MAX_ENCODED_SIZE
)));
}
}
let doc = match self.load_document(doc_id).await {
Ok(doc) => doc,
Err(CrdtError::DocumentNotFound(_)) => Doc::new(),
Err(e) => return Err(e),
};
for update_bytes in updates {
let update = Update::decode_v1(&update_bytes).map_err(|e| {
CrdtError::Deserialization(format!("Failed to decode update: {}", e))
})?;
let mut txn = doc.transact_mut();
txn.apply_update(update);
}
let parts: Vec<&str> = doc_id.split(':').collect();
if parts.len() >= 2 {
let entity_type = parts[0];
let entity_id = parts[1];
self.save_document(doc_id, entity_type, entity_id, &doc)
.await?;
}
Ok(doc)
}
pub async fn mark_deleted(&self, doc_id: &str, deleted_by: &str) -> CrdtResult<()> {
let doc = self.load_document(doc_id).await?;
{
let root = doc.get_or_insert_map("root");
let mut txn = doc.transact_mut();
let metadata = if let Some(existing) = root.get(&txn, "metadata") {
MapRef::try_from(existing).map_err(|e| {
CrdtError::Operation(format!("Invalid metadata structure: {:?}", e))
})?
} else {
let empty_prelim: MapPrelim = MapPrelim::from([("_", Any::Null)]);
let m = root.insert(&mut txn, "metadata", empty_prelim);
m.remove(&mut txn, "_");
m
};
metadata.insert(&mut txn, "deleted", true);
metadata.insert(&mut txn, "deleted_at", chrono::Utc::now().timestamp());
metadata.insert(&mut txn, "deleted_by", deleted_by);
}
let parts: Vec<&str> = doc_id.split(':').collect();
if parts.len() >= 2 {
let entity_type = parts[0];
let entity_id = parts[1];
self.save_document(doc_id, entity_type, entity_id, &doc)
.await?;
}
Ok(())
}
pub async fn is_deleted(&self, doc_id: &str) -> CrdtResult<bool> {
let doc = self.load_document(doc_id).await?;
let root = doc.get_or_insert_map("root");
let txn = doc.transact();
if let Some(metadata_val) = root.get(&txn, "metadata")
&& let Ok(metadata) = MapRef::try_from(metadata_val)
&& let Some(deleted_val) = metadata.get(&txn, "deleted")
&& let Ok(deleted) = bool::try_from(deleted_val)
{
return Ok(deleted);
}
Ok(false)
}
pub async fn delete_document(&self, doc_id: &str) -> CrdtResult<()> {
let parts: Vec<&str> = doc_id.split(':').collect();
if parts.len() < 2 {
return Err(CrdtError::InvalidDocumentId(format!(
"Invalid doc_id format (expected 'entity_type:entity_id:...'): {}",
doc_id
)));
}
let entity_type = parts[0];
let (yrs_path, meta_path) = self.doc_paths(entity_type, doc_id);
if !yrs_path.exists() {
return Err(CrdtError::DocumentNotFound(doc_id.to_string()));
}
fs::remove_file(&yrs_path)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to delete Yrs file: {}", e)))?;
if meta_path.exists() {
fs::remove_file(&meta_path).await.map_err(|e| {
CrdtError::FileSystem(format!("Failed to delete metadata file: {}", e))
})?;
}
Ok(())
}
pub fn get_map_value<T>(doc: &Doc, key: &str) -> CrdtResult<Option<T>>
where
T: TryFrom<Out>,
{
let map = doc.get_or_insert_map("root");
let txn = doc.transact();
Ok(map.get(&txn, key).and_then(|out| T::try_from(out).ok()))
}
pub fn set_map_value<V>(doc: &Doc, key: &str, value: V) -> CrdtResult<()>
where
V: Into<Any>,
{
let map = doc.get_or_insert_map("root");
let mut txn = doc.transact_mut();
map.insert(&mut txn, key, value);
Ok(())
}
pub fn get_map_bool(map: &MapRef, txn: &impl ReadTxn, key: &str) -> Option<bool> {
map.get(txn, key).and_then(|out| bool::try_from(out).ok())
}
pub fn get_map_string(map: &MapRef, txn: &impl ReadTxn, key: &str) -> Option<String> {
map.get(txn, key).and_then(|out| String::try_from(out).ok())
}
pub fn get_map_i64(map: &MapRef, txn: &impl ReadTxn, key: &str) -> Option<i64> {
map.get(txn, key).and_then(|out| i64::try_from(out).ok())
}
pub fn get_nested_map(map: &MapRef, txn: &impl ReadTxn, key: &str) -> Option<MapRef> {
map.get(txn, key).and_then(|out| MapRef::try_from(out).ok())
}
pub fn set_map_string(
map: &MapRef,
txn: &mut TransactionMut,
key: impl Into<String>,
value: impl Into<String>,
) {
map.insert(txn, key.into(), value.into());
}
pub fn set_map_i64(map: &MapRef, txn: &mut TransactionMut, key: impl Into<String>, value: i64) {
map.insert(txn, key.into(), Any::BigInt(value));
}
pub fn set_map_bool(
map: &MapRef,
txn: &mut TransactionMut,
key: impl Into<String>,
value: bool,
) {
map.insert(txn, key.into(), value);
}
pub fn get_or_create_nested_map(
parent: &MapRef,
txn: &mut TransactionMut,
key: impl Into<String>,
) -> MapRef {
let key_str = key.into();
if let Some(existing) = parent.get(txn, &key_str)
&& let Ok(m) = MapRef::try_from(existing)
{
return m;
}
let empty_prelim: MapPrelim = MapPrelim::from([("_", Any::Null)]);
let new_map: MapRef = parent.insert(txn, key_str.as_str(), empty_prelim);
new_map.remove(txn, "_");
new_map
}
pub fn map_contains_key(map: &MapRef, txn: &impl ReadTxn, key: &str) -> bool {
map.contains_key(txn, key)
}
pub async fn compact_tombstones(&self) -> CrdtResult<usize> {
let crdt_dir = self.storage_dir.join("crdt");
if !crdt_dir.exists() {
return Ok(0);
}
let mut removed_count = 0;
let retention_secs = self.config.tombstone_retention.as_secs() as i64;
let now = chrono::Utc::now().timestamp();
let mut read_dir = fs::read_dir(&crdt_dir)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read CRDT directory: {}", e)))?;
while let Some(entry) = read_dir
.next_entry()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read directory entry: {}", e)))?
{
if !entry
.file_type()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to get file type: {}", e)))?
.is_dir()
{
continue;
}
let entity_type_dir = entry.path();
let mut entity_dir = fs::read_dir(&entity_type_dir).await.map_err(|e| {
CrdtError::FileSystem(format!("Failed to read entity type directory: {}", e))
})?;
while let Some(file_entry) = entity_dir
.next_entry()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read file entry: {}", e)))?
{
let file_name = file_entry.file_name();
let file_name_str = file_name.to_string_lossy();
if !file_name_str.ends_with(".yrs") {
continue;
}
let hex_id = file_name_str.trim_end_matches(".yrs");
let doc_id = match hex::decode(hex_id)
.ok()
.and_then(|bytes| String::from_utf8(bytes).ok())
{
Some(id) => id,
None => {
tracing::warn!("Failed to decode doc_id from hex: {}", hex_id);
continue;
}
};
let doc = match self.load_document(&doc_id).await {
Ok(d) => d,
Err(e) => {
tracing::warn!("Failed to load document {}: {}", doc_id, e);
continue;
}
};
let root = doc.get_or_insert_map("root");
let txn = doc.transact();
let is_deleted = if let Some(metadata_val) = root.get(&txn, "metadata")
&& let Ok(metadata) = MapRef::try_from(metadata_val)
{
let deleted = Self::get_map_bool(&metadata, &txn, "deleted").unwrap_or(false);
let deleted_at = Self::get_map_i64(&metadata, &txn, "deleted_at");
if deleted && deleted_at.is_some() {
let age_secs = now - deleted_at.unwrap_or(now);
age_secs >= retention_secs
} else {
false
}
} else {
false
};
if is_deleted {
if let Err(e) = self.delete_document(&doc_id).await {
tracing::warn!("Failed to delete tombstoned document {}: {}", doc_id, e);
} else {
tracing::debug!("Compacted tombstone: {}", doc_id);
removed_count += 1;
}
}
}
}
Ok(removed_count)
}
pub async fn compact_document(&self, doc_id: &str, force: bool) -> CrdtResult<u64> {
let parts: Vec<&str> = doc_id.split(':').collect();
if parts.len() < 2 {
return Err(CrdtError::InvalidDocumentId(format!(
"Invalid doc_id format (expected 'entity_type:entity_id:...'): {}",
doc_id
)));
}
let entity_type = parts[0];
let entity_id = parts[1];
let (yrs_path, _) = self.doc_paths(entity_type, doc_id);
if !yrs_path.exists() {
return Err(CrdtError::DocumentNotFound(doc_id.to_string()));
}
let old_size = fs::metadata(&yrs_path)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to get file metadata: {}", e)))?
.len();
if !force && old_size < self.config.max_document_bytes as u64 {
return Ok(0);
}
let doc = self.load_document(doc_id).await?;
let new_state = {
const MAX_TXN_RETRIES: usize = 8;
const BASE_DELAY_MS: u64 = 5;
let mut attempt = 0usize;
loop {
match doc.try_transact() {
Ok(txn) => {
break txn.encode_state_as_update_v1(&yrs::StateVector::default());
}
Err(err) => {
if attempt >= MAX_TXN_RETRIES {
return Err(CrdtError::Operation(format!(
"Failed to acquire read transaction for {} after {} attempts: {}",
doc_id,
attempt + 1,
err
)));
}
let delay = BASE_DELAY_MS.saturating_mul((attempt + 1) as u64);
tokio::time::sleep(Duration::from_millis(delay)).await;
attempt += 1;
}
}
}
};
let new_size = new_state.len() as u64;
if new_size < old_size || force {
let new_doc = Doc::new();
{
let mut txn = new_doc.transact_mut();
let update = Update::decode_v1(&new_state).map_err(|e| {
CrdtError::Deserialization(format!("Failed to decode compacted state: {}", e))
})?;
txn.apply_update(update);
}
self.save_document(doc_id, entity_type, entity_id, &new_doc)
.await?;
let bytes_saved = old_size.saturating_sub(new_size);
tracing::debug!(
"Compacted document {}: {} bytes -> {} bytes (saved {})",
doc_id,
old_size,
new_size,
bytes_saved
);
Ok(bytes_saved)
} else {
Ok(0)
}
}
pub async fn compact_all(&self) -> CrdtResult<CompactionResult> {
let mut result = CompactionResult::default();
result.tombstones_removed = self.compact_tombstones().await?;
let crdt_dir = self.storage_dir.join("crdt");
if !crdt_dir.exists() {
return Ok(result);
}
let mut read_dir = fs::read_dir(&crdt_dir)
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read CRDT directory: {}", e)))?;
while let Some(entry) = read_dir
.next_entry()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read directory entry: {}", e)))?
{
if !entry
.file_type()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to get file type: {}", e)))?
.is_dir()
{
continue;
}
let entity_type_dir = entry.path();
let mut entity_dir = fs::read_dir(&entity_type_dir).await.map_err(|e| {
CrdtError::FileSystem(format!("Failed to read entity type directory: {}", e))
})?;
while let Some(file_entry) = entity_dir
.next_entry()
.await
.map_err(|e| CrdtError::FileSystem(format!("Failed to read file entry: {}", e)))?
{
let file_name = file_entry.file_name();
let file_name_str = file_name.to_string_lossy();
if !file_name_str.ends_with(".yrs") {
continue;
}
let hex_id = file_name_str.trim_end_matches(".yrs");
let doc_id = match hex::decode(hex_id)
.ok()
.and_then(|bytes| String::from_utf8(bytes).ok())
{
Some(id) => id,
None => {
tracing::warn!("Failed to decode doc_id from hex: {}", hex_id);
continue;
}
};
match self.compact_document(&doc_id, false).await {
Ok(bytes_saved) => {
if bytes_saved > 0 {
result.documents_compacted += 1;
result.bytes_saved += bytes_saved;
}
}
Err(e) => {
tracing::warn!("Failed to compact document {}: {}", doc_id, e);
}
}
}
}
tracing::debug!(
"Compaction complete: {} tombstones removed, {} documents compacted, {} bytes saved",
result.tombstones_removed,
result.documents_compacted,
result.bytes_saved
);
Ok(result)
}
pub async fn start_compaction_task(&self) {
let mut task_lock = self.compaction_task.write().await;
if task_lock.is_some() {
tracing::debug!("Compaction task already running");
return;
}
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel::<()>(1);
*self.compaction_shutdown_tx.write().await = Some(shutdown_tx);
let interval = self.config.compaction_interval;
let storage_dir = self.storage_dir.clone();
let config = self.config.clone();
let handle = tokio::spawn(async move {
tracing::debug!("Compaction task started with interval {:?}", interval);
loop {
tokio::select! {
_ = tokio::time::sleep(interval) => {
match CrdtManager::new_with_config(&storage_dir, config.clone()).await {
Ok(manager) => {
match manager.compact_all().await {
Ok(result) => {
tracing::debug!(
"Background compaction: {} tombstones, {} docs, {} bytes saved",
result.tombstones_removed,
result.documents_compacted,
result.bytes_saved
);
}
Err(e) => {
tracing::warn!("Background compaction failed: {}", e);
}
}
}
Err(e) => {
tracing::warn!("Failed to create manager for compaction: {}", e);
}
}
}
_ = shutdown_rx.recv() => {
tracing::debug!("Compaction task shutting down");
break;
}
}
}
});
*task_lock = Some(handle);
tracing::debug!("Compaction task spawned");
}
pub async fn stop_compaction_task(&self) {
let mut task_lock = self.compaction_task.write().await;
if let Some(handle) = task_lock.take() {
if let Some(tx) = self.compaction_shutdown_tx.write().await.take() {
let _ = tx.send(()).await;
}
let _ = handle.await;
tracing::debug!("Compaction task stopped");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_save_and_load_document() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Test Channel").expect("set name");
CrdtManager::set_map_value(&doc, "count", 42i64).expect("set count");
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let loaded_doc = manager
.load_document("channel:ch-1:doc")
.await
.expect("load document");
let name: String = CrdtManager::get_map_value(&loaded_doc, "name")
.expect("get value")
.expect("name exists");
assert_eq!(name, "Test Channel");
let count: i64 = CrdtManager::get_map_value(&loaded_doc, "count")
.expect("get value")
.expect("count exists");
assert_eq!(count, 42);
}
#[tokio::test]
async fn test_list_documents() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
for i in 1..=3 {
let doc = Doc::new();
manager
.save_document(
&format!("channel:ch-{}:doc", i),
"channel",
&format!("ch-{}", i),
&doc,
)
.await
.expect("save document");
}
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 3);
assert!(docs.contains(&"channel:ch-1:doc".to_string()));
assert!(docs.contains(&"channel:ch-2:doc".to_string()));
assert!(docs.contains(&"channel:ch-3:doc".to_string()));
}
#[tokio::test]
async fn test_delete_document() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 1);
manager
.delete_document("channel:ch-1:doc")
.await
.expect("delete document");
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 0);
}
#[tokio::test]
async fn test_document_not_found() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let result = manager.load_document("nonexistent").await;
assert!(result.is_err());
assert!(matches!(result, Err(CrdtError::InvalidDocumentId(_))));
let result = manager.load_document("channel:ch-1:doc").await;
assert!(result.is_err());
assert!(matches!(result, Err(CrdtError::DocumentNotFound(_))));
}
#[tokio::test]
async fn test_concurrent_saves() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
for i in 1..=5 {
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "version", i as i64).expect("set version");
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
}
let loaded_doc = manager
.load_document("channel:ch-1:doc")
.await
.expect("load document");
let version: i64 = CrdtManager::get_map_value(&loaded_doc, "version")
.expect("get value")
.expect("version exists");
assert_eq!(version, 5);
}
#[tokio::test]
async fn test_metadata_persistence() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let hex_doc_id = hex::encode("channel:ch-1:doc".as_bytes());
let meta_path = storage_path.join(format!("crdt/channel/{}.meta", hex_doc_id));
assert!(meta_path.exists());
let meta_json = tokio::fs::read_to_string(meta_path)
.await
.expect("read metadata");
let metadata: DocumentMetadata = serde_json::from_str(&meta_json).expect("parse metadata");
assert_eq!(metadata.doc_id, "channel:ch-1:doc");
assert_eq!(metadata.entity_type, "channel");
assert_eq!(metadata.entity_id, "ch-1");
}
#[tokio::test]
async fn test_apply_update() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Initial").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
let doc_peer = manager
.load_document("channel:ch-1:metadata")
.await
.expect("load document");
CrdtManager::set_map_value(&doc_peer, "name", "Updated").expect("set name");
let update_bytes = doc_peer
.transact()
.encode_state_as_update_v1(&yrs::StateVector::default());
manager
.apply_update("channel:ch-1:metadata", &update_bytes)
.await
.expect("apply update");
let loaded_doc = manager
.load_document("channel:ch-1:metadata")
.await
.expect("load document");
let name: String = CrdtManager::get_map_value(&loaded_doc, "name")
.expect("get value")
.expect("name exists");
assert_eq!(name, "Updated");
}
#[tokio::test]
async fn test_merge_updates() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc_initial = Doc::new();
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc_initial)
.await
.expect("save initial document");
let doc1 = Doc::new();
CrdtManager::set_map_value(&doc1, "field1", "value1").expect("set field1");
let update1 = doc1
.transact()
.encode_state_as_update_v1(&yrs::StateVector::default());
let doc2 = Doc::new();
CrdtManager::set_map_value(&doc2, "field2", "value2").expect("set field2");
let update2 = doc2
.transact()
.encode_state_as_update_v1(&yrs::StateVector::default());
let doc3 = Doc::new();
CrdtManager::set_map_value(&doc3, "field3", "value3").expect("set field3");
let update3 = doc3
.transact()
.encode_state_as_update_v1(&yrs::StateVector::default());
let merged_doc = manager
.merge_updates("channel:ch-1:metadata", vec![update1, update2, update3])
.await
.expect("merge updates");
let field1: String = CrdtManager::get_map_value(&merged_doc, "field1")
.expect("get value")
.expect("field1 exists");
let field2: String = CrdtManager::get_map_value(&merged_doc, "field2")
.expect("get value")
.expect("field2 exists");
let field3: String = CrdtManager::get_map_value(&merged_doc, "field3")
.expect("get value")
.expect("field3 exists");
assert_eq!(field1, "value1");
assert_eq!(field2, "value2");
assert_eq!(field3, "value3");
}
#[tokio::test]
async fn test_mark_deleted() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Test").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
manager
.mark_deleted("channel:ch-1:metadata", "deleter-id")
.await
.expect("mark deleted");
let is_deleted = manager
.is_deleted("channel:ch-1:metadata")
.await
.expect("check deleted");
assert!(is_deleted);
}
#[tokio::test]
async fn test_is_deleted_false_for_non_deleted() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Test").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
let is_deleted = manager
.is_deleted("channel:ch-1:metadata")
.await
.expect("check deleted");
assert!(!is_deleted);
}
#[tokio::test]
async fn test_doc_id_validation() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
let result = manager
.save_document("wrong:ch-1:doc", "channel", "ch-1", &doc)
.await;
assert!(matches!(result, Err(CrdtError::InvalidDocumentId(_))));
let result = manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_apply_update_creates_new_doc() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "field", "value").expect("set field");
let update_bytes = doc
.transact()
.encode_state_as_update_v1(&yrs::StateVector::default());
let result = manager
.apply_update("channel:ch-1:doc", &update_bytes)
.await;
assert!(
result.is_ok(),
"Expected success when applying update to non-existent document"
);
let loaded_doc = manager
.load_document("channel:ch-1:doc")
.await
.expect("load created doc");
let value = CrdtManager::get_map_value::<String>(&loaded_doc, "field")
.expect("get field")
.expect("field should exist");
assert_eq!(value, "value");
}
#[tokio::test]
async fn test_merge_updates_creates_new_doc() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "field", "value").expect("set field");
let update_bytes = doc
.transact()
.encode_state_as_update_v1(&yrs::StateVector::default());
let result = manager
.merge_updates("channel:ch-1:doc", vec![update_bytes])
.await;
assert!(
result.is_ok(),
"Expected success when merging updates to non-existent document"
);
let loaded_doc = manager
.load_document("channel:ch-1:doc")
.await
.expect("load created doc");
let value = CrdtManager::get_map_value::<String>(&loaded_doc, "field")
.expect("get field")
.expect("field should exist");
assert_eq!(value, "value");
}
#[tokio::test]
async fn test_update_size_limits() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let large_update = vec![0u8; 11 * 1024 * 1024];
let result = manager
.apply_update("channel:ch-1:doc", &large_update)
.await;
assert!(matches!(result, Err(CrdtError::Encoding(_))));
let result = manager
.merge_updates("channel:ch-1:doc", vec![large_update])
.await;
assert!(matches!(result, Err(CrdtError::Encoding(_))));
}
#[tokio::test]
async fn test_metadata_version_progression() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let hex_doc_id = hex::encode("channel:ch-1:doc".as_bytes());
let meta_path = storage_path.join(format!("crdt/channel/{}.meta", hex_doc_id));
let meta_json = tokio::fs::read_to_string(&meta_path)
.await
.expect("read metadata");
let metadata: DocumentMetadata = serde_json::from_str(&meta_json).expect("parse metadata");
assert_eq!(metadata.version, 1, "First save should be version 1");
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let meta_json = tokio::fs::read_to_string(&meta_path)
.await
.expect("read metadata");
let metadata: DocumentMetadata = serde_json::from_str(&meta_json).expect("parse metadata");
assert_eq!(metadata.version, 2, "Second save should be version 2");
manager
.save_document("channel:ch-1:doc", "channel", "ch-1", &doc)
.await
.expect("save document");
let meta_json = tokio::fs::read_to_string(&meta_path)
.await
.expect("read metadata");
let metadata: DocumentMetadata = serde_json::from_str(&meta_json).expect("parse metadata");
assert_eq!(metadata.version, 3, "Third save should be version 3");
}
#[test]
fn test_compaction_config_defaults() {
let config = CompactionConfig::default();
assert_eq!(config.tombstone_retention.as_secs(), 7 * 24 * 60 * 60);
assert_eq!(config.compaction_interval.as_secs(), 60 * 60);
assert_eq!(config.max_document_bytes, 5 * 1024 * 1024);
}
#[tokio::test]
async fn test_compact_tombstones_removes_old() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let config = CompactionConfig {
tombstone_retention: Duration::from_secs(0),
..Default::default()
};
let manager = CrdtManager::new_with_config(storage_path, config)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Test").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
manager
.mark_deleted("channel:ch-1:metadata", "deleter-id")
.await
.expect("mark deleted");
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 1);
let removed = manager
.compact_tombstones()
.await
.expect("compact tombstones");
assert_eq!(removed, 1);
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 0);
}
#[tokio::test]
async fn test_compact_tombstones_preserves_fresh() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let config = CompactionConfig {
tombstone_retention: Duration::from_secs(3600),
..Default::default()
};
let manager = CrdtManager::new_with_config(storage_path, config)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Test").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
manager
.mark_deleted("channel:ch-1:metadata", "deleter-id")
.await
.expect("mark deleted");
let removed = manager
.compact_tombstones()
.await
.expect("compact tombstones");
assert_eq!(removed, 0);
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 1);
}
#[tokio::test]
async fn test_compact_tombstones_ignores_non_deleted() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let config = CompactionConfig {
tombstone_retention: Duration::from_secs(0),
..Default::default()
};
let manager = CrdtManager::new_with_config(storage_path, config)
.await
.expect("create manager");
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", "Test").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
let removed = manager
.compact_tombstones()
.await
.expect("compact tombstones");
assert_eq!(removed, 0);
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 1);
}
#[tokio::test]
async fn test_compact_document_reduces_size() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let manager = CrdtManager::new(storage_path)
.await
.expect("create manager");
let doc = Doc::new();
for i in 0..100 {
CrdtManager::set_map_value(&doc, &format!("field_{}", i), format!("value_{}", i))
.expect("set field");
}
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
for i in 0..100 {
CrdtManager::set_map_value(&doc, &format!("field_{}", i), format!("updated_{}", i))
.expect("update field");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc)
.await
.expect("save document");
}
let bytes_saved = manager
.compact_document("channel:ch-1:metadata", true)
.await
.expect("compact document");
assert!(bytes_saved < u64::MAX);
let loaded_doc = manager
.load_document("channel:ch-1:metadata")
.await
.expect("load document");
let field_0: String = CrdtManager::get_map_value(&loaded_doc, "field_0")
.expect("get value")
.expect("field exists");
assert_eq!(field_0, "updated_0");
let field_99: String = CrdtManager::get_map_value(&loaded_doc, "field_99")
.expect("get value")
.expect("field exists");
assert_eq!(field_99, "updated_99");
}
#[tokio::test]
async fn test_compact_all_combined() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let config = CompactionConfig {
tombstone_retention: Duration::from_secs(0),
max_document_bytes: 0, ..Default::default()
};
let manager = CrdtManager::new_with_config(storage_path, config)
.await
.expect("create manager");
let doc1 = Doc::new();
CrdtManager::set_map_value(&doc1, "name", "Tombstone").expect("set name");
manager
.save_document("channel:ch-1:metadata", "channel", "ch-1", &doc1)
.await
.expect("save document");
manager
.mark_deleted("channel:ch-1:metadata", "deleter-id")
.await
.expect("mark deleted");
let doc2 = Doc::new();
CrdtManager::set_map_value(&doc2, "name", "Normal").expect("set name");
manager
.save_document("channel:ch-2:metadata", "channel", "ch-2", &doc2)
.await
.expect("save document");
let result = manager.compact_all().await.expect("compact all");
assert_eq!(result.tombstones_removed, 1);
assert!(result.documents_compacted < usize::MAX);
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 1);
assert!(docs.contains(&"channel:ch-2:metadata".to_string()));
}
#[tokio::test]
async fn test_start_stop_compaction_task() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let config = CompactionConfig {
compaction_interval: Duration::from_millis(100), ..Default::default()
};
let manager = CrdtManager::new_with_config(storage_path, config)
.await
.expect("create manager");
manager.start_compaction_task().await;
{
let task_lock = manager.compaction_task.read().await;
assert!(task_lock.is_some());
}
tokio::time::sleep(Duration::from_millis(200)).await;
manager.stop_compaction_task().await;
{
let task_lock = manager.compaction_task.read().await;
assert!(task_lock.is_none());
}
}
#[tokio::test]
async fn test_compaction_result_accumulates() {
let temp_dir = tempdir().expect("temp dir");
let storage_path = temp_dir.path();
let config = CompactionConfig {
tombstone_retention: Duration::from_secs(0),
max_document_bytes: 0,
..Default::default()
};
let manager = CrdtManager::new_with_config(storage_path, config)
.await
.expect("create manager");
for i in 1..=3 {
let doc = Doc::new();
CrdtManager::set_map_value(&doc, "name", format!("Doc {}", i)).expect("set name");
manager
.save_document(
&format!("channel:ch-{}:metadata", i),
"channel",
&format!("ch-{}", i),
&doc,
)
.await
.expect("save document");
manager
.mark_deleted(&format!("channel:ch-{}:metadata", i), "deleter-id")
.await
.expect("mark deleted");
}
let result = manager.compact_all().await.expect("compact all");
assert_eq!(result.tombstones_removed, 3);
assert_eq!(result.documents_compacted, 0);
let docs = manager
.list_documents("channel")
.await
.expect("list documents");
assert_eq!(docs.len(), 0);
}
}