use crate::crdt::*;
use crate::error::AppResult;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
#[derive(Debug)]
pub struct MessageSyncService {
peer_id: String,
entity_clocks: Arc<RwLock<HashMap<String, VectorClock>>>,
entity_messages: Arc<RwLock<HashMap<String, Vec<CRDTMessage>>>>,
pending_messages: Arc<RwLock<HashMap<String, Vec<CRDTMessage>>>>,
lamport_clock: Arc<RwLock<u64>>,
}
impl MessageSyncService {
pub fn new(peer_id: String) -> Self {
info!("🔄 MessageSyncService initialized for peer: {}", peer_id);
Self {
peer_id,
entity_clocks: Arc::new(RwLock::new(HashMap::new())),
entity_messages: Arc::new(RwLock::new(HashMap::new())),
pending_messages: Arc::new(RwLock::new(HashMap::new())),
lamport_clock: Arc::new(RwLock::new(0)),
}
}
pub async fn get_all_messages(&self, entity_id: &str) -> AppResult<SyncResponse> {
let messages_map = self.entity_messages.read().await;
let clocks_map = self.entity_clocks.read().await;
let mut messages = messages_map.get(entity_id).cloned().unwrap_or_default();
let vector_clock = clocks_map.get(entity_id).cloned().unwrap_or_default();
sort_messages_causally(&mut messages);
info!(
"📤 get_all_messages for {}: {} messages",
entity_id,
messages.len()
);
Ok(SyncResponse {
entity_id: entity_id.to_string(),
entity_type: self.infer_entity_type(entity_id),
messages,
vector_clock,
})
}
pub async fn receive_message(&self, message: CRDTMessage) -> AppResult<ReceiveResult> {
let entity_id = message.metadata.entity_id.clone();
let clocks_map = self.entity_clocks.read().await;
let local_clock = clocks_map.get(&entity_id).cloned().unwrap_or_default();
drop(clocks_map);
let has_deps = local_clock.has_dependencies(&message.metadata.vector_clock);
if !has_deps {
warn!("⚠️ Out-of-order message detected: {}", message.metadata.id);
let mut pending_map = self.pending_messages.write().await;
let pending = pending_map.entry(entity_id.clone()).or_default();
pending.push(message.clone());
let missing = local_clock.get_missing_ranges(&message.metadata.vector_clock);
return Ok(ReceiveResult {
accepted: false,
out_of_order: true,
missing_ranges: Some(missing),
});
}
self.add_message(message).await?;
self.process_pending_messages(&entity_id).await?;
Ok(ReceiveResult {
accepted: true,
out_of_order: false,
missing_ranges: None,
})
}
pub async fn send_message(
&self,
entity_id: String,
entity_type: EntityType,
content: MessageContent,
reply_to_id: Option<String>,
) -> AppResult<CRDTMessage> {
let mut clocks_map = self.entity_clocks.write().await;
let clock = clocks_map.entry(entity_id.clone()).or_default();
clock.increment(&self.peer_id);
let new_clock = clock.clone();
drop(clocks_map);
let mut lamport = self.lamport_clock.write().await;
*lamport += 1;
let lamport_value = *lamport;
drop(lamport);
let messages_map = self.entity_messages.read().await;
let previous_id = messages_map
.get(&entity_id)
.and_then(|msgs| msgs.last())
.map(|msg| msg.metadata.id.clone());
drop(messages_map);
let metadata = MessageMetadata {
id: format!(
"{}-{}-{}",
self.peer_id,
new_clock.0.get(&self.peer_id).copied().unwrap_or(0),
chrono::Utc::now().timestamp_millis()
),
entity_id: entity_id.clone(),
entity_type,
author_peer_id: self.peer_id.clone(),
vector_clock: new_clock,
lamport_clock: lamport_value,
timestamp: chrono::Utc::now().timestamp_millis() as u64,
previous_message_id: previous_id,
reply_to_id,
};
let message = CRDTMessage {
content,
metadata,
local_state: Some(LocalMessageState {
status: Some(MessageStatus::Sent),
reactions: Vec::new(),
edited_at: None,
thread_count: None,
latest_reply_by: None,
}),
};
self.add_message(message.clone()).await?;
Ok(message)
}
pub async fn request_sync(
&self,
entity_id: &str,
from_peer_id: &str,
) -> AppResult<SyncRequest> {
let clocks_map = self.entity_clocks.read().await;
let local_clock = clocks_map.get(entity_id).cloned().unwrap_or_default();
drop(clocks_map);
let pending_map = self.pending_messages.read().await;
let missing_ids = pending_map
.get(entity_id)
.map(|pending| pending.iter().map(|m| m.metadata.id.clone()).collect());
drop(pending_map);
debug!("🔄 Requesting sync for {} from {}", entity_id, from_peer_id);
debug!(" Local clock: {:?}", local_clock);
debug!(" Missing messages: {:?}", missing_ids);
Ok(SyncRequest {
entity_id: entity_id.to_string(),
entity_type: self.infer_entity_type(entity_id),
requester_peer_id: self.peer_id.clone(),
vector_clock: local_clock,
missing_message_ids: missing_ids,
request_id: chrono::Utc::now().timestamp_millis() as u64,
})
}
pub async fn handle_sync_response(&self, response: SyncResponse) -> AppResult<SyncResult> {
let entity_id = &response.entity_id;
let mut added = 0;
let mut rejected = 0;
info!(
"📥 Handling sync response for {}: {} messages",
entity_id,
response.messages.len()
);
for message in response.messages {
let result = self.receive_message(message).await?;
if result.accepted {
added += 1;
} else {
rejected += 1;
}
}
let mut clocks_map = self.entity_clocks.write().await;
let local_clock = clocks_map.entry(entity_id.clone()).or_default();
local_clock.merge(&response.vector_clock);
let merged_clock = local_clock.clone();
drop(clocks_map);
info!("✅ Sync complete: {} added, {} rejected", added, rejected);
debug!(" Updated clock: {:?}", merged_clock);
Ok(SyncResult {
messages_added: added,
messages_rejected: rejected,
})
}
pub async fn get_sync_state(&self, entity_id: &str) -> AppResult<EntitySyncState> {
let messages_map = self.entity_messages.read().await;
let pending_map = self.pending_messages.read().await;
let clocks_map = self.entity_clocks.read().await;
let messages = messages_map.get(entity_id).cloned().unwrap_or_default();
let pending = pending_map.get(entity_id).cloned().unwrap_or_default();
let clock = clocks_map.get(entity_id).cloned().unwrap_or_default();
Ok(EntitySyncState {
entity_id: entity_id.to_string(),
entity_type: self.infer_entity_type(entity_id),
vector_clock: clock,
last_sync_time: chrono::Utc::now().timestamp_millis() as u64,
message_count: messages.len(),
missing_messages: pending.iter().map(|m| m.metadata.id.clone()).collect(),
out_of_order_messages: pending.iter().map(|m| m.metadata.id.clone()).collect(),
})
}
pub async fn get_messages(&self, entity_id: &str) -> AppResult<Vec<CRDTMessage>> {
let messages_map = self.entity_messages.read().await;
let mut messages = messages_map.get(entity_id).cloned().unwrap_or_default();
sort_messages_causally(&mut messages);
Ok(messages)
}
pub async fn needs_sync(&self, entity_id: &str, remote_clock: &VectorClock) -> bool {
let clocks_map = self.entity_clocks.read().await;
let local_clock = clocks_map.get(entity_id).cloned().unwrap_or_default();
let missing = local_clock.get_missing_ranges(remote_clock);
!missing.is_empty()
}
pub async fn delete_message(&self, entity_id: &str, message_id: &str) -> AppResult<bool> {
let mut messages_map = self.entity_messages.write().await;
if let Some(messages) = messages_map.get_mut(entity_id) {
let original_len = messages.len();
messages.retain(|m| m.metadata.id != message_id);
let deleted = messages.len() < original_len;
if deleted {
info!("🗑️ Message deleted: {} (entity: {})", message_id, entity_id);
}
return Ok(deleted);
}
Ok(false)
}
pub async fn edit_message(
&self,
entity_id: &str,
message_id: &str,
new_text: String,
) -> AppResult<u64> {
let mut messages_map = self.entity_messages.write().await;
if let Some(messages) = messages_map.get_mut(entity_id) {
for message in messages.iter_mut() {
if message.metadata.id == message_id {
message.content.text = new_text;
let edited_at = chrono::Utc::now().timestamp_millis() as u64;
let local_state =
message
.local_state
.get_or_insert_with(|| LocalMessageState {
status: None,
reactions: Vec::new(),
edited_at: None,
thread_count: None,
latest_reply_by: None,
});
local_state.edited_at = Some(edited_at);
info!("✏️ Message edited: {} (entity: {})", message_id, entity_id);
return Ok(edited_at);
}
}
}
Err(crate::error::AppError::NotFound(format!(
"Message not found: {}",
message_id
)))
}
pub async fn add_reaction(
&self,
entity_id: &str,
message_id: &str,
emoji: String,
peer_id: String,
) -> AppResult<()> {
let mut messages_map = self.entity_messages.write().await;
if let Some(messages) = messages_map.get_mut(entity_id) {
for message in messages.iter_mut() {
if message.metadata.id == message_id {
let local_state =
message
.local_state
.get_or_insert_with(|| LocalMessageState {
status: None,
reactions: Vec::new(),
edited_at: None,
thread_count: None,
latest_reply_by: None,
});
if let Some(reaction) =
local_state.reactions.iter_mut().find(|r| r.emoji == emoji)
{
if !reaction.peer_ids.contains(&peer_id) {
reaction.peer_ids.push(peer_id.clone());
reaction.count += 1;
}
} else {
local_state.reactions.push(Reaction {
emoji: emoji.clone(),
count: 1,
user_reacted: Some(true),
peer_ids: vec![peer_id.clone()],
});
}
info!(
"👍 Reaction added: {} to {} (entity: {})",
emoji, message_id, entity_id
);
return Ok(());
}
}
}
Err(crate::error::AppError::NotFound(format!(
"Message not found: {}",
message_id
)))
}
pub async fn remove_reaction(
&self,
entity_id: &str,
message_id: &str,
emoji: String,
peer_id: String,
) -> AppResult<()> {
let mut messages_map = self.entity_messages.write().await;
if let Some(messages) = messages_map.get_mut(entity_id) {
for message in messages.iter_mut() {
if message.metadata.id == message_id {
if let Some(ref mut local_state) = message.local_state
&& let Some(reaction) =
local_state.reactions.iter_mut().find(|r| r.emoji == emoji)
{
reaction.peer_ids.retain(|p| p != &peer_id);
reaction.count = reaction.count.saturating_sub(1);
if reaction.count == 0 {
local_state.reactions.retain(|r| r.emoji != emoji);
}
info!(
"👎 Reaction removed: {} from {} (entity: {})",
emoji, message_id, entity_id
);
return Ok(());
}
return Err(crate::error::AppError::NotFound(format!(
"Reaction not found: {} on {}",
emoji, message_id
)));
}
}
}
Err(crate::error::AppError::NotFound(format!(
"Message not found: {}",
message_id
)))
}
pub async fn get_reactions(
&self,
entity_id: &str,
message_id: &str,
) -> AppResult<Vec<Reaction>> {
let messages_map = self.entity_messages.read().await;
if let Some(messages) = messages_map.get(entity_id) {
for message in messages.iter() {
if message.metadata.id == message_id {
return Ok(message
.local_state
.as_ref()
.map(|ls| ls.reactions.clone())
.unwrap_or_default());
}
}
}
Err(crate::error::AppError::NotFound(format!(
"Message not found: {}",
message_id
)))
}
async fn add_message(&self, message: CRDTMessage) -> AppResult<()> {
let entity_id = &message.metadata.entity_id;
let mut messages_map = self.entity_messages.write().await;
let messages = messages_map.entry(entity_id.clone()).or_default();
if messages
.iter()
.any(|m| m.metadata.id == message.metadata.id)
{
warn!("⚠️ Duplicate message ignored: {}", message.metadata.id);
return Ok(());
}
messages.push(message.clone());
drop(messages_map);
let mut clocks_map = self.entity_clocks.write().await;
let local_clock = clocks_map.entry(entity_id.clone()).or_default();
local_clock.merge(&message.metadata.vector_clock);
drop(clocks_map);
let mut lamport = self.lamport_clock.write().await;
*lamport = (*lamport).max(message.metadata.lamport_clock) + 1;
drop(lamport);
info!(
"📨 Message added: {} (entity: {})",
message.metadata.id, entity_id
);
Ok(())
}
async fn process_pending_messages(&self, entity_id: &str) -> AppResult<()> {
let mut pending_map = self.pending_messages.write().await;
let pending_messages = match pending_map.get_mut(entity_id) {
Some(p) if !p.is_empty() => std::mem::take(p),
_ => return Ok(()),
};
drop(pending_map);
let clocks_map = self.entity_clocks.read().await;
let local_clock = clocks_map.get(entity_id).cloned().unwrap_or_default();
drop(clocks_map);
let mut still_pending = Vec::new();
for message in pending_messages {
if local_clock.has_dependencies(&message.metadata.vector_clock) {
info!(
"✅ Pending message now has dependencies: {}",
message.metadata.id
);
self.add_message(message).await?;
} else {
still_pending.push(message);
}
}
let mut pending_map = self.pending_messages.write().await;
if !still_pending.is_empty() {
info!(
"⏳ Still pending: {} messages for {}",
still_pending.len(),
entity_id
);
pending_map.insert(entity_id.to_string(), still_pending);
} else {
pending_map.remove(entity_id);
info!("✨ All pending messages processed for {}", entity_id);
}
Ok(())
}
fn infer_entity_type(&self, entity_id: &str) -> EntityType {
if entity_id.starts_with("contact-")
|| entity_id.starts_with("ben-")
|| entity_id.starts_with("lauren")
{
EntityType::Person
} else if entity_id.contains("-org") {
EntityType::Organisation
} else if entity_id.starts_with("project-") {
EntityType::Project
} else if entity_id.contains("general") || entity_id.contains("channel") {
EntityType::Channel
} else {
EntityType::Group
}
}
}
#[derive(Debug)]
pub struct ReceiveResult {
pub accepted: bool,
pub out_of_order: bool,
pub missing_ranges: Option<Vec<MissingRange>>,
}
#[derive(Debug)]
pub struct SyncResult {
pub messages_added: usize,
pub messages_rejected: usize,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::crdt::VectorClock;
#[tokio::test]
async fn test_message_sync_service_creation() {
let service = MessageSyncService::new("test-peer".to_string());
let state = service.get_sync_state("test-entity").await.unwrap();
assert_eq!(state.entity_id, "test-entity");
assert_eq!(state.message_count, 0);
assert!(state.missing_messages.is_empty());
}
#[tokio::test]
async fn test_get_all_messages_empty() {
let service = MessageSyncService::new("test-peer".to_string());
let response = service.get_all_messages("unknown-entity").await.unwrap();
assert_eq!(response.entity_id, "unknown-entity");
assert!(response.messages.is_empty());
}
#[tokio::test]
async fn test_get_messages_empty() {
let service = MessageSyncService::new("test-peer".to_string());
let messages = service.get_messages("unknown-entity").await.unwrap();
assert!(messages.is_empty());
}
#[tokio::test]
async fn test_needs_sync_initially() {
let service = MessageSyncService::new("test-peer".to_string());
let remote_clock = VectorClock::default();
let needs = service.needs_sync("test-entity", &remote_clock).await;
assert!(!needs);
}
#[tokio::test]
async fn test_delete_message_nonexistent() {
let service = MessageSyncService::new("test-peer".to_string());
let deleted = service
.delete_message("test-entity", "nonexistent")
.await
.unwrap();
assert!(!deleted);
}
#[tokio::test]
async fn test_edit_message_nonexistent() {
let service = MessageSyncService::new("test-peer".to_string());
let result = service
.edit_message("test-entity", "nonexistent", "new text".to_string())
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_add_reaction_nonexistent() {
let service = MessageSyncService::new("test-peer".to_string());
let result = service
.add_reaction(
"test-entity",
"nonexistent",
"thumbsup".to_string(),
"peer".to_string(),
)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_remove_reaction_nonexistent() {
let service = MessageSyncService::new("test-peer".to_string());
let result = service
.remove_reaction(
"test-entity",
"nonexistent",
"thumbsup".to_string(),
"peer".to_string(),
)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_get_reactions_nonexistent() {
let service = MessageSyncService::new("test-peer".to_string());
let result = service.get_reactions("test-entity", "nonexistent").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_request_sync() {
let service = MessageSyncService::new("test-peer".to_string());
let request = service
.request_sync("test-entity", "remote-peer")
.await
.unwrap();
assert_eq!(request.entity_id, "test-entity");
assert_eq!(request.requester_peer_id, "test-peer");
}
#[test]
fn test_receive_result_structure() {
let result = ReceiveResult {
accepted: true,
out_of_order: false,
missing_ranges: None,
};
assert!(result.accepted);
assert!(!result.out_of_order);
}
#[test]
fn test_sync_result_structure() {
let result = SyncResult {
messages_added: 5,
messages_rejected: 2,
};
assert_eq!(result.messages_added, 5);
assert_eq!(result.messages_rejected, 2);
}
}