fuchsia_bluetooth/
detachable_map.rs

Help
1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use fuchsia_sync::RwLock;
6use std::borrow::Borrow;
7use std::collections::hash_map::Entry;
8use std::collections::HashMap;
9use std::hash::Hash;
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, Weak};
12
13// TODO(https://fxbug.dev/42154109): This would be so so much easier with Arc::new_cyclic.
14
15/// A weak reference to an entry in a DetachableMap. This is a weak reference, if the entry is
16/// detached before it is upgraded, the entry can be gone.
17pub struct DetachableWeak<K, V> {
18    /// Weak reference to the value, upgrading this doesn't guarantee it's in the map, but usually
19    /// it is.
20    inner: Weak<V>,
21    parent: Arc<RwLock<HashMap<K, (Arc<V>, Arc<AtomicBool>)>>>,
22    /// Key reference, shared with all DetachableWeaks that reference the same entry in `parent`.
23    /// false if anyone detached the entry or the key was replaced in `parent`, to prevent removing
24    /// an entry at the same key.
25    key_ref: Arc<AtomicBool>,
26    /// A copy of `key` that can be used whenever.
27    key: K,
28}
29
30impl<K: Clone, V> Clone for DetachableWeak<K, V> {
31    fn clone(&self) -> Self {
32        Self {
33            inner: self.inner.clone(),
34            parent: self.parent.clone(),
35            key_ref: self.key_ref.clone(),
36            key: self.key.clone(),
37        }
38    }
39}
40
41impl<K: Hash + Eq, V> DetachableWeak<K, V> {
42    fn new(
43        item: &Arc<V>,
44        parent: Arc<RwLock<HashMap<K, (Arc<V>, Arc<AtomicBool>)>>>,
45        key: K,
46        key_ref: Arc<AtomicBool>,
47    ) -> Self {
48        Self { inner: Arc::downgrade(item), parent, key, key_ref }
49    }
50
51    /// Attempt to upgrade the weak pointer to an Arc, extending the lifetime of the value if
52    /// successful.  Returns None if the item has been dropped (by another client detaching this)
53    pub fn upgrade(&self) -> Option<Arc<V>> {
54        self.inner.upgrade()
55    }
56
57    /// Destroys the original reference to this vended item.
58    /// If other references to the item exist (from `upgrade`), it will not be dropped until those
59    /// references are dropped.
60    pub fn detach(self) {
61        let mut lock = self.parent.write();
62        if self.key_ref.swap(false, Ordering::Relaxed) {
63            // Dropping this reference only frees the key_ref and value if no other references are
64            // held.
65            let _ = lock.remove(&self.key);
66        }
67    }
68
69    /// Get a reference to the key for this entry.
70    pub fn key(&self) -> &K {
71        &self.key
72    }
73}
74
75pub struct LazyEntry<K, V> {
76    parent: Arc<RwLock<HashMap<K, (Arc<V>, Arc<AtomicBool>)>>>,
77    key: K,
78}
79
80impl<K: Clone, V> Clone for LazyEntry<K, V> {
81    fn clone(&self) -> Self {
82        Self { parent: self.parent.clone(), key: self.key.clone() }
83    }
84}
85
86impl<K: Hash + Eq + Clone, V> LazyEntry<K, V> {
87    fn new(parent: Arc<RwLock<HashMap<K, (Arc<V>, Arc<AtomicBool>)>>>, key: K) -> Self {
88        Self { parent, key }
89    }
90
91    /// Get a reference to the key that this entry is built for.
92    pub fn key(&self) -> &K {
93        &self.key
94    }
95
96    /// Attempt to insert into the map at `key`. Returns a detachable weak entry if the value was
97    /// inserted, and Err(value) if the item already existed.
98    /// Err(value) if the value was not able to be inserted.
99    pub fn try_insert(&self, value: V) -> Result<DetachableWeak<K, V>, V> {
100        match self.parent.write().entry(self.key.clone()) {
101            Entry::Occupied(_) => return Err(value),
102            Entry::Vacant(entry) => {
103                let _ = entry.insert((Arc::new(value), Arc::new(AtomicBool::new(true))));
104            }
105        };
106        Ok(self.get().unwrap())
107    }
108
109    /// Attempt to resolve the entry to a weak reference, as returned by `DetachableMap::get`.
110    /// Returns None if the key does not exist.
111    pub fn get(&self) -> Option<DetachableWeak<K, V>> {
112        self.parent.read().get(&self.key).and_then(|(v, key_ref)| {
113            let map = self.parent.clone();
114            let key_ref = key_ref.clone();
115            let key = self.key.clone();
116            Some(DetachableWeak::new(&v, map, key, key_ref))
117        })
118    }
119}
120
121/// A Map with detachable entries.  After retrieval, entries can be "detached", removing them from
122/// the map, allowing any client to expire the key in the map.  They are weak, and can be upgraded
123/// to a strong reference to the stored object.
124pub struct DetachableMap<K, V> {
125    /// The map. For each key, it holds a reference to the value, and a shared bool indicating if
126    /// the paired value is still in the map.
127    map: Arc<RwLock<HashMap<K, (Arc<V>, Arc<AtomicBool>)>>>,
128}
129
130impl<K: Hash + Eq + Clone, V> Default for DetachableMap<K, V> {
131    fn default() -> Self {
132        Self { map: Arc::new(RwLock::new(HashMap::new())) }
133    }
134}
135
136impl<K: Hash + Eq + Clone, V> DetachableMap<K, V> {
137    /// Creates an empty `DetachableMap`.   The map is initially empty.
138    pub fn new() -> DetachableMap<K, V> {
139        Default::default()
140    }
141
142    /// Inserts a new item into the map at `key`
143    /// Returns a reference to the old item at `key` if one existed or None otherwise.
144    pub fn insert(&mut self, key: K, value: V) -> Option<Arc<V>> {
145        let mut lock = self.map.write();
146        let new_pair = (Arc::new(value), Arc::new(AtomicBool::new(true)));
147        if let Some((prev_val, old_key_ref)) = lock.insert(key, new_pair) {
148            old_key_ref.store(false, Ordering::Relaxed);
149            Some(prev_val)
150        } else {
151            None
152        }
153    }
154
155    /// True if the map contains a value for the specified key  The key may be any borrowed form of
156    /// the key's type, with `Hash` and `Eq` matching the type.
157    pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool
158    where
159        K: Borrow<Q>,
160        Q: Hash + Eq,
161    {
162        self.map.read().contains_key(key)
163    }
164
165    /// Returns a detachable reference to the value at the given key, if it exists.
166    pub fn get(&self, key: &K) -> Option<DetachableWeak<K, V>> {
167        self.map.read().get(key).and_then(|(v, key_ref)| {
168            let map = self.map.clone();
169            let key_ref = key_ref.clone();
170            let key = key.clone();
171            Some(DetachableWeak::new(&v, map, key, key_ref))
172        })
173    }
174
175    /// Returns a lazy entry. Lazy Entries can be used later to attempt to insert into the map if
176    /// the key doesn't exist.
177    /// They can also be resolved to a detachable reference (as returned by `DetachableMap::get`) if
178    /// the key already exists.
179    pub fn lazy_entry(&self, key: &K) -> LazyEntry<K, V> {
180        LazyEntry::new(self.map.clone(), key.clone())
181    }
182}
183
184#[cfg(test)]
185mod test {
186    use super::*;
187
188    #[derive(Default, Debug, PartialEq)]
189    struct TestStruct {
190        data: u32,
191    }
192
193    #[test]
194    fn contains_keys() {
195        let mut map = DetachableMap::default();
196
197        assert!(map.insert(0, TestStruct { data: 45 }).is_none());
198
199        assert!(map.contains_key(&0));
200
201        let detached = map.get(&0);
202        assert!(detached.is_some());
203        // Detaching removes it from the map.
204        detached.unwrap().detach();
205
206        assert_eq!(false, map.contains_key(&0));
207    }
208
209    #[test]
210    fn upgrade_detached() {
211        let mut map = DetachableMap::default();
212
213        assert!(map.insert(0, TestStruct { data: 45 }).is_none());
214
215        let detached = map.get(&0);
216        assert!(detached.is_some());
217        let detached = detached.unwrap();
218        let detached_clone = detached.clone();
219
220        let second = map.get(&0);
221        assert!(second.is_some());
222        let second = second.unwrap();
223
224        let upgraded = detached.upgrade().expect("should be able to upgrade");
225
226        // Detaching should mean we can't get it from the map anymore (and consumes second)
227        second.detach();
228
229        assert!(map.get(&0).is_none());
230
231        // We can still upgrade because it's still around from the strong ref.
232        let second_up = detached.upgrade().expect("should be able to upgrade");
233        let third_up = detached_clone.upgrade().expect("should be able to upgrade");
234
235        // Dropping all the strong refs means neither can upgrade anymore though.
236        drop(upgraded);
237        drop(second_up);
238        drop(third_up);
239
240        assert!(detached.upgrade().is_none());
241
242        // Detaching twice doesn't do anything (and doesn't panic)
243        detached.detach();
244    }
245
246    #[test]
247    fn cant_detach_replaced_key() {
248        let mut map = DetachableMap::default();
249
250        assert!(map.insert(0, TestStruct { data: 12 }).is_none());
251
252        let detached = map.get(&0).expect("key is in");
253
254        // Replace the struct in the original map, keeping the data around.
255        let replaced = map.insert(0, TestStruct { data: 34 });
256
257        // Can still upgrade to the old data, since it's still alive.
258        assert_eq!(TestStruct { data: 12 }, *detached.upgrade().expect("data is still alive"));
259
260        // dropping the old data means upgrade fails though.
261        drop(replaced);
262
263        assert!(detached.upgrade().is_none());
264
265        // Trying to detach from an earlier generation will not remove the new data.
266        detached.detach();
267
268        let new_detached = map.get(&0).expect("key was replaced");
269
270        assert_eq!(TestStruct { data: 34 }, *new_detached.upgrade().expect("should be there"));
271
272        // Detaching from a cloned one has no effect, even if the key was re-added before we try.
273        let cloned = new_detached.clone();
274        let another = map.get(&0).expect("still there");
275
276        new_detached.detach();
277
278        assert!(map.get(&0).is_none());
279
280        // We don't replace any data, it got detached above.
281        let replaced = map.insert(0, TestStruct { data: 45 });
282        assert!(replaced.is_none());
283
284        // This doesn't remove the new data either though.
285        cloned.detach();
286        assert!(map.get(&0).is_some());
287        // Even if we didn't clone it from the same DetachedWeak.
288        another.detach();
289        assert!(map.get(&0).is_some());
290    }
291
292    #[test]
293    fn lazy_entry() {
294        let map = DetachableMap::default();
295
296        // Should be able to get an entry before the key exists.
297        let entry = map.lazy_entry(&1);
298
299        // Can't get a reference if the key doesn't exist.
300        assert!(entry.get().is_none());
301
302        // We can insert though.
303        let detachable =
304            entry.try_insert(TestStruct { data: 45 }).expect("should be able to insert");
305
306        // Can't insert if there's something there though.
307        let second_val = TestStruct { data: 56 };
308        let returned_val =
309            entry.try_insert(second_val).err().expect("should get an error when trying to insert");
310        assert_eq!(56, returned_val.data);
311
312        assert!(entry.get().is_some());
313
314        // If we detach though, the entry is empty again, and we can insert again.
315        detachable.detach();
316
317        assert!(entry.get().is_none());
318
319        let new = entry.try_insert(returned_val).expect("should be able to insert after removal");
320
321        // Deopping the new entry doesn't remove it from the map.
322        drop(new);
323
324        let still_there = map.get(&1).expect("should be there");
325        assert!(still_there.upgrade().is_some());
326    }
327}
fuchsia_bluetooth/detachable_map.rs

fuchsia_bluetooth/
detachable_map.rs
