sim: Track messages received in simulator

This is going to be useful for testing message amplification in the
network protocol.

We also reset the state of a few things everytime the simulator is run.
This commit is contained in:
cloudhead 2024-02-09 11:31:02 +01:00
parent 286c639e26
commit d1f4161ee9
No known key found for this signature in database
1 changed files with 15 additions and 2 deletions

View File

@ -63,7 +63,7 @@ pub enum Input {
}, },
/// Disconnected from peer. /// Disconnected from peer.
Disconnected(NodeId, Rc<DisconnectReason>), Disconnected(NodeId, Rc<DisconnectReason>),
/// Received a message from a remote peer. /// Received messages from a remote peer.
Received(NodeId, Vec<Message>), Received(NodeId, Vec<Message>),
/// Fetch completed for a node. /// Fetch completed for a node.
Fetched(RepoId, NodeId, Rc<Result<fetch::FetchResult, FetchError>>), Fetched(RepoId, NodeId, Rc<Result<fetch::FetchResult, FetchError>>),
@ -175,6 +175,8 @@ pub struct Simulation<S, G> {
inbox: Inbox, inbox: Inbox,
/// Events emitted during simulation. /// Events emitted during simulation.
events: BTreeMap<NodeId, VecDeque<Event>>, events: BTreeMap<NodeId, VecDeque<Event>>,
/// Messages received during simulation.
messages: Vec<Message>,
/// Priority events that should happen immediately. /// Priority events that should happen immediately.
priority: VecDeque<Scheduled>, priority: VecDeque<Scheduled>,
/// Simulated latencies between nodes. /// Simulated latencies between nodes.
@ -207,6 +209,7 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
messages: BTreeMap::new(), messages: BTreeMap::new(),
}, },
events: BTreeMap::new(), events: BTreeMap::new(),
messages: Vec::new(),
priority: VecDeque::new(), priority: VecDeque::new(),
partitions: BTreeSet::new(), partitions: BTreeSet::new(),
latencies: BTreeMap::new(), latencies: BTreeMap::new(),
@ -246,6 +249,11 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
self.events.entry(*node).or_default().drain(..) self.events.entry(*node).or_default().drain(..)
} }
/// Get all messages received by nodes during the simulation.
pub fn messages(&mut self) -> &[Message] {
&self.messages
}
/// Get the latency between two nodes. The minimum latency between nodes is 1 millisecond. /// Get the latency between two nodes. The minimum latency between nodes is 1 millisecond.
pub fn latency(&self, from: NodeId, to: NodeId) -> LocalDuration { pub fn latency(&self, from: NodeId, to: NodeId) -> LocalDuration {
self.latencies self.latencies
@ -294,6 +302,10 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
{ {
let mut nodes: BTreeMap<_, _> = peers.into_iter().map(|p| (p.id(), p)).collect(); let mut nodes: BTreeMap<_, _> = peers.into_iter().map(|p| (p.id(), p)).collect();
self.messages.clear();
self.events.clear();
self.start_time = self.time;
while self.step_(&mut nodes) { while self.step_(&mut nodes) {
if !pred(self) { if !pred(self) {
break; break;
@ -406,9 +418,10 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
} }
Input::Wake => p.wake(), Input::Wake => p.wake(),
Input::Received(id, msgs) => { Input::Received(id, msgs) => {
for msg in msgs { for msg in msgs.clone() {
p.received_message(id, msg); p.received_message(id, msg);
} }
self.messages.extend(msgs);
} }
Input::Fetched(rid, nid, result) => { Input::Fetched(rid, nid, result) => {
let result = Rc::try_unwrap(result).unwrap(); let result = Rc::try_unwrap(result).unwrap();