dark_std/sync/
map_btree.rs

1use parking_lot::Mutex;
2use serde::{Deserializer, Serialize, Serializer};
3use std::borrow::Borrow;
4use std::cell::UnsafeCell;
5use std::collections::{
6    btree_map::IntoIter as MapIntoIter, btree_map::Iter as MapIter,
7    btree_map::IterMut as MapIterMut, BTreeMap,
8};
9use std::fmt::{Debug, Display, Formatter};
10use std::hash::Hash;
11use std::marker::PhantomData;
12use std::ops::{Deref, DerefMut};
13use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
14use std::sync::Arc;
15
16use super::{ReadGuard, ReadMapGuard, WriteGuard, WriteLock};
17
18/// Read guard returned by [`SyncBtreeMap::get`].
19pub type BtreeMapGet<'a, V> = ReadGuard<'a, V>;
20
21/// Write guard returned by [`SyncBtreeMap::get_mut`].
22pub struct BtreeMapRefMut<'a, K, V> {
23    inner: WriteGuard<'a, V>,
24    _k: PhantomData<&'a K>,
25}
26
27impl<'a, K, V> BtreeMapRefMut<'a, K, V> {
28    #[inline]
29    pub(crate) fn new(inner: WriteGuard<'a, V>) -> Self {
30        BtreeMapRefMut {
31            inner,
32            _k: PhantomData,
33        }
34    }
35}
36
37impl<'a, K, V> Deref for BtreeMapRefMut<'a, K, V> {
38    type Target = V;
39
40    #[inline]
41    fn deref(&self) -> &Self::Target {
42        &self.inner
43    }
44}
45
46impl<'a, K, V> DerefMut for BtreeMapRefMut<'a, K, V> {
47    #[inline]
48    fn deref_mut(&mut self) -> &mut Self::Target {
49        &mut self.inner
50    }
51}
52
53impl<'a, K, V: Debug> Debug for BtreeMapRefMut<'a, K, V> {
54    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
55        Debug::fmt(&*self.inner, f)
56    }
57}
58
59impl<'a, K, V: Display> Display for BtreeMapRefMut<'a, K, V> {
60    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
61        Display::fmt(&*self.inner, f)
62    }
63}
64
65impl<'a, K, V: PartialEq> PartialEq for BtreeMapRefMut<'a, K, V> {
66    fn eq(&self, other: &Self) -> bool {
67        *self.inner == *other.inner
68    }
69}
70
71impl<'a, K, V: Eq> Eq for BtreeMapRefMut<'a, K, V> {}
72
73/// Read iterator returned by [`SyncBtreeMap::iter`].
74pub struct BtreeMapIter<'a, K, V> {
75    count: &'a AtomicUsize,
76    inner: MapIter<'a, K, V>,
77    _not_send: PhantomData<*const ()>,
78}
79
80impl<'a, K, V> Drop for BtreeMapIter<'a, K, V> {
81    fn drop(&mut self) {
82        self.count.fetch_sub(1, Ordering::Release);
83    }
84}
85
86impl<'a, K, V> Iterator for BtreeMapIter<'a, K, V> {
87    type Item = (&'a K, &'a V);
88
89    fn next(&mut self) -> Option<Self::Item> {
90        self.inner.next()
91    }
92}
93
94/// Write iterator returned by [`SyncBtreeMap::iter_mut`].
95pub struct BtreeMapIterMut<'a, K, V> {
96    _w: WriteLock<'a>,
97    inner: MapIterMut<'a, K, V>,
98}
99
100impl<'a, K, V> Iterator for BtreeMapIterMut<'a, K, V> {
101    type Item = (&'a K, &'a mut V);
102
103    fn next(&mut self) -> Option<Self::Item> {
104        self.inner.next()
105    }
106}
107
108/// this sync map used to many reader,writer less.space-for-time strategy
109///
110/// Reads are lock-free: `get`/`iter`/`dirty_ref`/`len`/`contains_key` only
111/// register a reader slot with an atomic counter and then read the map without
112/// any lock (readers never block each other and never touch a lock word).
113/// Writes take a mutex, raise a `writing` flag and wait until all in-flight
114/// readers are gone before mutating the map in place — O(log n), no
115/// whole-container copy and no `Clone` requirement on `K`/`V`.
116///
117/// # Deadlock note
118/// A read guard makes writers wait until it is dropped. Do not call a write
119/// method while a read/write guard is alive in the same scope: drop the guard
120/// first (e.g. `drop(g)` before `insert`/`remove`/`get_mut`), otherwise the
121/// writer waits for its own guard and deadlocks.
122pub struct SyncBtreeMap<K: Eq + Hash, V> {
123    dirty: UnsafeCell<BTreeMap<K, V>>,
124    write: Mutex<()>,
125    id: usize,
126    writing: AtomicBool,
127    registry: Mutex<Vec<std::boxed::Box<AtomicUsize>>>,
128}
129
130// SAFETY: all writers hold `write` and wait for `readers` to drain before
131// touching `dirty`; readers either see a consistent snapshot or retry while a
132// writer is active, so concurrent access to `dirty` is race-free.
133unsafe impl<K: Eq + Hash, V: Send> Send for SyncBtreeMap<K, V> {}
134unsafe impl<K: Eq + Hash, V: Sync> Sync for SyncBtreeMap<K, V> {}
135
136impl<K, V> SyncBtreeMap<K, V>
137where
138    K: Eq + Hash,
139{
140    #[inline]
141    fn begin_read(&self) -> &AtomicUsize {
142        // The counter lives in thread-local storage: concurrent readers only
143        // touch their own cache line and never contend with each other. SeqCst
144        // closes the store-buffering window with the writer's all-zero scan.
145        let count = super::reader_count_for(self.id, &self.registry);
146        loop {
147            count.fetch_add(1, Ordering::SeqCst);
148            if !self.writing.load(Ordering::SeqCst) {
149                return count;
150            }
151            count.fetch_sub(1, Ordering::SeqCst);
152            std::thread::yield_now();
153        }
154    }
155
156    #[inline]
157    fn begin_write(&self) -> WriteLock<'_> {
158        let lock = self.write.lock();
159        self.writing.store(true, Ordering::SeqCst);
160        loop {
161            let registry = self.registry.lock();
162            let all_zero = registry.iter().all(|c| c.load(Ordering::SeqCst) == 0);
163            if all_zero {
164                break;
165            }
166            drop(registry);
167            std::thread::yield_now();
168        }
169        WriteLock::new(lock, &self.writing)
170    }
171
172    pub fn new_arc() -> Arc<Self> {
173        Arc::new(Self::new())
174    }
175
176    pub fn new() -> Self {
177        Self {
178            dirty: UnsafeCell::new(BTreeMap::new()),
179            write: Mutex::new(()),
180            id: super::CONTAINER_ID.fetch_add(1, Ordering::Relaxed),
181            writing: AtomicBool::new(false),
182            registry: Mutex::new(Vec::new()),
183        }
184    }
185
186    pub fn with_capacity(_capacity: usize) -> Self {
187        Self::new()
188    }
189
190    pub fn with_map(map: BTreeMap<K, V>) -> Self
191    where
192        K: Ord,
193    {
194        Self {
195            dirty: UnsafeCell::new(map),
196            write: Mutex::new(()),
197            id: super::CONTAINER_ID.fetch_add(1, Ordering::Relaxed),
198            writing: AtomicBool::new(false),
199            registry: Mutex::new(Vec::new()),
200        }
201    }
202
203    pub fn insert(&self, k: K, v: V) -> Option<V>
204    where
205        K: Ord,
206    {
207        let _w = self.begin_write();
208        unsafe { &mut *self.dirty.get() }.insert(k, v)
209    }
210
211    pub fn insert_mut(&mut self, k: K, v: V) -> Option<V>
212    where
213        K: Ord,
214    {
215        unsafe { &mut *self.dirty.get() }.insert(k, v)
216    }
217
218    pub fn remove(&self, k: &K) -> Option<V>
219    where
220        K: Ord,
221    {
222        let _w = self.begin_write();
223        unsafe { &mut *self.dirty.get() }.remove(k)
224    }
225
226    pub fn remove_mut(&mut self, k: &K) -> Option<V>
227    where
228        K: Ord,
229    {
230        unsafe { &mut *self.dirty.get() }.remove(k)
231    }
232
233    pub fn len(&self) -> usize {
234        let count = self.begin_read();
235        let n = unsafe { &*self.dirty.get() }.len();
236        count.fetch_sub(1, Ordering::Release);
237        n
238    }
239
240    pub fn is_empty(&self) -> bool {
241        let count = self.begin_read();
242        let b = unsafe { &*self.dirty.get() }.is_empty();
243        count.fetch_sub(1, Ordering::Release);
244        b
245    }
246
247    pub fn clear(&self) {
248        let _w = self.begin_write();
249        unsafe { &mut *self.dirty.get() }.clear();
250    }
251
252    pub fn clear_mut(&mut self) {
253        unsafe { &mut *self.dirty.get() }.clear();
254    }
255
256    pub fn shrink_to_fit(&self) {}
257
258    pub fn shrink_to_fit_mut(&mut self) {}
259
260    pub fn from(map: BTreeMap<K, V>) -> Self
261    where
262        K: Eq + Hash + Ord,
263    {
264        Self::with_map(map)
265    }
266
267    /// Returns a read-guarded reference to the value corresponding to the key.
268    ///
269    /// The key may be any borrowed form of the map's key type.
270    ///
271    /// The read is lock-free: it only registers a reader slot, so concurrent
272    /// reads never block each other and never take a lock. Writers wait for
273    /// the returned guard to be dropped before mutating the map.
274    ///
275    /// # Examples
276    ///
277    /// ```
278    /// use dark_std::sync::{SyncBtreeMap};
279    ///
280    /// let mut map = SyncBtreeMap::new();
281    /// map.insert_mut(1, "a");
282    /// assert_eq!(*map.get(&1).unwrap(), "a");
283    /// assert_eq!(map.get(&2).is_none(), true);
284    /// ```
285    #[inline]
286    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<BtreeMapGet<'_, V>>
287    where
288        K: Borrow<Q> + Ord,
289        Q: Ord,
290    {
291        let count = self.begin_read();
292        let m = unsafe { &*self.dirty.get() };
293        match m.get(k) {
294            Some(v) => Some(ReadGuard::new(count, v)),
295            None => {
296                count.fetch_sub(1, Ordering::Release);
297                None
298            }
299        }
300    }
301
302    /// Returns a write-guarded mutable reference to the value of the key.
303    ///
304    /// The guard holds the writer lock (writers are mutually exclusive and
305    /// wait for in-flight readers) until it is dropped, so the mutable
306    /// reference can never race with concurrent readers or writers. Drop it
307    /// before calling another method from the same scope.
308    #[inline]
309    pub fn get_mut(&self, k: &K) -> Option<BtreeMapRefMut<'_, K, V>>
310    where
311        K: Ord,
312    {
313        let w = self.begin_write();
314        let m = unsafe { &mut *self.dirty.get() };
315        match m.get_mut(k) {
316            Some(v) => Some(BtreeMapRefMut::new(WriteGuard::new(w, v))),
317            None => None,
318        }
319    }
320
321    #[inline]
322    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
323    where
324        K: Borrow<Q> + Ord,
325        Q: Ord,
326    {
327        let count = self.begin_read();
328        let b = unsafe { &*self.dirty.get() }.contains_key(k);
329        count.fetch_sub(1, Ordering::Release);
330        b
331    }
332
333    pub fn iter(&self) -> BtreeMapIter<'_, K, V> {
334        let count = self.begin_read();
335        let m = unsafe { &*self.dirty.get() };
336        BtreeMapIter {
337            count,
338            inner: m.iter(),
339            _not_send: PhantomData,
340        }
341    }
342
343    pub fn iter_mut(&self) -> BtreeMapIterMut<'_, K, V> {
344        let w = self.begin_write();
345        let m = unsafe { &mut *self.dirty.get() };
346        BtreeMapIterMut {
347            _w: w,
348            inner: m.iter_mut(),
349        }
350    }
351
352    pub fn into_iter(self) -> MapIntoIter<K, V>
353    where
354        K: Ord,
355    {
356        self.into_inner().into_iter()
357    }
358
359    pub fn dirty_ref(&self) -> ReadMapGuard<'_, BTreeMap<K, V>> {
360        let count = self.begin_read();
361        let m = unsafe { &*self.dirty.get() };
362        ReadMapGuard::new(count, m)
363    }
364
365    pub fn into_inner(self) -> BTreeMap<K, V>
366    where
367        K: Ord,
368    {
369        self.dirty.into_inner()
370    }
371}
372
373impl<K: Eq + Hash + Ord, V> IntoIterator for SyncBtreeMap<K, V> {
374    type Item = (K, V);
375    type IntoIter = MapIntoIter<K, V>;
376
377    fn into_iter(self) -> Self::IntoIter {
378        self.into_iter()
379    }
380}
381
382impl<'a, K: Eq + Hash, V> IntoIterator for &'a SyncBtreeMap<K, V> {
383    type Item = (&'a K, &'a V);
384    type IntoIter = BtreeMapIter<'a, K, V>;
385
386    fn into_iter(self) -> Self::IntoIter {
387        self.iter()
388    }
389}
390
391impl<K: Eq + Hash + Ord, V> From<BTreeMap<K, V>> for SyncBtreeMap<K, V> {
392    fn from(arg: BTreeMap<K, V>) -> Self {
393        Self::from(arg)
394    }
395}
396
397impl<K, V> serde::Serialize for SyncBtreeMap<K, V>
398where
399    K: Eq + Hash + Serialize + Ord,
400    V: Serialize,
401{
402    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
403    where
404        S: Serializer,
405    {
406        self.dirty_ref().serialize(serializer)
407    }
408}
409
410impl<'de, K, V> serde::Deserialize<'de> for SyncBtreeMap<K, V>
411where
412    K: Eq + Hash + Ord + serde::Deserialize<'de>,
413    V: serde::Deserialize<'de>,
414{
415    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
416    where
417        D: Deserializer<'de>,
418    {
419        let m = BTreeMap::deserialize(deserializer)?;
420        Ok(Self::from(m))
421    }
422}
423
424impl<K, V> Debug for SyncBtreeMap<K, V>
425where
426    K: Eq + Hash + Debug,
427    V: Debug,
428{
429    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
430        Debug::fmt(&*self.dirty_ref(), f)
431    }
432}
433
434impl<K, V> Display for SyncBtreeMap<K, V>
435where
436    K: Eq + Hash + Debug,
437    V: Debug,
438{
439    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
440        Debug::fmt(&*self.dirty_ref(), f)
441    }
442}
443
444impl<K: Clone + Eq + Hash + Ord, V: Clone> Clone for SyncBtreeMap<K, V> {
445    fn clone(&self) -> Self {
446        let c = (*self.dirty_ref()).clone();
447        SyncBtreeMap::from(c)
448    }
449}
450
451impl<K: Eq + Hash, V> Default for SyncBtreeMap<K, V> {
452    fn default() -> Self {
453        SyncBtreeMap::new()
454    }
455}
dark_std/sync/map_btree.rs

dark_std/sync/
map_btree.rs
