cliquenet/
delay.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    mem,
4    sync::Arc,
5    time::Duration,
6};
7
8use bytes::Bytes;
9use parking_lot::Mutex;
10use tokio::time::Instant;
11
12use crate::{
13    Config, RetryPolicy,
14    msg::{MsgId, Slot},
15};
16
17#[derive(Clone, Debug)]
18pub struct DelayQueue {
19    conf: Arc<Config>,
20    inner: Arc<Mutex<Inner>>,
21}
22
23#[derive(Debug)]
24struct Inner {
25    map: BTreeMap<(Slot, MsgId), Entry>,
26    due: BTreeSet<(Instant, Slot, MsgId)>,
27}
28
29#[derive(Debug)]
30struct Entry {
31    msg: Bytes,
32    pol: RetryPolicy,
33    num: usize,
34    due: Instant,
35}
36
37impl DelayQueue {
38    pub fn new(conf: Arc<Config>) -> Self {
39        Self {
40            conf,
41            inner: Arc::new(Mutex::new(Inner {
42                map: BTreeMap::new(),
43                due: BTreeSet::new(),
44            })),
45        }
46    }
47
48    pub fn add(&self, s: Slot, i: MsgId, msg: Bytes, pol: RetryPolicy, now: Instant) {
49        let num = 0;
50        let due = timeout(&self.conf, now, num);
51        let mut inner = self.inner.lock();
52        inner.map.insert((s, i), Entry { msg, pol, num, due });
53        inner.due.insert((due, s, i));
54    }
55
56    pub fn due(&self, now: Instant) -> Option<(Bytes, RetryPolicy)> {
57        let mut inner = self.inner.lock();
58        while let Some(&(due, slot, id)) = inner.due.first() {
59            if due > now {
60                return None;
61            }
62            let _ = inner.due.pop_first();
63            let Some(entry) = inner.map.get_mut(&(slot, id)) else {
64                continue;
65            };
66            entry.num = entry.num.saturating_add(1);
67            let t = timeout(&self.conf, now, entry.num);
68            entry.due = t;
69            let item = (entry.msg.clone(), entry.pol);
70            inner.due.insert((t, slot, id));
71            return Some(item);
72        }
73        None
74    }
75
76    pub fn remove(&self, s: Slot, i: MsgId) {
77        let mut inner = self.inner.lock();
78        if let Some(entry) = inner.map.remove(&(s, i)) {
79            inner.due.remove(&(entry.due, s, i));
80        }
81    }
82
83    pub fn reset(&self, now: Instant) {
84        let mut inner = self.inner.lock();
85        inner.due.clear();
86        let mut due = mem::take(&mut inner.due);
87        for (&(slot, id), entry) in &mut inner.map {
88            entry.num = 0;
89            entry.due = now;
90            due.insert((now, slot, id));
91        }
92        inner.due = due
93    }
94
95    pub fn gc(&self, s: Slot) {
96        let mut inner = self.inner.lock();
97        inner.map = inner.map.split_off(&(s, MsgId(0)));
98    }
99
100    pub fn len(&self) -> usize {
101        self.inner.lock().map.len()
102    }
103
104    pub fn is_due(&self, now: Instant) -> bool {
105        let inner = self.inner.lock();
106        let Some(&(due, ..)) = inner.due.first() else {
107            return false;
108        };
109        due <= now
110    }
111}
112
113fn timeout(cfg: &Config, now: Instant, at: usize) -> Instant {
114    let d = *cfg
115        .send_retry_delays
116        .get(at)
117        .unwrap_or_else(|| cfg.send_retry_delays.last());
118    now + Duration::from_secs(d.into())
119}
cliquenet/delay.rs

cliquenet/
delay.rs
