feat: state storage RAII

1 files changed, 19 insertions, 22 deletions

diff --git a/src/state_store.rs b/src/state_store.rs
index 7e6e6a1..d45001f 100644
--- a/src/state_store.rs
+++ b/src/state_store.rs
@@ -112,8 +112,9 @@ where
     V: Send + Sync + 'static,
 {
     pub fn store(&self, key: K, value: V) -> StateListener<K, V> {
-        let store = self.inner.try_get_value().unwrap();
-        let mut store = store.store.write().unwrap();
+        let arc_store = self.inner.try_get_value().unwrap();
+        let store = arc_store.store.clone();
+        let mut store = store.write().unwrap();
         debug!("store state: {:?}", store);
         if let Some((v, count)) = store.get_mut(&key) {
             debug!("updating old value already in store");
@@ -123,7 +124,7 @@ where
                 value: v.clone(),
                 cleaner: StateCleaner {
                     key,
-                    state_store: self.clone(),
+                    state_store: arc_store,
                 },
             }
         } else {
@@ -134,15 +135,16 @@ where
                 value: v.into(),
                 cleaner: StateCleaner {
                     key,
-                    state_store: self.clone(),
+                    state_store: arc_store,
                 },
             }
         }
     }
 
     pub fn get_listener(&self, key: K) -> Option<StateListener<K, V>> {
-        let store = self.inner.try_get_value().unwrap();
-        let mut store = store.store.write().unwrap();
+        let arc_store = self.inner.try_get_value().unwrap();
+        let store = arc_store.store.clone();
+        let mut store = store.write().unwrap();
         debug!("store state: {:?}", store);
         if let Some((v, count)) = store.get_mut(&key) {
             *count += 1;
@@ -150,7 +152,7 @@ where
                 value: v.clone(),
                 cleaner: StateCleaner {
                     key,
-                    state_store: self.clone(),
+                    state_store: arc_store,
                 },
             })
         } else {
@@ -180,17 +182,6 @@ where
         }
     }
 
-    fn remove(&self, key: &K) {
-        // let store = self.inner.try_get_value().unwrap();
-        // let mut store = store.store.write().unwrap();
-        // if let Some((_v, count)) = store.get_mut(key) {
-        //     *count -= 1;
-        //     if *count == 0 {
-        //         store.remove(key);
-        //         debug!("dropped item from store");
-        //     }
-        // }
-    }
 }
 
 #[derive(Clone)]
@@ -234,7 +225,7 @@ where
     V: Send + Sync + 'static,
 {
     key: K,
-    state_store: StateStore<K, V>,
+    state_store: ArcStateStore<K, V>,
 }
 
 impl<K, V> Clone for StateCleaner<K, V>
@@ -244,8 +235,7 @@ where
 {
     fn clone(&self) -> Self {
         {
-            let store = self.state_store.inner.try_get_value().unwrap();
-            let mut store = store.store.write().unwrap();
+            let mut store = self.state_store.store.write().unwrap();
             if let Some((_v, count)) = store.get_mut(&self.key) {
                 *count += 1;
             }
@@ -261,6 +251,13 @@ impl<K: Eq + std::hash::Hash + Send + Sync + 'static, V: Send + Sync + 'static>
     for StateCleaner<K, V>
 {
     fn drop(&mut self) {
-        self.state_store.remove(&self.key);
+        let mut store = self.state_store.store.write().unwrap();
+        if let Some((_v, count)) = store.get_mut(&self.key) {
+            *count -= 1;
+            if *count == 0 {
+                store.remove(&self.key);
+                debug!("dropped item from store");
+            }
+        }
     }
 }