cliquenet/
queue.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use parking_lot::Mutex;
4use tokio::sync::Notify;
5
6use crate::msg::{MsgId, Slot};
7
8#[derive(Debug)]
9pub struct Queue<T>(Arc<Inner<T>>);
10
11#[derive(Debug)]
12struct Inner<T> {
13    sig: Notify,
14    map: Mutex<BTreeMap<(Slot, MsgId), T>>,
15}
16
17impl<T> Default for Queue<T> {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl<T> Clone for Queue<T> {
24    fn clone(&self) -> Self {
25        Self(self.0.clone())
26    }
27}
28
29impl<T> Queue<T> {
30    pub fn new() -> Self {
31        Self(Arc::new(Inner {
32            sig: Notify::new(),
33            map: Mutex::new(BTreeMap::new()),
34        }))
35    }
36
37    pub fn enqueue(&self, s: Slot, i: MsgId, val: T) {
38        self.0.map.lock().insert((s, i), val);
39        self.0.sig.notify_waiters();
40    }
41
42    pub fn gc(&self, s: Slot) {
43        let mut map = self.0.map.lock();
44        *map = map.split_off(&(s, MsgId(0)))
45    }
46
47    pub fn remove(&self, s: Slot, i: MsgId) {
48        self.0.map.lock().remove(&(s, i));
49    }
50
51    pub fn len(&self) -> usize {
52        self.0.map.lock().len()
53    }
54}
55
56impl<T: Clone> Queue<T> {
57    pub fn try_next(&self) -> Option<(Slot, MsgId, T)> {
58        let map = self.0.map.lock();
59        let (&(s, i), v) = map.first_key_value()?;
60        Some((s, i, v.clone()))
61    }
62
63    pub async fn next(&self) -> (Slot, MsgId, T) {
64        loop {
65            let future = self.0.sig.notified();
66            if let Some(v) = self.try_next() {
67                return v;
68            }
69            future.await;
70        }
71    }
72}
cliquenet/queue.rs

cliquenet/
queue.rs
