Fix stores deadlock (#4950)

* Fix stores deadlock

* fix clippy

* more clippy fixes

* move panic for reading invalid keys into userland instead of in the
hashmap iter method

* apply the same fixes to btreemap

* return a dropped error instead of panicing if the item has been removed

Cargo.lock

@@ -6446,6 +6446,7 @@ dependencies = [
  "dioxus-core",
  "dioxus-signals",
  "dioxus-stores-macro",
+ "generational-box",
 ]
 
 [[package]]

packages/stores/Cargo.toml

@@ -13,6 +13,7 @@ description = "Server function macros for Dioxus"
 dioxus-core = { workspace = true }
 dioxus-signals = { workspace = true }
 dioxus-stores-macro = { workspace = true, optional = true }
+generational-box.workspace = true
 
 [dev-dependencies]
 dioxus = { workspace = true }

packages/stores/src/impls/btreemap.rs

@@ -1,12 +1,15 @@
 //! Additional utilities for `BTreeMap` stores.
 
-use std::{borrow::Borrow, collections::BTreeMap, hash::Hash, iter::FusedIterator};
+use std::{
+    borrow::Borrow, collections::BTreeMap, hash::Hash, iter::FusedIterator, panic::Location,
+};
 
 use crate::{store::Store, ReadStore};
 use dioxus_signals::{
     AnyStorage, BorrowError, BorrowMutError, ReadSignal, Readable, ReadableExt, UnsyncStorage,
     Writable, WriteLock, WriteSignal,
 };
+use generational_box::ValueDroppedError;
 
 impl<Lens: Readable<Target = BTreeMap<K, V>> + 'static, K: 'static, V: 'static>
     Store<BTreeMap<K, V>, Lens>
@@ -77,7 +80,7 @@ impl<Lens: Readable<Target = BTreeMap<K, V>> + 'static, K: 'static, V: 'static>
         self.selector().track_shallow();
         let keys: Vec<_> = self.selector().peek_unchecked().keys().cloned().collect();
         keys.into_iter().map(move |key| {
-            let value = self.clone().get(key.clone()).unwrap();
+            let value = self.clone().get_unchecked(key.clone());
             (key, value)
         })
     }
@@ -110,7 +113,7 @@ impl<Lens: Readable<Target = BTreeMap<K, V>> + 'static, K: 'static, V: 'static>
         self.selector().track_shallow();
         let keys = self.selector().peek().keys().cloned().collect::<Vec<_>>();
         keys.into_iter()
-            .map(move |key| self.clone().get(key).unwrap())
+            .map(move |key| self.clone().get_unchecked(key))
     }
 
     /// Insert a new key-value pair into the BTreeMap. This method will mark the store as shallowly dirty, causing
@@ -256,15 +259,38 @@ impl<Lens: Readable<Target = BTreeMap<K, V>> + 'static, K: 'static, V: 'static>
         Q: Hash + Ord + 'static,
         K: Borrow<Q> + Ord,
     {
-        self.contains_key(&key).then(|| {
-            self.into_selector()
-                .hash_child_unmapped(key.borrow())
-                .map_writer(move |writer| GetWrite {
-                    index: key,
-                    write: writer,
-                })
-                .into()
-        })
+        self.contains_key(&key).then(|| self.get_unchecked(key))
+    }
+
+    /// Get a store for the value associated with the given key without checking if the key exists.
+    /// This method creates a new store scope that tracks just changes to the value associated with the key.
+    ///
+    /// This is not unsafe, but it will panic when you try to read the value if it does not exist.
+    ///
+    /// # Example
+    /// ```rust, no_run
+    /// use dioxus_stores::*;
+    /// use dioxus::prelude::*;
+    /// use std::collections::BTreeMap;
+    /// let mut store = use_store(|| BTreeMap::new());
+    /// store.insert(0, "value".to_string());
+    /// assert_eq!(store.get_unchecked(0).cloned(), "value".to_string());
+    /// ```
+    #[track_caller]
+    pub fn get_unchecked<Q>(self, key: Q) -> Store<V, GetWrite<Q, Lens>>
+    where
+        Q: Hash + Ord + 'static,
+        K: Borrow<Q> + Ord,
+    {
+        let created = std::panic::Location::caller();
+        self.into_selector()
+            .hash_child_unmapped(key.borrow())
+            .map_writer(move |writer| GetWrite {
+                index: key,
+                write: writer,
+                created,
+            })
+            .into()
     }
 }
 
@@ -273,6 +299,7 @@ impl<Lens: Readable<Target = BTreeMap<K, V>> + 'static, K: 'static, V: 'static>
 pub struct GetWrite<Index, Write> {
     index: Index,
     write: Write,
+    created: &'static Location<'static>,
 }
 
 impl<Index, Write, K, V> Readable for GetWrite<Index, Write>
@@ -289,12 +316,9 @@ where
     where
         Self::Target: 'static,
     {
-        self.write.try_read_unchecked().map(|value| {
-            Self::Storage::map(value, |value: &Write::Target| {
-                value
-                    .get(&self.index)
-                    .expect("Tried to access a key that does not exist")
-            })
+        self.write.try_read_unchecked().and_then(|value| {
+            Self::Storage::try_map(value, |value: &Write::Target| value.get(&self.index))
+                .ok_or_else(|| BorrowError::Dropped(ValueDroppedError::new(self.created)))
         })
     }
 
@@ -302,12 +326,9 @@ where
     where
         Self::Target: 'static,
     {
-        self.write.try_peek_unchecked().map(|value| {
-            Self::Storage::map(value, |value: &Write::Target| {
-                value
-                    .get(&self.index)
-                    .expect("Tried to access a key that does not exist")
-            })
+        self.write.try_peek_unchecked().and_then(|value| {
+            Self::Storage::try_map(value, |value: &Write::Target| value.get(&self.index))
+                .ok_or_else(|| BorrowError::Dropped(ValueDroppedError::new(self.created)))
         })
     }
 
@@ -333,12 +354,11 @@ where
     where
         Self::Target: 'static,
     {
-        self.write.try_write_unchecked().map(|value| {
-            WriteLock::map(value, |value: &mut Write::Target| {
-                value
-                    .get_mut(&self.index)
-                    .expect("Tried to access a key that does not exist")
+        self.write.try_write_unchecked().and_then(|value| {
+            WriteLock::filter_map(value, |value: &mut Write::Target| {
+                value.get_mut(&self.index)
             })
+            .ok_or_else(|| BorrowMutError::Dropped(ValueDroppedError::new(self.created)))
         })
     }
 }

packages/stores/src/impls/hashmap.rs

@@ -5,6 +5,7 @@ use std::{
     collections::HashMap,
     hash::{BuildHasher, Hash},
     iter::FusedIterator,
+    panic::Location,
 };
 
 use crate::{store::Store, ReadStore};
@@ -12,6 +13,7 @@ use dioxus_signals::{
     AnyStorage, BorrowError, BorrowMutError, ReadSignal, Readable, ReadableExt, UnsyncStorage,
     Writable, WriteLock, WriteSignal,
 };
+use generational_box::ValueDroppedError;
 
 impl<Lens: Readable<Target = HashMap<K, V, St>> + 'static, K: 'static, V: 'static, St: 'static>
     Store<HashMap<K, V, St>, Lens>
@@ -82,10 +84,8 @@ impl<Lens: Readable<Target = HashMap<K, V, St>> + 'static, K: 'static, V: 'stati
     {
         self.selector().track_shallow();
         let keys: Vec<_> = self.selector().peek_unchecked().keys().cloned().collect();
-        keys.into_iter().map(move |key| {
-            let value = self.clone().get(key.clone()).unwrap();
-            (key, value)
-        })
+        keys.into_iter()
+            .map(move |key| (key.clone(), self.clone().get_unchecked(key)))
     }
 
     /// Get an iterator over the values in the HashMap. This method will track the store shallowly and only cause
@@ -117,7 +117,7 @@ impl<Lens: Readable<Target = HashMap<K, V, St>> + 'static, K: 'static, V: 'stati
         self.selector().track_shallow();
         let keys = self.selector().peek().keys().cloned().collect::<Vec<_>>();
         keys.into_iter()
-            .map(move |key| self.clone().get(key).unwrap())
+            .map(move |key| self.clone().get_unchecked(key))
     }
 
     /// Insert a new key-value pair into the HashMap. This method will mark the store as shallowly dirty, causing
@@ -264,15 +264,39 @@ impl<Lens: Readable<Target = HashMap<K, V, St>> + 'static, K: 'static, V: 'stati
         K: Borrow<Q> + Eq + Hash,
         St: BuildHasher,
     {
-        self.contains_key(&key).then(|| {
-            self.into_selector()
-                .hash_child_unmapped(key.borrow())
-                .map_writer(move |writer| GetWrite {
-                    index: key,
-                    write: writer,
-                })
-                .into()
-        })
+        self.contains_key(&key).then(|| self.get_unchecked(key))
+    }
+
+    /// Get a store for the value associated with the given key without checking if the key exists.
+    /// This method creates a new store scope that tracks just changes to the value associated with the key.
+    ///
+    /// This is not unsafe, but it will panic when you try to read the value if it does not exist.
+    ///
+    /// # Example
+    /// ```rust, no_run
+    /// use dioxus_stores::*;
+    /// use dioxus::prelude::*;
+    /// use std::collections::HashMap;
+    /// let mut store = use_store(|| HashMap::new());
+    /// store.insert(0, "value".to_string());
+    /// assert_eq!(store.get_unchecked(0).cloned(), "value".to_string());
+    /// ```
+    #[track_caller]
+    pub fn get_unchecked<Q>(self, key: Q) -> Store<V, GetWrite<Q, Lens>>
+    where
+        Q: Hash + Eq + 'static,
+        K: Borrow<Q> + Eq + Hash,
+        St: BuildHasher,
+    {
+        let location = Location::caller();
+        self.into_selector()
+            .hash_child_unmapped(key.borrow())
+            .map_writer(move |writer| GetWrite {
+                index: key,
+                write: writer,
+                created: location,
+            })
+            .into()
     }
 }
 
@@ -281,6 +305,7 @@ impl<Lens: Readable<Target = HashMap<K, V, St>> + 'static, K: 'static, V: 'stati
 pub struct GetWrite<Index, Write> {
     index: Index,
     write: Write,
+    created: &'static Location<'static>,
 }
 
 impl<Index, Write, K, V, St> Readable for GetWrite<Index, Write>
@@ -298,12 +323,9 @@ where
     where
         Self::Target: 'static,
     {
-        self.write.try_read_unchecked().map(|value| {
-            Self::Storage::map(value, |value: &Write::Target| {
-                value
-                    .get(&self.index)
-                    .expect("Tried to access a key that does not exist")
-            })
+        self.write.try_read_unchecked().and_then(|value| {
+            Self::Storage::try_map(value, |value: &Write::Target| value.get(&self.index))
+                .ok_or_else(|| BorrowError::Dropped(ValueDroppedError::new(self.created)))
         })
     }
 
@@ -311,12 +333,9 @@ where
     where
         Self::Target: 'static,
     {
-        self.write.try_peek_unchecked().map(|value| {
-            Self::Storage::map(value, |value: &Write::Target| {
-                value
-                    .get(&self.index)
-                    .expect("Tried to access a key that does not exist")
-            })
+        self.write.try_peek_unchecked().and_then(|value| {
+            Self::Storage::try_map(value, |value: &Write::Target| value.get(&self.index))
+                .ok_or_else(|| BorrowError::Dropped(ValueDroppedError::new(self.created)))
         })
     }
 
@@ -343,12 +362,11 @@ where
     where
         Self::Target: 'static,
     {
-        self.write.try_write_unchecked().map(|value| {
-            WriteLock::map(value, |value: &mut Write::Target| {
-                value
-                    .get_mut(&self.index)
-                    .expect("Tried to access a key that does not exist")
+        self.write.try_write_unchecked().and_then(|value| {
+            WriteLock::filter_map(value, |value: &mut Write::Target| {
+                value.get_mut(&self.index)
             })
+            .ok_or_else(|| BorrowMutError::Dropped(ValueDroppedError::new(self.created)))
         })
     }
 }

packages/stores/src/subscriptions.rs

@@ -49,67 +49,38 @@ impl SelectorNode {
             .get_mut_or_default(rest)
     }
 
-    /// Visit this node and all of its children in depth-first order, calling the provided function on each node.
+    /// Get the path to each node under this node
     ///
     /// This is used to mark nodes dirty recursively when a Store is written to.
-    fn visit_depth_first_mut(&mut self, f: &mut dyn FnMut(&mut SelectorNode)) {
-        f(self);
-        for child in self.root.values_mut() {
-            child.visit_depth_first_mut(f);
+    fn paths_under(&self, current_path: &[PathKey], paths: &mut Vec<Box<[PathKey]>>) {
+        paths.push(current_path.into());
+        for child in self.root.values() {
+            child.paths_under(current_path, paths);
         }
     }
 
-    /// Mark this selector and all children as dirty. This should be called any time a raw mutable reference to a store
-    /// is exposed to the user. They could write to any level of the store, so we need to mark all nodes as dirty.
-    fn mark_children_dirty(&mut self, path: &[PathKey]) {
-        let Some(node) = self.get_mut(path) else {
-            return;
-        };
-
-        // Mark the node and all its children as dirty
-        node.visit_depth_first_mut(&mut |node| {
-            node.mark_dirty();
-        });
-    }
-
-    /// Mark only children after a certain index as dirty. This is used when inserting a new item into a list.
+    /// Get paths to only children before a certain index.
+    ///
+    /// This is used when inserting a new item into a list.
     /// Items after the index that is inserted need to be marked dirty because the value that index points to may have changed.
-    fn mark_dirty_at_and_after_index(&mut self, path: &[PathKey], index: usize) {
-        let Some(node) = self.get_mut(path) else {
+    fn paths_at_and_after_index(
+        &self,
+        path: &[PathKey],
+        index: usize,
+        paths: &mut Vec<Box<[PathKey]>>,
+    ) {
+        let Some(node) = self.get(path) else {
             return;
         };
 
         // Mark the nodes before the index as dirty
-        for (i, child) in node.root.iter_mut() {
+        for (i, child) in node.root.iter() {
             if *i as usize >= index {
-                child.visit_depth_first_mut(&mut |node| {
-                    node.mark_dirty();
-                });
+                child.paths_under(path, paths);
             }
         }
     }
 
-    /// Mark a specific node as dirty without marking its children. This is used for data structures like HashMaps
-    /// when inserting or removing items. Inserting an item into a HashMap only changes the length of the map and the
-    /// specific value that was inserted or removed.
-    fn mark_dirty_shallow(&mut self, path: &[PathKey]) {
-        let Some(node) = self.get_mut(path) else {
-            return;
-        };
-
-        node.mark_dirty();
-    }
-
-    /// Mark this node as dirty, which will notify all subscribers that the value has changed.
-    fn mark_dirty(&mut self) {
-        // We cannot hold the subscribers lock while calling mark_dirty, because mark_dirty can run user code which may cause a new subscriber to be added. If we hold the lock, we will deadlock.
-        #[allow(clippy::mutable_key_type)]
-        let mut subscribers = std::mem::take(&mut self.subscribers);
-        subscribers.retain(|reactive_context| reactive_context.mark_dirty());
-        // Extend the subscribers list instead of overwriting it in case a subscriber is added while reactive contexts are marked dirty
-        self.subscribers.extend(subscribers);
-    }
-
     /// Remove a path from the subscription tree
     fn remove(&mut self, path: &[PathKey]) {
         let [first, rest @ ..] = path else {
@@ -253,19 +224,50 @@ impl StoreSubscriptions {
         }
     }
 
+    /// Mark the node and all its children as dirty
     pub(crate) fn mark_dirty(&self, key: &[PathKey]) {
-        self.inner.write_unchecked().root.mark_children_dirty(key);
+        let paths = {
+            let read = self.inner.read();
+            let mut nodes = Vec::new();
+            read.root.paths_under(key, &mut nodes);
+            nodes
+        };
+        for path in paths {
+            self.mark_dirty_shallow(&path);
+        }
     }
 
+    /// Mark a single node as dirty
     pub(crate) fn mark_dirty_shallow(&self, key: &[PathKey]) {
-        self.inner.write_unchecked().root.mark_dirty_shallow(key);
+        // We cannot hold the subscribers lock while calling mark_dirty, because mark_dirty can run user code which may cause a new subscriber to be added. If we hold the lock, we will deadlock.
+        #[allow(clippy::mutable_key_type)]
+        let mut subscribers = {
+            let mut write = self.inner.write_unchecked();
+            let Some(node) = write.root.get_mut(key) else {
+                return;
+            };
+            std::mem::take(&mut node.subscribers)
+        };
+        subscribers.retain(|reactive_context| reactive_context.mark_dirty());
+        // Extend the subscribers list instead of overwriting it in case a subscriber is added while reactive contexts are marked dirty
+        let mut write = self.inner.write_unchecked();
+        let Some(node) = write.root.get_mut(key) else {
+            return;
+        };
+        node.subscribers.extend(subscribers);
     }
 
+    /// Mark all nodes after the index and their children as dirty
     pub(crate) fn mark_dirty_at_and_after_index(&self, key: &[PathKey], index: usize) {
-        self.inner
-            .write_unchecked()
-            .root
-            .mark_dirty_at_and_after_index(key, index);
+        let paths = {
+            let read = self.inner.read_unchecked();
+            let mut paths = Vec::new();
+            read.root.paths_at_and_after_index(key, index, &mut paths);
+            paths
+        };
+        for path in paths {
+            self.mark_dirty_shallow(&path);
+        }
     }
 
     /// Get a subscriber list for a specific path in the store. This is used to subscribe to changes