Use single threaded version of tokio::sync::watch

Since we are guaranteed to be single threaded within the handler fns for a
single hash, we can use a much more lightweight watcher impl.

Also, tokio::sync::watch is optimized for the case where there are a lot
of parallel watchers, using a giant BigNotify internally. We want to
optimize for the case where there are just 0 or a few watchers, so we are
not that much concerned with the thundering herd problem.

Author: rklaehn

src/store/fs/bao_file.rs

@@ -20,7 +20,7 @@ use bao_tree::{
 use bytes::{Bytes, BytesMut};
 use derive_more::Debug;
 use irpc::channel::mpsc;
-use tokio::sync::watch;
+use super::util::watch;
 use tracing::{debug, error, info, trace};
 
 use super::{

src/store/fs/util.rs

@@ -2,7 +2,7 @@ use std::future::Future;
 
 use tokio::{select, sync::mpsc};
 pub(crate) mod entity_manager;
-pub(crate) mod watcher;
+pub(crate) mod watch;
 
 /// A wrapper for a tokio mpsc receiver that allows peeking at the next message.
 #[derive(Debug)]

src/store/fs/util/watch.rs

@@ -0,0 +1,87 @@
+use std::{ops::Deref, sync::Arc};
+
+use atomic_refcell::{AtomicRef, AtomicRefCell};
+
+struct State<T> {
+    value: T,
+    dropped: bool,
+}
+
+struct Shared<T> {
+    value: AtomicRefCell<State<T>>,
+    notify: tokio::sync::Notify,
+}
+
+pub struct Sender<T>(Arc<Shared<T>>);
+
+pub struct Receiver<T>(Arc<Shared<T>>);
+
+impl<T> Sender<T> {
+    pub fn new(value: T) -> Self {
+        Self(Arc::new(Shared {
+            value: AtomicRefCell::new(State {
+                value,
+                dropped: false,
+            }),
+            notify: tokio::sync::Notify::new(),
+        }))
+    }
+
+    pub fn send_if_modified<F>(&self, modify: F) -> bool
+    where
+        F: FnOnce(&mut T) -> bool,
+    {
+        let mut state = self.0.value.borrow_mut();
+        let modified = modify(&mut state.value);
+        if modified {
+            self.0.notify.notify_waiters();
+        }
+        modified
+    }
+
+    pub fn borrow(&self) -> impl Deref<Target = T> + '_ {
+        AtomicRef::map(self.0.value.borrow(), |state| &state.value)
+    }
+
+    pub fn subscribe(&self) -> Receiver<T> {
+        Receiver(self.0.clone())
+    }
+}
+
+impl<T> Drop for Sender<T> {
+    fn drop(&mut self) {
+        self.0.value.borrow_mut().dropped = true;
+        self.0.notify.notify_waiters();
+    }
+}   
+
+impl<T> Receiver<T> {
+    pub async fn changed(&self) -> Result<(), error::RecvError> {
+        self.0.notify.notified().await;
+        if self.0.value.borrow().dropped {
+            Err(error::RecvError(()))
+        } else {
+            Ok(())
+        }
+    }
+
+    pub fn borrow(&self) -> impl Deref<Target = T> + '_ {
+        AtomicRef::map(self.0.value.borrow(), |state| &state.value)
+    }
+}
+    
+pub mod error {
+    use std::{error::Error, fmt};
+
+    /// Error produced when receiving a change notification.
+    #[derive(Debug, Clone)]
+    pub struct RecvError(pub(super) ());
+
+    impl fmt::Display for RecvError {
+        fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+            write!(fmt, "channel closed")
+        }
+    }
+
+    impl Error for RecvError {}
+}