feat: add criterion bench for RouteProvider implementations

ic-agent/Cargo.toml

@@ -97,6 +97,7 @@ serde_json.workspace = true
 [target.'cfg(not(target_family = "wasm"))'.dev-dependencies]
 tokio = { workspace = true, features = ["full"] }
 mockito = "1.0.2"
+criterion = "0.5"
 
 [target.'cfg(target_family = "wasm")'.dev-dependencies]
 wasm-bindgen-test = "0.3.34"
@@ -133,8 +134,15 @@ wasm-bindgen = [
     "backoff/wasm-bindgen",
     "cached/wasm",
 ]
+bench = []
 
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu", "wasm32-unknown-unknown"]
 rustdoc-args = ["--cfg=docsrs"]
 features = ["hyper"]
+
+[[bench]]
+name = "perf_route_provider"
+path = "benches/perf_route_provider.rs"
+required-features = ["bench"]
+harness = false

ic-agent/benches/perf_route_provider.rs

@@ -0,0 +1,139 @@
+use std::{sync::Arc, time::Duration};
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use ic_agent::agent::http_transport::{
+    dynamic_routing::{
+        dynamic_route_provider::DynamicRouteProviderBuilder,
+        node::Node,
+        snapshot::{
+            latency_based_routing::LatencyRoutingSnapshot,
+            round_robin_routing::RoundRobinRoutingSnapshot, routing_snapshot::RoutingSnapshot,
+        },
+        test_utils::{NodeHealthCheckerMock, NodesFetcherMock},
+    },
+    route_provider::{RoundRobinRouteProvider, RouteProvider},
+};
+use reqwest::Client;
+use tokio::{runtime::Handle, sync::oneshot};
+use tokio_util::sync::CancellationToken;
+
+// To run the benchmark use the command:
+// $ cargo bench --bench perf_route_provider --features bench
+
+// Benchmarking function
+fn benchmark_route_providers(c: &mut Criterion) {
+    // For displaying trace messages of the inner running tasks in dynamic route providers, enable the subscriber below
+
+    // use tracing::Level;
+    // use tracing_subscriber::FmtSubscriber;
+    // FmtSubscriber::builder().with_max_level(Level::TRACE).init();
+
+    // Number of different domains for each route provider
+    let nodes_count = 100;
+
+    let mut group = c.benchmark_group("RouteProviders");
+    group.sample_size(10000);
+
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("failed to create runtime");
+
+    // Setup all route providers
+    let token = CancellationToken::new();
+    let route_providers =
+        setup_route_providers(nodes_count, runtime.handle().clone(), token.clone());
+
+    for (name, instance) in route_providers {
+        group.bench_function(name, |b| {
+            b.iter(|| {
+                let _url = instance.route().unwrap();
+            })
+        });
+    }
+    token.cancel();
+    group.finish();
+}
+
+criterion_group!(benches, benchmark_route_providers);
+criterion_main!(benches);
+
+fn setup_static_route_provider(nodes_count: usize) -> Arc<dyn RouteProvider> {
+    let urls: Vec<_> = (0..nodes_count)
+        .map(|idx| format!("https://domain_{idx}.app"))
+        .collect();
+    Arc::new(RoundRobinRouteProvider::new(urls).unwrap())
+}
+
+async fn setup_dynamic_route_provider<S: RoutingSnapshot + 'static>(
+    nodes_count: usize,
+    snapshot: S,
+) -> Arc<dyn RouteProvider> {
+    let client = Client::builder().build().expect("failed to build a client");
+
+    let nodes: Vec<_> = (0..nodes_count)
+        .map(|idx| Node::new(&format!("https://domain_{idx}.app")).unwrap())
+        .collect();
+
+    let fetcher = Arc::new(NodesFetcherMock::new());
+    let checker = Arc::new(NodeHealthCheckerMock::new());
+    let fetch_interval = Duration::from_secs(1);
+    let check_interval = Duration::from_millis(50);
+
+    fetcher.overwrite_nodes(nodes.clone());
+    checker.overwrite_healthy_nodes(nodes.clone());
+
+    // Use e.g. a couple of nodes as seeds.
+    let seeds = nodes[..2].to_vec();
+
+    let route_provider = DynamicRouteProviderBuilder::new(snapshot, seeds, client.clone())
+        .with_fetch_period(fetch_interval)
+        .with_fetcher(fetcher)
+        .with_check_period(check_interval)
+        .with_checker(checker)
+        .build()
+        .await;
+
+    Arc::new(route_provider)
+}
+
+fn setup_route_providers(
+    nodes_count: usize,
+    runtime: Handle,
+    cancellation_token: CancellationToken,
+) -> Vec<(String, Arc<dyn RouteProvider>)> {
+    // Assemble all instances for benching.
+    let mut route_providers = vec![];
+    // Setup static round-robin route provider
+    route_providers.push((
+        "Static round-robin RouteProvider".to_string(),
+        setup_static_route_provider(nodes_count),
+    ));
+    // Setup dynamic round-robin route provider
+    let (tx, rx) = oneshot::channel();
+    let token_cloned = cancellation_token.clone();
+    runtime.spawn(async move {
+        let rp = setup_dynamic_route_provider(nodes_count, RoundRobinRoutingSnapshot::new()).await;
+        tx.send(rp).unwrap();
+        token_cloned.cancelled().await;
+    });
+    let route_provider = runtime.block_on(async { rx.await.unwrap() });
+    route_providers.push((
+        "Dynamic round-robin RouteProvider".to_string(),
+        route_provider,
+    ));
+    // Setup dynamic latency-based route provider
+    let (tx, rx) = oneshot::channel();
+    let token_cloned = cancellation_token.clone();
+    runtime.spawn(async move {
+        let rp = setup_dynamic_route_provider(nodes_count, LatencyRoutingSnapshot::new()).await;
+        tx.send(rp).unwrap();
+        token_cloned.cancelled().await;
+    });
+    let route_provider = runtime.block_on(async { rx.await.unwrap() });
+    route_providers.push((
+        "Dynamic latency-based RouteProvider".to_string(),
+        route_provider,
+    ));
+    route_providers
+}

ic-agent/src/agent/http_transport/dynamic_routing/mod.rs

@@ -10,7 +10,8 @@ pub mod node;
 pub mod nodes_fetch;
 /// Routing snapshot implementation.
 pub mod snapshot;
-#[cfg(test)]
-pub(super) mod test_utils;
+/// Testing and benchmarking helpers.
+#[cfg(any(test, feature = "bench"))]
+pub mod test_utils;
 /// Type aliases used in dynamic routing.
 pub(super) mod type_aliases;

ic-agent/src/agent/http_transport/dynamic_routing/test_utils.rs

@@ -17,17 +17,16 @@ use crate::agent::http_transport::{
     route_provider::RouteProvider,
 };
 
-pub(super) fn route_n_times(n: usize, f: Arc<impl RouteProvider + ?Sized>) -> Vec<String> {
+///
+pub fn route_n_times(n: usize, f: Arc<impl RouteProvider + ?Sized>) -> Vec<String> {
     (0..n)
         .map(|_| f.route().unwrap().domain().unwrap().to_string())
         .collect()
 }
 
-pub(super) fn assert_routed_domains<T>(
-    actual: Vec<T>,
-    expected: Vec<T>,
-    expected_repetitions: usize,
-) where
+///
+pub fn assert_routed_domains<T>(actual: Vec<T>, expected: Vec<T>, expected_repetitions: usize)
+where
     T: AsRef<str> + Eq + Hash + Debug + Ord,
 {
     fn build_count_map<T>(items: &[T]) -> HashMap<&T, usize>
@@ -56,9 +55,10 @@ pub(super) fn assert_routed_domains<T>(
         .all(|&x| x == &expected_repetitions));
 }
 
+/// A mock implementation of the nodes Fetch trait, without http calls.
 #[derive(Debug)]
-pub(super) struct NodesFetcherMock {
-    // A set of nodes, existing in the topology.
+pub struct NodesFetcherMock {
+    /// A set of nodes, existing in the topology.
     pub nodes: AtomicSwap<Vec<Node>>,
 }
 
@@ -77,19 +77,22 @@ impl Default for NodesFetcherMock {
 }
 
 impl NodesFetcherMock {
+    /// Create a new instance.
     pub fn new() -> Self {
         Self {
             nodes: Arc::new(ArcSwap::from_pointee(vec![])),
         }
     }
 
+    /// Sets the existing nodes in the topology.
     pub fn overwrite_nodes(&self, nodes: Vec<Node>) {
         self.nodes.store(Arc::new(nodes));
     }
 }
 
+/// A mock implementation of the node's HealthCheck trait, without http calls.
 #[derive(Debug)]
-pub(super) struct NodeHealthCheckerMock {
+pub struct NodeHealthCheckerMock {
     healthy_nodes: Arc<ArcSwap<HashSet<Node>>>,
 }
 
@@ -112,12 +115,14 @@ impl HealthCheck for NodeHealthCheckerMock {
 }
 
 impl NodeHealthCheckerMock {
+    /// Creates a new instance
     pub fn new() -> Self {
         Self {
             healthy_nodes: Arc::new(ArcSwap::from_pointee(HashSet::new())),
         }
     }
 
+    /// Sets healthy nodes in the topology.
     pub fn overwrite_healthy_nodes(&self, healthy_nodes: Vec<Node>) {
         self.healthy_nodes
             .store(Arc::new(HashSet::from_iter(healthy_nodes)));