├── .github └── workflows │ └── rust.yml ├── .gitignore ├── Cargo.toml ├── README.md ├── src ├── lib.rs ├── wait.rs └── waker_set.rs └── tests └── smoke.rs /.github/workflows/rust.yml: -------------------------------------------------------------------------------- 1 | name: Rust 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | branches: [ master ] 8 | 9 | jobs: 10 | build: 11 | 12 | runs-on: ubuntu-latest 13 | 14 | steps: 15 | - uses: actions/checkout@v2 16 | - name: Build 17 | run: cargo build --verbose 18 | - name: Run tests 19 | run: cargo test --verbose 20 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "waitmap" 3 | description = "an awaitable concurrent hash map" 4 | documentation = "https://docs.rs/waitmap" 5 | repository = "https://github.com/withoutboats/waitmap" 6 | license = "MIT OR Apache-2.0" 7 | version = "1.1.0" 8 | authors = ["Without Boats "] 9 | edition = "2018" 10 | 11 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 12 | 13 | [dependencies] 14 | smallvec = "1.2.0" 15 | dashmap = "3.7.0" 16 | 17 | [dev-dependencies.async-std] 18 | version = "1.5.0" 19 | features = ["unstable", "attributes"] 20 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # waitmap 2 | 3 | [![waitmap](https://docs.rs/waitmap/badge.svg)](https://docs.rs/waitmap/) 4 | [![version](https://img.shields.io/crates/v/waitmap)](https://crates.io/crates/waitmap/) 5 | 6 | Wait Map is an async/await concurrency primitive implemented as a concurrent hashmap. It is built 7 | on top of the [dashmap](https://github.com/xacrimon/dashmap) concurrent hashmap, with an additional "wait" API. 8 | 9 | The wait API lets users wait on one task for an entry to be filled by another task. For example: 10 | 11 | ```rust 12 | let map: WaitMap; 13 | 14 | // This will wait until a value is put under the key "Rosa Luxemburg" 15 | if let Some(value) = map.wait("Rosa Luxemburg").await { 16 | // ... 17 | } 18 | ``` 19 | 20 | It also supports a cancellation API, to cause any task waiting on an entry being filled to stop 21 | waiting (the future evaluating to `None`, just as if they had called `get` and the key was empty): 22 | 23 | ```rust 24 | // This will cause the other task to stop waiting, it receives a `None` value: 25 | map.cancel("Rosa Luxemburg"); 26 | ``` 27 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! Async concurrent hashmap built on top of [dashmap](https://docs.rs/dashmap/). 2 | //! 3 | //! # Wait 4 | //! [`WaitMap`](crate::WaitMap) is a concurrent hashmap with an asynchronous `wait` operation. 5 | //! ``` 6 | //! # extern crate async_std; 7 | //! # extern crate waitmap; 8 | //! # use async_std::main; 9 | //! # use waitmap::WaitMap; 10 | //! # #[async_std::main] 11 | //! # async fn main() -> std::io::Result<()> { 12 | //! let map: WaitMap = WaitMap::new(); 13 | //! # map.insert(String::from("Rosa Luxemburg"), 1); 14 | //! 15 | //! // This will wait until a value is put under the key "Rosa Luxemburg" 16 | //! if let Some(value) = map.wait("Rosa Luxemburg").await { 17 | //! // ... 18 | //! } 19 | //! # Ok(()) 20 | //! # } 21 | //! ``` 22 | //! 23 | //! Waits are cancellable. Cancelled waits evaluate to `None`. 24 | //! ``` 25 | //! # extern crate async_std; 26 | //! # extern crate waitmap; 27 | //! # use async_std::{main, task}; 28 | //! # use std::time::Duration; 29 | //! # use std::sync::Arc; 30 | //! # use waitmap::WaitMap; 31 | //! # #[async_std::main] 32 | //! # async fn main() -> std::io::Result<()> { 33 | //! let map: Arc> = Arc::new(WaitMap::new()); 34 | //! let map1 = map.clone(); 35 | //! 36 | //! let handle = task::spawn(async move { 37 | //! let result = map.wait("Voltairine de Cleyre").await; 38 | //! assert!(result.is_none()); 39 | //! }); 40 | //! 41 | //! task::spawn(async move { 42 | //! task::sleep(Duration::from_millis(100)).await; // avoid deadlock 43 | //! map1.cancel("Voltairine de Cleyre"); 44 | //! }); 45 | //! 46 | //! task::block_on(handle); 47 | //! # Ok(()) 48 | //! # } 49 | //! ``` 50 | 51 | mod wait; 52 | mod waker_set; 53 | 54 | use std::borrow::Borrow; 55 | use std::collections::hash_map::RandomState; 56 | use std::future::Future; 57 | use std::hash::{Hash, BuildHasher}; 58 | use std::mem; 59 | 60 | use dashmap::DashMap; 61 | use dashmap::mapref::entry::Entry::*; 62 | use dashmap::mapref::one; 63 | 64 | use WaitEntry::*; 65 | use wait::{Wait, WaitMut}; 66 | use waker_set::WakerSet; 67 | 68 | /// An asynchronous concurrent hashmap. 69 | pub struct WaitMap { 70 | map: DashMap, S>, 71 | } 72 | 73 | impl WaitMap { 74 | /// Make a new `WaitMap` using the default hasher. 75 | pub fn new() -> WaitMap { 76 | WaitMap { map: DashMap::with_hasher(RandomState::default()) } 77 | } 78 | } 79 | 80 | impl WaitMap { 81 | /// Make a new `WaitMap` using a custom hasher. 82 | /// ``` 83 | /// # extern crate async_std; 84 | /// # extern crate waitmap; 85 | /// # use async_std::main; 86 | /// # use waitmap::WaitMap; 87 | /// use std::collections::hash_map::RandomState; 88 | /// # #[async_std::main] 89 | /// # async fn main() -> std::io::Result<()> { 90 | /// let map: WaitMap = WaitMap::with_hasher(RandomState::new()); 91 | /// # Ok(()) 92 | /// # } 93 | /// ``` 94 | pub fn with_hasher(hasher: S) -> WaitMap { 95 | WaitMap { map: DashMap::with_hasher(hasher) } 96 | } 97 | 98 | /// Inserts a key-value pair into the map. 99 | /// 100 | /// If the map did not have this key present, `None` is returned. 101 | /// 102 | /// If there are any pending `wait` calls for this key, they are woken up. 103 | /// 104 | /// If the map did have this key present, the value is updated and the old value is returned. 105 | /// ``` 106 | /// # extern crate async_std; 107 | /// # extern crate waitmap; 108 | /// # use async_std::{main, sync::Arc, prelude::*}; 109 | /// # use waitmap::WaitMap; 110 | /// # #[async_std::main] 111 | /// # async fn main() -> std::io::Result<()> { 112 | /// let map: Arc> = Arc::new(WaitMap::new()); 113 | /// 114 | /// let insert_fut = async { map.insert("hi".to_string(), 0) }; 115 | /// let wait_fut = map.wait("hi"); 116 | /// 117 | /// let (insert_res, wait_res) = insert_fut.join(wait_fut).await; 118 | /// assert!(insert_res.is_none()); 119 | /// assert!(wait_res.is_some()); 120 | /// # Ok(()) 121 | /// # } 122 | /// ``` 123 | pub fn insert(&self, key: K, value: V) -> Option { 124 | match self.map.entry(key) { 125 | Occupied(mut entry) => { 126 | match mem::replace(entry.get_mut(), Filled(value)) { 127 | Waiting(wakers) => { 128 | drop(entry); // drop early to release lock before waking other tasks 129 | wakers.wake(); 130 | None 131 | } 132 | Filled(value) => Some(value), 133 | } 134 | } 135 | Vacant(slot) => { 136 | slot.insert(Filled(value)); 137 | None 138 | } 139 | } 140 | } 141 | 142 | pub fn get(&self, key: &Q) -> Option> 143 | where K: Borrow 144 | { 145 | Some(Ref { inner: self.map.get(key)? }) 146 | } 147 | 148 | pub fn get_mut(&self, key: &Q) -> Option> 149 | where K: Borrow 150 | { 151 | Some(RefMut { inner: self.map.get_mut(key)? }) 152 | } 153 | 154 | pub fn wait<'a: 'f, 'b: 'f, 'f, Q: ?Sized + Hash + Eq>(&'a self, qey: &'b Q) 155 | -> impl Future>> + 'f 156 | where 157 | K: Borrow + From<&'b Q>, 158 | { 159 | let key = K::from(qey); 160 | self.map.entry(key).or_insert(Waiting(WakerSet::new())); 161 | Wait::new(&self.map, qey) 162 | } 163 | 164 | pub fn wait_mut<'a: 'f, 'b: 'f, 'f, Q: ?Sized + Hash + Eq>(&'a self, qey: &'b Q) 165 | -> impl Future>> + 'f 166 | where 167 | K: Borrow + From<&'b Q>, 168 | { 169 | let key = K::from(qey); 170 | self.map.entry(key).or_insert(Waiting(WakerSet::new())); 171 | WaitMut::new(&self.map, qey) 172 | } 173 | 174 | pub fn cancel(&self, key: &Q) -> bool 175 | where K: Borrow 176 | { 177 | if let Some((_, entry)) = self.map.remove_if(key, |_, entry| { 178 | if let Waiting(_) = entry { true } else { false } 179 | }) { 180 | if let Waiting(wakers) = entry { 181 | wakers.wake(); 182 | } 183 | true 184 | } else { false } 185 | } 186 | 187 | /// Cancels all outstanding `waits` on the map. 188 | /// ``` 189 | /// # extern crate async_std; 190 | /// # extern crate waitmap; 191 | /// # use async_std::{main, stream, prelude::*}; 192 | /// # use waitmap::WaitMap; 193 | /// # #[async_std::main] 194 | /// # async fn main() -> std::io::Result<()> { 195 | /// let map: WaitMap = WaitMap::new(); 196 | /// let mut waitstream = 197 | /// stream::from_iter(vec![map.wait("we"), map.wait("are"), map.wait("waiting")]); 198 | /// 199 | /// map.cancel_all(); 200 | /// 201 | /// let mut num_cancelled = 0; 202 | /// while let Some(wait_fut) = waitstream.next().await { 203 | /// assert!(wait_fut.await.is_none()); 204 | /// num_cancelled += 1; 205 | /// } 206 | /// 207 | /// assert!(num_cancelled == 3); 208 | /// # Ok(()) 209 | /// # } 210 | /// ``` 211 | pub fn cancel_all(&self) { 212 | self.map.retain(|_, entry| { 213 | if let Waiting(wakers) = entry { 214 | // NB: In theory, there is a deadlock risk: if a task is awoken before the 215 | // retain is completed, it may see a waiting entry with an empty waker set, 216 | // rather than a missing entry. 217 | // 218 | // However, this is prevented by the memory guards already present in DashMap. 219 | // No other task will be able to view this entry until the guard on this shard 220 | // has been dropped, which will not occur until this shard's unretained members 221 | // have actually been removed. 222 | mem::replace(wakers, WakerSet::new()).wake(); 223 | false 224 | } else { true } 225 | }) 226 | } 227 | } 228 | 229 | enum WaitEntry { 230 | Waiting(WakerSet), 231 | Filled(V), 232 | } 233 | 234 | /// A shared reference to a `WaitMap` key-value pair. 235 | /// ``` 236 | /// # extern crate async_std; 237 | /// # extern crate waitmap; 238 | /// # use async_std::main; 239 | /// # use waitmap::{Ref, WaitMap}; 240 | /// # #[async_std::main] 241 | /// # async fn main() -> std::io::Result<()> { 242 | /// let map: WaitMap = WaitMap::new(); 243 | /// let emma = "Emma Goldman".to_string(); 244 | /// 245 | /// map.insert(emma.clone(), 0); 246 | /// let kv: Ref = map.get(&emma).unwrap(); 247 | /// 248 | /// assert!(*kv.key() == emma); 249 | /// assert!(*kv.value() == 0); 250 | /// assert!(kv.pair() == (&"Emma Goldman".to_string(), &0)); 251 | /// # Ok(()) 252 | /// # } 253 | /// ``` 254 | pub struct Ref<'a, K, V, S> { 255 | inner: one::Ref<'a, K, WaitEntry, S>, 256 | } 257 | 258 | impl<'a, K: Eq + Hash, V, S: BuildHasher> Ref<'a, K, V, S> { 259 | pub fn key(&self) -> &K { 260 | self.inner.key() 261 | } 262 | 263 | pub fn value(&self) -> &V { 264 | match self.inner.value() { 265 | Filled(value) => value, 266 | _ => panic!() 267 | } 268 | } 269 | 270 | pub fn pair(&self) -> (&K, &V) { 271 | (self.key(), self.value()) 272 | } 273 | } 274 | 275 | /// An exclusive reference to a `WaitMap` key-value pair. 276 | pub struct RefMut<'a, K, V, S> { 277 | inner: one::RefMut<'a, K, WaitEntry, S>, 278 | } 279 | 280 | impl<'a, K: Eq + Hash, V, S: BuildHasher> RefMut<'a, K, V, S> { 281 | pub fn key(&self) -> &K { 282 | self.inner.key() 283 | } 284 | 285 | pub fn value(&self) -> &V { 286 | match self.inner.value() { 287 | Filled(value) => value, 288 | _ => panic!() 289 | } 290 | } 291 | 292 | pub fn value_mut(&mut self) -> &mut V { 293 | match self.inner.value_mut() { 294 | Filled(value) => value, 295 | _ => panic!() 296 | } 297 | } 298 | 299 | pub fn pair(&self) -> (&K, &V) { 300 | (self.key(), self.value()) 301 | } 302 | 303 | pub fn pair_mut(&mut self) -> (&K, &mut V) { 304 | match self.inner.pair_mut() { 305 | (key, Filled(value)) => (key, value), 306 | _ => panic!(), 307 | } 308 | } 309 | } 310 | -------------------------------------------------------------------------------- /src/wait.rs: -------------------------------------------------------------------------------- 1 | use std::borrow::Borrow; 2 | use std::future::Future; 3 | use std::hash::{Hash, BuildHasher}; 4 | use std::pin::Pin; 5 | use std::task::{Context, Poll}; 6 | 7 | use dashmap::DashMap; 8 | 9 | use crate::WaitEntry; 10 | use crate::WaitEntry::*; 11 | use crate::{Ref, RefMut}; 12 | 13 | pub struct Wait<'a, 'b, K, V, S, Q> where 14 | K: Hash + Eq + Borrow, 15 | S: BuildHasher + Clone, 16 | Q: ?Sized + Hash + Eq, 17 | { 18 | map: &'a DashMap, S>, 19 | key: &'b Q, 20 | idx: usize, 21 | } 22 | 23 | impl<'a, 'b, K, V, S, Q> Wait<'a, 'b, K, V, S, Q> where 24 | K: Hash + Eq + Borrow, 25 | S: BuildHasher + Clone, 26 | Q: ?Sized + Hash + Eq, 27 | { 28 | pub(crate) fn new(map: &'a DashMap, S>, key: &'b Q) -> Self { 29 | Wait { map, key, idx: std::usize::MAX } 30 | } 31 | } 32 | 33 | impl<'a, 'b, K, V, S, Q> Future for Wait<'a, 'b, K, V, S, Q> where 34 | K: Hash + Eq + Borrow, 35 | S: BuildHasher + Clone, 36 | Q: ?Sized + Hash + Eq, 37 | { 38 | type Output = Option>; 39 | 40 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 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 | 58 | impl<'a, 'b, K, V, S, Q> Drop for Wait<'a, 'b, K, V, S, Q> where 59 | K: Hash + Eq + Borrow, 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 | 73 | pub struct WaitMut<'a, 'b, K, V, S, Q> where 74 | K: Hash + Eq + Borrow, 75 | S: BuildHasher + Clone, 76 | Q: ?Sized + Hash + Eq, 77 | { 78 | map: &'a DashMap, S>, 79 | key: &'b Q, 80 | idx: usize, 81 | } 82 | 83 | impl<'a, 'b, K, V, S, Q> WaitMut<'a, 'b, K, V, S, Q> where 84 | K: Hash + Eq + Borrow, 85 | S: BuildHasher + Clone, 86 | Q: ?Sized + Hash + Eq, 87 | { 88 | pub(crate) fn new(map: &'a DashMap, S>, key: &'b Q) -> Self { 89 | WaitMut { map, key, idx: std::usize::MAX } 90 | } 91 | } 92 | 93 | impl<'a, 'b, K, V, S, Q> Future for WaitMut<'a, 'b, K, V, S, Q> where 94 | K: Hash + Eq + Borrow, 95 | S: BuildHasher + Clone, 96 | Q: ?Sized + Hash + Eq, 97 | { 98 | type Output = Option>; 99 | 100 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 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 | 117 | impl<'a, 'b, K, V, S, Q> Drop for WaitMut<'a, 'b, K, V, S, Q> where 118 | K: Hash + Eq + Borrow, 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 | } 131 | -------------------------------------------------------------------------------- /src/waker_set.rs: -------------------------------------------------------------------------------- 1 | use std::task::Waker; 2 | 3 | use smallvec::SmallVec; 4 | 5 | pub struct WakerSet { 6 | wakers: SmallVec<[Option; 1]>, 7 | } 8 | 9 | impl WakerSet { 10 | pub fn new() -> WakerSet { 11 | WakerSet { 12 | wakers: SmallVec::new(), 13 | } 14 | } 15 | 16 | pub fn replace(&mut self, waker: Waker, idx: &mut usize) { 17 | let len = self.wakers.len(); 18 | if *idx >= len { 19 | debug_assert!(len != std::usize::MAX); // usize::MAX is used as a sentinel 20 | *idx = len; 21 | self.wakers.push(Some(waker)); 22 | } else { 23 | self.wakers[*idx] = Some(waker); 24 | } 25 | } 26 | 27 | pub fn remove(&mut self, idx: usize) { 28 | self.wakers[idx] = None; 29 | } 30 | 31 | pub fn wake(self) { 32 | for waker in self.wakers { 33 | if let Some(waker) = waker { 34 | waker.wake() 35 | } 36 | } 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /tests/smoke.rs: -------------------------------------------------------------------------------- 1 | use std::sync::Arc; 2 | use std::time::Duration; 3 | 4 | use waitmap::WaitMap; 5 | 6 | use async_std::task; 7 | 8 | #[test] 9 | fn works_like_a_normal_map() { 10 | let map = WaitMap::new(); 11 | assert!(map.get("Rosa Luxemburg").is_none()); 12 | map.insert(String::from("Rosa Luxemburg"), 0); 13 | assert_eq!(map.get("Rosa Luxemburg").unwrap().value(), &0); 14 | assert!(map.get("Voltairine de Cleyre").is_none()); 15 | } 16 | 17 | #[test] 18 | fn simple_waiting() { 19 | let map: Arc> = Arc::new(WaitMap::new()); 20 | let map2 = map.clone(); 21 | 22 | let handle = task::spawn(async move { 23 | let rosa = map.wait("Rosa Luxemburg").await; 24 | assert_eq!(rosa.unwrap().value(), &0); 25 | assert!(map.wait("Voltairine de Cleyre").await.is_none()); 26 | }); 27 | 28 | task::spawn(async move { 29 | task::sleep(Duration::from_millis(140)).await; 30 | map2.insert(String::from("Rosa Luxemburg"), 0); 31 | task::sleep(Duration::from_millis(140)).await; 32 | map2.cancel("Voltairine de Cleyre"); 33 | }); 34 | 35 | task::block_on(handle); 36 | } 37 | 38 | #[test] 39 | fn simple_waiting_mut() { 40 | let map: Arc> = Arc::new(WaitMap::new()); 41 | let map2 = map.clone(); 42 | 43 | let handle = task::spawn(async move { 44 | let rosa = map.wait_mut("Rosa Luxemburg").await; 45 | assert_eq!(rosa.unwrap().value(), &0); 46 | assert!(map.wait_mut("Voltairine de Cleyre").await.is_none()); 47 | }); 48 | 49 | task::spawn(async move { 50 | task::sleep(Duration::from_millis(140)).await; 51 | map2.insert(String::from("Rosa Luxemburg"), 0); 52 | task::sleep(Duration::from_millis(140)).await; 53 | map2.cancel("Voltairine de Cleyre"); 54 | }); 55 | 56 | task::block_on(handle); 57 | } 58 | 59 | #[test] 60 | fn cancel_all_cancels_all() { 61 | let map: Arc> = Arc::new(WaitMap::new()); 62 | let map2 = map.clone(); 63 | 64 | let handle = task::spawn(async move { 65 | let rosa = map.wait("Rosa Luxemburg"); 66 | let voltairine = map.wait("Voltairine de Cleyre"); 67 | assert!(rosa.await.is_none()); 68 | assert!(voltairine.await.is_none()); 69 | }); 70 | 71 | task::spawn(async move { 72 | task::sleep(Duration::from_millis(140)).await; 73 | map2.cancel_all(); 74 | }); 75 | 76 | task::block_on(handle); 77 | } 78 | 79 | #[test] 80 | fn multiple_tasks_can_wait_one_key() { 81 | let map: Arc> = Arc::new(WaitMap::new()); 82 | let map1 = map.clone(); 83 | let map2 = map.clone(); 84 | 85 | task::spawn(async move { 86 | map.insert(String::from("Rosa Luxemburg"), 0); 87 | }); 88 | 89 | let handle1 = task::spawn(async move { 90 | let rosa = map1.wait("Rosa Luxemburg").await; 91 | assert_eq!(rosa.unwrap().value(), &0); 92 | }); 93 | 94 | let handle2 = task::spawn(async move { 95 | let rosa = map2.wait("Rosa Luxemburg").await; 96 | assert_eq!(rosa.unwrap().value(), &0); 97 | }); 98 | 99 | task::block_on(handle1); 100 | task::block_on(handle2); 101 | } 102 | --------------------------------------------------------------------------------