diff --git a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs index f62edd64b06..8b81ce8ef75 100644 --- a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs +++ b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/main.rs @@ -13,6 +13,7 @@ mod message_test; mod handlers; mod message; +mod message_index_detector; pub mod metrics; mod protocol; mod stress_test_node; diff --git a/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message_index_detector.rs b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message_index_detector.rs new file mode 100644 index 00000000000..7c8c0a5e593 --- /dev/null +++ b/crates/apollo_network_benchmark/src/bin/broadcast_network_stress_test_node/message_index_detector.rs @@ -0,0 +1,31 @@ +// TODO(AndrewL): remove this once the struct is used +#![allow(dead_code)] + +#[derive(Default, Clone, Copy)] +pub struct MessageIndexTracker { + seen_messages_count: u64, + max_message_index: Option, + min_message_index: Option, +} + +impl MessageIndexTracker { + pub fn seen_message(&mut self, message_index: u64) { + self.seen_messages_count += 1; + if self.max_message_index.is_none() || self.max_message_index.unwrap() < message_index { + self.max_message_index = Some(message_index); + } + if self.min_message_index.is_none() || self.min_message_index.unwrap() > message_index { + self.min_message_index = Some(message_index); + } + } + + pub fn pending_messages_count(&self) -> u64 { + if self.seen_messages_count == 0 { + return 0; + } + + let min_message_index = self.min_message_index.unwrap(); + let max_message_index = self.max_message_index.unwrap(); + max_message_index - min_message_index + 1 - self.seen_messages_count + } +}