Add version of tokio::sync::Watcher that only works locally and copy the relevant API

Author: rklaehn

Cargo.toml

@@ -44,6 +44,7 @@ iroh-base = "0.90"
 reflink-copy = "0.1.24"
 irpc = { version = "0.5.0", features = ["rpc", "quinn_endpoint_setup", "message_spans", "stream", "derive"], default-features = false }
 iroh-metrics = { version = "0.35" }
+atomic_refcell = "0.1.13"
 
 [dev-dependencies]
 clap = { version = "4.5.31", features = ["derive"] }
@@ -58,7 +59,6 @@ testresult = "0.4.1"
 tracing-subscriber = { version = "0.3.19", features = ["fmt"] }
 tracing-test = "0.2.5"
 walkdir = "2.5.0"
-atomic_refcell = "0.1.13"
 
 [features]
 hide-proto-docs = []

src/store/fs.rs

@@ -106,7 +106,7 @@ use crate::{
         ApiClient,
     },
     store::{
-        fs::util::entity_manager::{self, ActiveEntityState},
+        fs::util::entity_manager::{self, ActiveEntityState, ShutdownCause},
         util::{BaoTreeSender, FixedSize, MemOrFile, ValueOrPoisioned},
         Hash,
     },
@@ -217,10 +217,17 @@ impl entity_manager::Params for EmParams {
 
     type EntityState = Slot;
 
-    async fn on_shutdown(
-        _state: entity_manager::ActiveEntityState<Self>,
-        _cause: entity_manager::ShutdownCause,
-    ) {
+    async fn on_shutdown(state: HashContext, cause: ShutdownCause) {
+        // this isn't strictly necessary. Drop will run anyway as soon as the
+        // state is reset to it's default value. Doing it here means that we
+        // have exact control over where it happens.
+        if let Some(handle) = state.state.0.lock().await.take() {
+            trace!(
+                "shutting down entity manager for hash: {}, cause: {cause:?}",
+                state.id
+            );
+            drop(handle);
+        }
     }
 }
 

src/store/fs/util.rs

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

src/store/fs/util/watcher.rs

@@ -0,0 +1,51 @@
+use std::{ops::Deref, sync::Arc};
+
+use atomic_refcell::AtomicRefCell;
+
+struct Shared<T> {
+    value: AtomicRefCell<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(value),
+            notify: tokio::sync::Notify::new(),
+        }))
+    }
+
+    pub fn send_if_modified<F>(&self, modify: F) -> bool
+    where
+        F: FnOnce(&mut T) -> bool,
+    {
+        let mut value = self.0.value.borrow_mut();
+        let modified = modify(&mut value);
+        if modified {
+            self.0.notify.notify_waiters();
+        }
+        modified
+    }
+
+    pub fn borrow(&self) -> impl Deref<Target = T> + '_ {
+        self.0.value.borrow()
+    }
+
+    pub fn subscribe(&self) -> Receiver<T> {
+        Receiver(self.0.clone())
+    }
+}
+
+impl<T> Receiver<T> {
+    pub async fn changed(&self) {
+        self.0.notify.notified().await;
+    }
+
+    pub fn borrow(&self) -> impl Deref<Target = T> + '_ {
+        self.0.value.borrow()
+    }
+}