waitmap/
wait.rs

1use std::borrow::Borrow;
2use std::future::Future;
3use std::hash::{Hash, BuildHasher};
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7use dashmap::DashMap;
8
9use crate::WaitEntry;
10use crate::WaitEntry::*;
11use crate::{Ref, RefMut};
12
13pub struct Wait<'a, 'b, K, V, S, Q> where
14    K: Hash + Eq + Borrow<Q>,
15    S: BuildHasher + Clone,
16    Q: ?Sized + Hash + Eq,
17{
18    map: &'a DashMap<K, WaitEntry<V>, S>,
19    key: &'b Q,
20    idx: usize,
21}
22
23impl<'a, 'b, K, V, S, Q> Wait<'a, 'b, K, V, S, Q> where
24    K: Hash + Eq + Borrow<Q>,
25    S: BuildHasher + Clone,
26    Q: ?Sized + Hash + Eq,
27{
28    pub(crate) fn new(map: &'a DashMap<K, WaitEntry<V>, S>, key: &'b Q) -> Self {
29        Wait { map, key, idx: std::usize::MAX }
30    }
31}
32
33impl<'a, 'b, K, V, S, Q> Future for Wait<'a, 'b, K, V, S, Q> where
34    K: Hash + Eq + Borrow<Q>,
35    S: BuildHasher + Clone,
36    Q: ?Sized + Hash + Eq,
37{
38    type Output = Option<Ref<'a, K, V, S>>;
39
40    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
41        match self.map.get_mut(self.key) {
42            Some(mut entry) => match entry.value_mut() {
43                Waiting(wakers)  => {
44                    wakers.replace(ctx.waker().clone(), &mut self.idx);
45                    Poll::Pending
46                }
47                Filled(_)        => {
48                    let inner = entry.downgrade();
49                    self.idx = std::usize::MAX;
50                    Poll::Ready(Some(Ref { inner }))
51                }
52            }
53            None        => Poll::Ready(None),
54        }
55    }
56}
57
58impl<'a, 'b, K, V, S, Q> Drop for Wait<'a, 'b, K, V, S, Q> where
59    K: Hash + Eq + Borrow<Q>,
60    S: BuildHasher + Clone,
61    Q: ?Sized + Hash + Eq,
62{
63    fn drop(&mut self) {
64        if self.idx == std::usize::MAX { return; }
65        if let Some(mut entry) = self.map.get_mut(self.key) {
66            if let Waiting(wakers) = entry.value_mut() {
67                wakers.remove(self.idx);
68            }
69        }
70    }
71}
72
73pub struct WaitMut<'a, 'b, K, V, S, Q> where
74    K: Hash + Eq + Borrow<Q>,
75    S: BuildHasher + Clone,
76    Q: ?Sized + Hash + Eq,
77{
78    map: &'a DashMap<K, WaitEntry<V>, S>,
79    key: &'b Q,
80    idx: usize,
81}
82
83impl<'a, 'b, K, V, S, Q> WaitMut<'a, 'b, K, V, S, Q> where
84    K: Hash + Eq + Borrow<Q>,
85    S: BuildHasher + Clone,
86    Q: ?Sized + Hash + Eq,
87{
88    pub(crate) fn new(map: &'a DashMap<K, WaitEntry<V>, S>, key: &'b Q) -> Self {
89        WaitMut { map, key, idx: std::usize::MAX }
90    }
91}
92
93impl<'a, 'b, K, V, S, Q> Future for WaitMut<'a, 'b, K, V, S, Q> where
94    K: Hash + Eq + Borrow<Q>,
95    S: BuildHasher + Clone,
96    Q: ?Sized + Hash + Eq,
97{
98    type Output = Option<RefMut<'a, K, V, S>>;
99
100    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
101        match self.map.get_mut(self.key) {
102            Some(mut entry) => match entry.value_mut() {
103                Waiting(wakers)  => {
104                    wakers.replace(ctx.waker().clone(), &mut self.idx);
105                    Poll::Pending
106                }
107                Filled(_)        => {
108                    self.idx = std::usize::MAX;
109                    Poll::Ready(Some(RefMut { inner: entry }))
110                }
111            }
112            None        => Poll::Ready(None),
113        }
114    }
115}
116
117impl<'a, 'b, K, V, S, Q> Drop for WaitMut<'a, 'b, K, V, S, Q> where
118    K: Hash + Eq + Borrow<Q>,
119    S: BuildHasher + Clone,
120    Q: ?Sized + Hash + Eq,
121{
122    fn drop(&mut self) {
123        if self.idx == std::usize::MAX { return; }
124        if let Some(mut entry) = self.map.get_mut(self.key) {
125            if let Waiting(wakers) = entry.value_mut() {
126                wakers.remove(self.idx);
127            }
128        }
129    }
130}
waitmap/wait.rs

waitmap/
wait.rs
