fix: automatically remove the sender when the receiver is removed

src/watch/mod.rs

@@ -11,17 +11,22 @@ pub(crate) use filter::*;
 pub(crate) use request::*;
 pub(crate) use sender::*;
 
-use std::sync::{Arc, RwLock, TryLockError};
+use std::{
+    sync::{Arc, RwLock},
+    vec,
+};
+
+#[cfg(not(feature = "tokio"))]
+use std::sync::mpsc::SendError;
+#[cfg(feature = "tokio")]
+use tokio::sync::mpsc::error::SendError;
+
 use thiserror::Error;
 
 #[derive(Error, Debug)]
 pub enum WatchEventError {
-    #[error("TryLockErrorPoisoned")]
-    // TryLockErrorPoisoned(Batch<'a>), // TODO: remove 'a lifetime from Batch Error
-    TryLockErrorPoisoned,
-    #[error("TryLockErrorWouldBlock")]
-    // TryLockErrorWouldBlock(Batch<'a>), // TODO: remove 'a lifetime from Batch Error
-    TryLockErrorWouldBlock,
+    #[error("LockErrorPoisoned")]
+    LockErrorPoisoned,
     #[cfg(not(feature = "tokio"))]
     #[error("SendError")]
     SendError(#[from] std::sync::mpsc::SendError<Event>),
@@ -44,15 +49,28 @@ pub(crate) fn push_batch(
     senders: Arc<RwLock<Watchers>>,
     batch: Batch,
 ) -> Result<(), WatchEventError> {
-    let watchers = senders.try_read().map_err(|err| match err {
-        TryLockError::Poisoned(_) => WatchEventError::TryLockErrorPoisoned,
-        TryLockError::WouldBlock => WatchEventError::TryLockErrorWouldBlock,
+    let watchers = senders.read().map_err(|err| match err {
+        _ => WatchEventError::LockErrorPoisoned,
     })?;
 
+    let mut unused_watchers = vec![];
     for (watcher_request, event) in batch {
-        for sender in watchers.find_senders(&watcher_request) {
-            let sender = sender.lock().unwrap();
-            sender.send(event.clone())?;
+        for (id, sender) in watchers.find_senders(&watcher_request) {
+            let l_sender = sender.lock().unwrap();
+            if let Err(SendError(_)) = l_sender.send(event.clone()) {
+                println!("Failed to send event to watcher {}", id);
+                unused_watchers.push(id);
+            }
+        }
+    }
+    // Drop the lock before removing the watchers to avoid deadlock
+    drop(watchers);
+
+    // Remove unused watchers
+    if unused_watchers.len() > 0 {
+        let mut w = senders.write().unwrap();
+        for id in unused_watchers {
+            w.remove_sender(id);
         }
     }
 

src/watch/sender.rs

@@ -28,18 +28,18 @@ impl Watchers {
     pub(crate) fn find_senders(
         &self,
         request: &WatcherRequest,
-    ) -> Vec<Arc<Mutex<MpscSender<Event>>>> {
+    ) -> Vec<(u64, Arc<Mutex<MpscSender<Event>>>)> {
         let mut event_senders = Vec::new();
-        for (_, (filter, event_sender)) in &self.0 {
+        for (id, (filter, event_sender)) in &self.0 {
             if filter.table_name == request.table_name {
                 match &filter.key_filter {
                     KeyFilter::Primary(value) => {
                         if let Some(key) = &value {
                             if key == &request.primary_key {
-                                event_senders.push(Arc::clone(event_sender));
+                                event_senders.push((*id, Arc::clone(event_sender)));
                             }
                         } else {
-                            event_senders.push(Arc::clone(event_sender));
+                            event_senders.push((*id, Arc::clone(event_sender)));
                         }
                     }
                     KeyFilter::PrimaryStartWith(key_prefix) => {
@@ -48,7 +48,7 @@ impl Watchers {
                             .as_slice()
                             .starts_with(key_prefix.as_slice())
                         {
-                            event_senders.push(Arc::clone(event_sender));
+                            event_senders.push((*id, Arc::clone(event_sender)));
                         }
                     }
                     KeyFilter::Secondary(key_def, key) => {
@@ -60,19 +60,20 @@ impl Watchers {
                                     match request_secondary_key {
                                         DatabaseKeyValue::Default(value) => {
                                             if value == filter_value {
-                                                event_senders.push(Arc::clone(event_sender));
+                                                event_senders.push((*id, Arc::clone(event_sender)));
                                             }
                                         }
                                         DatabaseKeyValue::Optional(value) => {
                                             if let Some(value) = value {
                                                 if value == filter_value {
-                                                    event_senders.push(Arc::clone(event_sender));
+                                                    event_senders
+                                                        .push((*id, Arc::clone(event_sender)));
                                                 }
                                             }
                                         }
                                     }
                                 } else {
-                                    event_senders.push(Arc::clone(event_sender));
+                                    event_senders.push((*id, Arc::clone(event_sender)));
                                 }
                             }
                         }
@@ -85,15 +86,15 @@ impl Watchers {
                                 DatabaseKeyValue::Default(value) => {
                                     if key_def == request_secondary_key_def {
                                         if value.as_slice().starts_with(key_prefix.as_slice()) {
-                                            event_senders.push(Arc::clone(event_sender));
+                                            event_senders.push((*id, Arc::clone(event_sender)));
                                         }
                                     }
                                 }
                                 DatabaseKeyValue::Optional(value) => {
                                     if let Some(value) = value {
                                         if key_def == request_secondary_key_def {
                                             if value.as_slice().starts_with(key_prefix.as_slice()) {
-                                                event_senders.push(Arc::clone(event_sender));
+                                                event_senders.push((*id, Arc::clone(event_sender)));
                                             }
                                         }
                                     }

tests/watch/mod.rs

@@ -292,6 +292,42 @@ fn unwatch() {
     assert!(recv.try_recv().is_err());
 }
 
+#[test]
+fn unwatch_by_deleted_receiver() {
+    let tf = TmpFs::new().unwrap();
+
+    let mut builder = DatabaseBuilder::new();
+    builder.define::<ItemA>().unwrap();
+    let db = builder.create(tf.path("test").as_std_path()).unwrap();
+
+    let item_a = ItemA { id: 1 };
+
+    let (recv, recv_id) = db.watch().get().primary::<ItemA>(item_a.id).unwrap();
+
+    let rw = db.rw_transaction().unwrap();
+    rw.insert(item_a.clone()).unwrap();
+    rw.commit().unwrap();
+
+    for _ in 0..1 {
+        let inner_event: ItemA = if let Event::Insert(event) = recv.recv_timeout(TIMEOUT).unwrap() {
+            event.inner()
+        } else {
+            panic!("wrong event")
+        };
+        assert_eq!(inner_event, item_a);
+    }
+
+    drop(recv);
+
+    let rw = db.rw_transaction().unwrap();
+    rw.insert(item_a.clone()).unwrap();
+    // The watcher is removed during the commit because the receiver is dropped
+    rw.commit().unwrap();
+
+    // Check if the watcher is removed
+    assert!(!db.unwatch(recv_id).unwrap());
+}
+
 #[derive(Serialize, Deserialize, Eq, PartialEq, Debug, Clone)]
 #[native_model(id = 4, version = 1)]
 #[native_db]