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

dark_std/sync/
map_btree.rs
