use std::collections::HashMap;
use std::sync::{Arc, Mutex};
#[derive(Debug, Clone)]
pub enum DocumentMessage {
Json(String),
Bytes(Vec<u8>),
Custom { type_id: String, payload: Vec<u8> },
}
pub trait DocumentMessaging: Send + Sync + 'static {
fn post_message(&self, source_doc: usize, target_doc: usize, message: DocumentMessage);
fn register_handler(
&self,
doc_id: usize,
handler: Box<dyn Fn(usize, DocumentMessage) + Send + Sync>,
);
fn unregister_handler(&self, doc_id: usize);
fn has_handler(&self, doc_id: usize) -> bool;
fn handler_count(&self) -> usize;
}
pub struct InMemoryDocumentMessaging {
handlers: Mutex<HashMap<usize, Box<dyn Fn(usize, DocumentMessage) + Send + Sync>>>,
}
impl InMemoryDocumentMessaging {
pub fn new() -> Self {
Self {
handlers: Mutex::new(HashMap::new()),
}
}
pub fn new_shared() -> Arc<Self> {
Arc::new(Self::new())
}
}
impl Default for InMemoryDocumentMessaging {
fn default() -> Self {
Self::new()
}
}
impl DocumentMessaging for InMemoryDocumentMessaging {
fn post_message(&self, source_doc: usize, target_doc: usize, message: DocumentMessage) {
let handlers = self.handlers.lock().unwrap_or_else(|e| e.into_inner());
if let Some(handler) = handlers.get(&target_doc) {
handler(source_doc, message);
}
}
fn register_handler(
&self,
doc_id: usize,
handler: Box<dyn Fn(usize, DocumentMessage) + Send + Sync>,
) {
let mut handlers = self.handlers.lock().unwrap_or_else(|e| e.into_inner());
handlers.insert(doc_id, handler);
}
fn unregister_handler(&self, doc_id: usize) {
let mut handlers = self.handlers.lock().unwrap_or_else(|e| e.into_inner());
handlers.remove(&doc_id);
}
fn has_handler(&self, doc_id: usize) -> bool {
let handlers = self.handlers.lock().unwrap_or_else(|e| e.into_inner());
handlers.contains_key(&doc_id)
}
fn handler_count(&self) -> usize {
let handlers = self.handlers.lock().unwrap_or_else(|e| e.into_inner());
handlers.len()
}
}
pub struct FilteringDocumentMessaging<F> {
inner: Arc<dyn DocumentMessaging>,
filter: F,
}
impl<F> FilteringDocumentMessaging<F>
where
F: Fn(usize, usize, &DocumentMessage) -> bool + Send + Sync + 'static,
{
pub fn new(inner: Arc<dyn DocumentMessaging>, filter: F) -> Self {
Self { inner, filter }
}
}
impl<F> DocumentMessaging for FilteringDocumentMessaging<F>
where
F: Fn(usize, usize, &DocumentMessage) -> bool + Send + Sync + 'static,
{
fn post_message(&self, source_doc: usize, target_doc: usize, message: DocumentMessage) {
if (self.filter)(source_doc, target_doc, &message) {
self.inner.post_message(source_doc, target_doc, message);
}
}
fn register_handler(
&self,
doc_id: usize,
handler: Box<dyn Fn(usize, DocumentMessage) + Send + Sync>,
) {
self.inner.register_handler(doc_id, handler);
}
fn unregister_handler(&self, doc_id: usize) {
self.inner.unregister_handler(doc_id);
}
fn has_handler(&self, doc_id: usize) -> bool {
self.inner.has_handler(doc_id)
}
fn handler_count(&self) -> usize {
self.inner.handler_count()
}
}
pub fn create_same_origin_filter(inner: Arc<dyn DocumentMessaging>) -> Arc<dyn DocumentMessaging> {
Arc::new(FilteringDocumentMessaging::new(
inner,
|_source, _target, _msg| {
true },
))
}
pub type SharedDocumentMessaging = Arc<dyn DocumentMessaging>;
pub fn create_default_messaging() -> SharedDocumentMessaging {
InMemoryDocumentMessaging::new_shared()
}
#[derive(Debug, Clone)]
pub struct MessageEnvelope {
pub source: usize,
pub target: usize,
pub timestamp: std::time::SystemTime,
pub message: DocumentMessage,
}
pub mod message_helpers {
use super::*;
pub fn bytes(data: Vec<u8>) -> DocumentMessage {
DocumentMessage::Bytes(data)
}
pub fn text(text: impl Into<String>) -> DocumentMessage {
DocumentMessage::Json(format!("\"{}\"", text.into()))
}
pub fn event(event_type: impl Into<String>) -> DocumentMessage {
DocumentMessage::Json(format!("{{\"type\":\"{}\"}}", event_type.into()))
}
pub fn json(json_str: impl Into<String>) -> DocumentMessage {
DocumentMessage::Json(json_str.into())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_in_memory_messaging() {
let messaging = InMemoryDocumentMessaging::new();
let received = Arc::new(Mutex::new(None));
let received_clone = received.clone();
messaging.register_handler(
2,
Box::new(move |source, msg| {
*received_clone.lock().unwrap_or_else(|e| e.into_inner()) = Some((source, msg));
}),
);
let message = DocumentMessage::Json(r#"{"hello": "world"}"#.to_string());
messaging.post_message(1, 2, message.clone());
let result = received.lock().unwrap_or_else(|e| e.into_inner()).clone();
assert!(result.is_some());
let (source, received_msg) = result.unwrap();
assert_eq!(source, 1);
match received_msg {
DocumentMessage::Json(json) => {
assert_eq!(json, r#"{"hello": "world"}"#);
}
_ => panic!("Expected JSON message"),
}
}
#[test]
fn test_unregister_handler() {
let messaging = InMemoryDocumentMessaging::new();
messaging.register_handler(1, Box::new(|_, _| {}));
assert!(messaging.has_handler(1));
messaging.unregister_handler(1);
assert!(!messaging.has_handler(1));
}
#[test]
fn test_message_to_nonexistent_handler() {
let messaging = InMemoryDocumentMessaging::new();
let message = DocumentMessage::Json(r#"{"test": true}"#.to_string());
messaging.post_message(1, 999, message);
assert_eq!(messaging.handler_count(), 0);
}
#[test]
fn test_handler_count() {
let messaging = InMemoryDocumentMessaging::new();
assert_eq!(messaging.handler_count(), 0);
messaging.register_handler(1, Box::new(|_, _| {}));
assert_eq!(messaging.handler_count(), 1);
messaging.register_handler(2, Box::new(|_, _| {}));
assert_eq!(messaging.handler_count(), 2);
messaging.unregister_handler(1);
assert_eq!(messaging.handler_count(), 1);
}
#[test]
fn test_filtering_messaging() {
let inner = create_default_messaging();
let filter = FilteringDocumentMessaging::new(inner.clone(), |_, target, _| target != 3);
let received_doc2 = Arc::new(Mutex::new(false));
let received_doc2_clone = received_doc2.clone();
let received_doc3 = Arc::new(Mutex::new(false));
let received_doc3_clone = received_doc3.clone();
inner.register_handler(
2,
Box::new(move |_, _| {
*received_doc2_clone.lock().unwrap_or_else(|e| e.into_inner()) = true;
}),
);
inner.register_handler(
3,
Box::new(move |_, _| {
*received_doc3_clone.lock().unwrap_or_else(|e| e.into_inner()) = true;
}),
);
let message = DocumentMessage::Json(r#"{}"#.to_string());
filter.post_message(1, 2, message.clone());
filter.post_message(1, 3, message);
assert!(*received_doc2.lock().unwrap_or_else(|e| e.into_inner()));
assert!(!*received_doc3.lock().unwrap_or_else(|e| e.into_inner())); }
#[test]
fn test_message_helpers() {
let msg = message_helpers::text("hello");
match msg {
DocumentMessage::Json(json) => assert!(json.contains("hello")),
_ => panic!("Expected JSON"),
}
let msg = message_helpers::bytes(vec![1, 2, 3]);
match msg {
DocumentMessage::Bytes(bytes) => assert_eq!(bytes, vec![1, 2, 3]),
_ => panic!("Expected Bytes"),
}
let msg = message_helpers::event("click");
match msg {
DocumentMessage::Json(json) => assert!(json.contains("click")),
_ => panic!("Expected JSON"),
}
let msg = message_helpers::json(r#"{"custom": true}"#);
match msg {
DocumentMessage::Json(json) => assert_eq!(json, r#"{"custom": true}"#),
_ => panic!("Expected JSON"),
}
}
}