63 lines
1.5 KiB
Rust
63 lines
1.5 KiB
Rust
use collections::BTreeMap;
|
|
use core::mem;
|
|
use spin::Mutex;
|
|
|
|
use sync::WaitCondition;
|
|
|
|
#[derive(Debug)]
|
|
pub struct WaitMap<K, V> {
|
|
inner: Mutex<BTreeMap<K, V>>,
|
|
condition: WaitCondition
|
|
}
|
|
|
|
impl<K, V> WaitMap<K, V> where K: Clone + Ord {
|
|
pub fn new() -> WaitMap<K, V> {
|
|
WaitMap {
|
|
inner: Mutex::new(BTreeMap::new()),
|
|
condition: WaitCondition::new()
|
|
}
|
|
}
|
|
|
|
pub fn receive_nonblock(&self, key: &K) -> Option<V> {
|
|
self.inner.lock().remove(key)
|
|
}
|
|
|
|
pub fn receive(&self, key: &K) -> V {
|
|
loop {
|
|
if let Some(value) = self.receive_nonblock(key) {
|
|
return value;
|
|
}
|
|
self.condition.wait();
|
|
}
|
|
}
|
|
|
|
pub fn receive_any_nonblock(&self) -> Option<(K, V)> {
|
|
let mut inner = self.inner.lock();
|
|
if let Some(key) = inner.keys().next().map(|key| key.clone()) {
|
|
inner.remove(&key).map(|value| (key, value))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
pub fn receive_any(&self) -> (K, V) {
|
|
loop {
|
|
if let Some(entry) = self.receive_any_nonblock() {
|
|
return entry;
|
|
}
|
|
self.condition.wait();
|
|
}
|
|
}
|
|
|
|
pub fn receive_all(&self) -> BTreeMap<K, V> {
|
|
let mut ret = BTreeMap::new();
|
|
mem::swap(&mut ret, &mut *self.inner.lock());
|
|
ret
|
|
}
|
|
|
|
pub fn send(&self, key: K, value: V) {
|
|
self.inner.lock().insert(key, value);
|
|
self.condition.notify();
|
|
}
|
|
}
|