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