dusk_consensus/
queue.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4//
5// Copyright (c) DUSK NETWORK. All rights reserved.
6
7use std::collections::{BTreeMap, VecDeque};
8use std::fmt::Debug;
9
10use node_data::message::Message;
11use thiserror::Error;
12use tracing::warn;
13
14type StepMap<T> = BTreeMap<u8, VecDeque<T>>;
15type RoundMap<T> = BTreeMap<u64, StepMap<T>>;
16
17const MAX_MESSAGES_PER_QUEUE: usize = 1000;
18
19#[derive(Debug, Default)]
20pub struct MsgRegistry<T: QueueMessage>(RoundMap<T>);
21
22pub trait QueueMessage: Debug + Clone {
23    fn step(&self) -> u8;
24
25    fn round(&self) -> u64;
26
27    fn signer(&self) -> Option<node_data::bls::PublicKeyBytes>;
28}
29
30impl QueueMessage for Message {
31    fn round(&self) -> u64 {
32        self.header.round
33    }
34    fn step(&self) -> u8 {
35        self.get_step()
36    }
37    fn signer(&self) -> Option<node_data::bls::PublicKeyBytes> {
38        self.get_signer().map(|s| *s.bytes())
39    }
40}
41
42#[derive(Debug, Error)]
43pub enum MsgRegistryError<T> {
44    #[error("Msg already enqueued")]
45    SignerAlreadyEnqueue(T),
46    #[error("This msg has no signer")]
47    NoSigner(T),
48}
49
50/// A message registry that stores messages based on their round and step.
51impl<T: QueueMessage> MsgRegistry<T> {
52    /// Inserts a message into the registry based on its round and step.
53    pub fn put_msg(&mut self, msg: T) -> Result<T, MsgRegistryError<T>> {
54        let round = msg.round();
55        let step = msg.step();
56        let vec = self
57            .0
58            .entry(round)
59            .or_default()
60            .entry(step)
61            .or_insert(VecDeque::with_capacity(MAX_MESSAGES_PER_QUEUE));
62        if msg.signer().is_none() {
63            return Err(MsgRegistryError::NoSigner(msg));
64        }
65        if vec.iter().any(|m| m.signer() == msg.signer()) {
66            return Err(MsgRegistryError::SignerAlreadyEnqueue(msg));
67        }
68
69        if vec.len() == vec.capacity() {
70            warn!("queue ({}, {}) is full, dropping", round, step);
71            vec.pop_front();
72        }
73
74        let ret = msg.clone();
75        vec.push_back(msg);
76        Ok(ret)
77    }
78
79    /// Drains and returns all messages that belong to the specified round and
80    /// step.
81    pub fn drain_msg_by_round_step(
82        &mut self,
83        round: u64,
84        step: u8,
85    ) -> Option<VecDeque<T>> {
86        self.0
87            .get_mut(&round)
88            .and_then(|r| r.remove_entry(&step).map(|(_, v)| v))
89    }
90
91    /// Removes all messages that belong to the specified round.
92    pub fn remove_msgs_by_round(&mut self, round: u64) {
93        if let Some(r) = self.0.get_mut(&round) {
94            r.clear();
95        };
96
97        self.0.remove(&round);
98    }
99
100    /// Removes all messages that do not belong to the range (closed interval)
101    /// of keys
102    pub fn remove_msgs_out_of_range(&mut self, start_round: u64, offset: u64) {
103        let end_round = start_round + offset;
104
105        self.0 = self
106            .0
107            .split_off(&start_round)
108            .into_iter()
109            .filter(|(k, _)| *k <= end_round)
110            .collect();
111    }
112
113    /// Returns the total number of messages in the registry.
114    pub fn msg_count(&self) -> usize {
115        self.0
116            .values()
117            .map(|round| round.values().map(|items| items.len()).sum::<usize>())
118            .sum()
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use node_data::bls::PUBLIC_BLS_SIZE;
125
126    use super::QueueMessage;
127    use crate::queue::MsgRegistry;
128
129    #[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
130    struct Item(u64, u8, i32, node_data::bls::PublicKeyBytes);
131
132    impl Item {
133        fn new(round: u64, step: u8, data: i32) -> Self {
134            let mut buf = [0u8; PUBLIC_BLS_SIZE];
135            let data_bytes = data.to_le_bytes();
136
137            buf[0] = data_bytes[0];
138            buf[1] = data_bytes[1];
139            buf[2] = data_bytes[2];
140            buf[3] = data_bytes[3];
141            Self(round, step, data, node_data::bls::PublicKeyBytes(buf))
142        }
143    }
144
145    impl QueueMessage for Item {
146        fn round(&self) -> u64 {
147            self.0
148        }
149        fn step(&self) -> u8 {
150            self.1
151        }
152        fn signer(&self) -> Option<node_data::bls::PublicKeyBytes> {
153            Some(self.3)
154        }
155    }
156    #[test]
157    fn test_push_event() -> Result<(), super::MsgRegistryError<Item>> {
158        let round = 55555;
159
160        let mut reg = MsgRegistry::<Item>::default();
161        reg.put_msg(Item::new(round, 2, 5))?;
162        reg.put_msg(Item::new(round, 2, 4))?;
163        reg.put_msg(Item::new(round, 2, 3))?;
164
165        assert_eq!(reg.msg_count(), 3);
166        assert!(reg.drain_msg_by_round_step(round, 3).is_none());
167        assert!(reg.drain_msg_by_round_step(4444, 2).is_none());
168
169        for i in 1..100 {
170            reg.put_msg(Item::new(4444, i as u8, i))?;
171        }
172
173        assert_eq!(reg.msg_count(), 100 + 2);
174        assert_eq!(
175            reg.drain_msg_by_round_step(round, 2).unwrap(),
176            vec![
177                Item::new(round, 2, 5),
178                Item::new(round, 2, 4),
179                Item::new(round, 2, 3)
180            ],
181        );
182        assert_eq!(reg.msg_count(), 99);
183
184        reg.remove_msgs_by_round(4444);
185        assert_eq!(reg.msg_count(), 0);
186        assert!(reg.drain_msg_by_round_step(round, 2).is_none());
187        Ok(())
188    }
189
190    #[test]
191    fn test_remove_msgs_out_of_range()
192    -> Result<(), super::MsgRegistryError<Item>> {
193        let round = 100;
194
195        let mut reg = MsgRegistry::<Item>::default();
196        reg.put_msg(Item::new(round + 1, 1, 1))?;
197        reg.put_msg(Item::new(round + 2, 1, 1))?;
198        reg.put_msg(Item::new(round * 3, 1, 1))?;
199        reg.put_msg(Item::new(round, 1, 1))?;
200        assert_eq!(reg.msg_count(), 4);
201
202        reg.remove_msgs_out_of_range(round + 1, 1);
203        assert_eq!(reg.msg_count(), 2);
204
205        assert!(reg.drain_msg_by_round_step(round, 1).is_none());
206        assert!(reg.drain_msg_by_round_step(round * 3, 1).is_none());
207
208        assert!(reg.drain_msg_by_round_step(round + 1, 1).is_some());
209        assert!(reg.drain_msg_by_round_step(round + 2, 1).is_some());
210        Ok(())
211    }
212}
dusk_consensus/queue.rs

dusk_consensus/
queue.rs
