feat(pocket-ic): support custom blockmaker metrics at each round (dfinity#3685)

Currently PocketIC at each tick() executes a
[round](https://github.com/dfinity/ic/blob/master/rs/state_machine_tests/src/lib.rs#L1338)
with state_machine_tests. However a dummy
[blockmaker](https://github.com/dfinity/ic/blob/0468e6f90e52fab4be7b0bf722dacb520a589c55/rs/types/types/src/batch.rs#L199-L200)
is used.

The metrics_collector_canister I am building store and serve node
metrics from all the subnet. To test this canister I need then to modify
PocketIC so that at each "round" it stores blockmaker metrics that I can
pass in from tests

---------

Co-authored-by: mraszyk <[email protected]>

Author: pietrodimarco-dfinity

Cargo.lock

@@ -356,6 +356,23 @@ impl From<Principal> for RawNodeId {
     }
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema, Default)]
+pub struct TickConfigs {
+    pub blockmakers: Option<BlockmakerConfigs>,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema)]
+pub struct BlockmakerConfigs {
+    pub blockmakers_per_subnet: Vec<RawSubnetBlockmaker>,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema)]
+pub struct RawSubnetBlockmaker {
+    pub subnet: RawSubnetId,
+    pub blockmaker: RawNodeId,
+    pub failed_blockmakers: Vec<RawNodeId>,
+}
+
 #[derive(Serialize, Deserialize, JsonSchema)]
 pub struct RawVerifyCanisterSigArg {
     #[serde(deserialize_with = "base64::deserialize")]

packages/pocket-ic/src/common/rest.rs

@@ -522,6 +522,14 @@ impl PocketIc {
         runtime.block_on(async { self.pocket_ic.tick().await })
     }
 
+    /// Make the IC produce and progress by one block with custom
+    /// configs for the round.
+    #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id))]
+    pub fn tick_with_configs(&self, configs: crate::common::rest::TickConfigs) {
+        let runtime = self.runtime.clone();
+        runtime.block_on(async { self.pocket_ic.tick_with_configs(configs).await })
+    }
+
     /// Configures the IC to make progress automatically,
     /// i.e., periodically update the time of the IC
     /// to the real time and execute rounds on the subnets.

packages/pocket-ic/src/lib.rs

@@ -5,7 +5,7 @@ use crate::common::rest::{
     MockCanisterHttpResponse, RawAddCycles, RawCanisterCall, RawCanisterHttpRequest, RawCanisterId,
     RawCanisterResult, RawCycles, RawEffectivePrincipal, RawIngressStatusArgs, RawMessageId,
     RawMockCanisterHttpResponse, RawPrincipalId, RawSetStableMemory, RawStableMemory, RawSubnetId,
-    RawTime, RawVerifyCanisterSigArg, SubnetId, Topology,
+    RawTime, RawVerifyCanisterSigArg, SubnetId, TickConfigs, Topology,
 };
 use crate::management_canister::{
     CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterInstallModeUpgradeInner,
@@ -290,8 +290,15 @@ impl PocketIc {
     /// inter-canister calls or heartbeats.
     #[instrument(skip(self), fields(instance_id=self.instance_id))]
     pub async fn tick(&self) {
+        self.tick_with_configs(TickConfigs::default()).await;
+    }
+
+    /// Make the IC produce and progress by one block with custom
+    /// configs for the round.
+    #[instrument(skip(self), fields(instance_id=self.instance_id))]
+    pub async fn tick_with_configs(&self, configs: TickConfigs) {
         let endpoint = "update/tick";
-        self.post::<(), _>(endpoint, "").await;
+        self.post::<(), _>(endpoint, configs).await;
     }
 
     /// Configures the IC to make progress automatically,

packages/pocket-ic/src/nonblocking.rs

@@ -98,12 +98,29 @@ type TransformArgs = record {
   context : blob;
 };
 
+type NodeMetricsHistoryArgs = record {
+    subnet_id: principal;
+    start_at_timestamp_nanos: nat64;
+};
+
+type NodeMetrics = record {
+    node_id: principal;
+    num_blocks_proposed_total: nat64;
+    num_block_failures_total: nat64;
+};
+
+type NodeMetricsHistoryResponse = record {
+  timestamp_nanos: nat64;
+  node_metrics: vec NodeMetrics;
+};
+
 service : {
   http_request: (request: HttpGatewayRequest) -> (HttpGatewayResponse) query;
   schnorr_public_key : (opt principal, vec blob, SchnorrKeyId) -> (SchnorrPublicKeyResult);
   sign_with_schnorr : (blob, vec blob, SchnorrKeyId) -> (SignWithSchnorrResult);
   ecdsa_public_key : (opt principal, vec blob, text) -> (EcdsaPublicKeyResult);
   sign_with_ecdsa : (blob, vec blob, text) -> (SignWithEcdsaResult);
+  node_metrics_history_proxy: (NodeMetricsHistoryArgs) -> (vec NodeMetricsHistoryResponse);
   canister_http : () -> (HttpResponseResult);
   canister_http_with_transform : () -> (HttpResponse);
   transform : (TransformArgs) -> (HttpResponse) query;

packages/pocket-ic/test_canister/canister.did

@@ -293,6 +293,40 @@ async fn call_with_large_blob(canister: Principal, blob_len: usize) -> usize {
         .0
 }
 
+#[derive(CandidType, Deserialize)]
+pub struct NodeMetrics {
+    pub node_id: Principal,
+    pub num_blocks_proposed_total: u64,
+    pub num_block_failures_total: u64,
+}
+
+#[derive(CandidType, Deserialize)]
+pub struct NodeMetricsHistoryResponse {
+    pub timestamp_nanos: u64,
+    pub node_metrics: Vec<NodeMetrics>,
+}
+
+#[derive(CandidType, Deserialize)]
+pub struct NodeMetricsHistoryArgs {
+    pub start_at_timestamp_nanos: u64,
+    pub subnet_id: Principal,
+}
+
+#[update]
+async fn node_metrics_history_proxy(
+    args: NodeMetricsHistoryArgs,
+) -> Vec<NodeMetricsHistoryResponse> {
+    ic_cdk::api::call::call_with_payment128::<_, (Vec<NodeMetricsHistoryResponse>,)>(
+        candid::Principal::management_canister(),
+        "node_metrics_history",
+        (args,),
+        0_u128,
+    )
+    .await
+    .unwrap()
+    .0
+}
+
 // executing many instructions
 
 #[update]

packages/pocket-ic/test_canister/src/canister.rs

@@ -2,10 +2,12 @@ use candid::{decode_one, encode_one, CandidType, Decode, Deserialize, Encode, Pr
 use ic_certification::Label;
 use ic_transport_types::Envelope;
 use ic_transport_types::EnvelopeContent::ReadState;
+use pocket_ic::common::rest::{BlockmakerConfigs, RawSubnetBlockmaker, TickConfigs};
 use pocket_ic::management_canister::{
     CanisterIdRecord, CanisterInstallMode, CanisterSettings, EcdsaPublicKeyResult,
-    HttpRequestResult, ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm,
-    SchnorrPublicKeyArgsKeyId, SchnorrPublicKeyResult, SignWithBip341Aux, SignWithSchnorrAux,
+    HttpRequestResult, NodeMetricsHistoryArgs, NodeMetricsHistoryResultItem,
+    ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm, SchnorrPublicKeyArgsKeyId,
+    SchnorrPublicKeyResult, SignWithBip341Aux, SignWithSchnorrAux,
 };
 use pocket_ic::{
     common::rest::{
@@ -2214,3 +2216,82 @@ fn test_reject_response_type() {
         assert!(!err.certified);
     }
 }
+
+#[test]
+fn test_custom_blockmaker_metrics() {
+    const HOURS_IN_SECONDS: u64 = 60 * 60;
+
+    let pocket_ic = PocketIcBuilder::new().with_application_subnet().build();
+    let topology = pocket_ic.topology();
+    let application_subnet = topology.get_app_subnets()[0];
+
+    // Create and install test canister.
+    let canister = pocket_ic.create_canister_on_subnet(None, None, application_subnet);
+    pocket_ic.add_cycles(canister, INIT_CYCLES);
+    pocket_ic.install_canister(canister, test_canister_wasm(), vec![], None);
+
+    let nodes = topology
+        .subnet_configs
+        .get(&application_subnet)
+        .unwrap()
+        .clone();
+
+    let blockmaker_1 = nodes.node_ids[0].clone();
+    let blockmaker_2 = nodes.node_ids[1].clone();
+
+    let subnets_blockmakers = vec![RawSubnetBlockmaker {
+        subnet: application_subnet.into(),
+        blockmaker: blockmaker_1.clone(),
+        failed_blockmakers: vec![blockmaker_2.clone()],
+    }];
+
+    let tick_configs = TickConfigs {
+        blockmakers: Some(BlockmakerConfigs {
+            blockmakers_per_subnet: subnets_blockmakers,
+        }),
+    };
+    let daily_blocks = 5;
+
+    // Blockmaker metrics are recorded in the management canister
+    for _ in 0..daily_blocks {
+        pocket_ic.tick_with_configs(tick_configs.clone());
+    }
+    // Advance time until next day so that management canister can record blockmaker metrics
+    pocket_ic.advance_time(std::time::Duration::from_secs(HOURS_IN_SECONDS * 24));
+    pocket_ic.tick();
+
+    let response = pocket_ic
+        .update_call(
+            canister,
+            Principal::anonymous(),
+            // Calls the node_metrics_history method on the management canister
+            "node_metrics_history_proxy",
+            Encode!(&NodeMetricsHistoryArgs {
+                subnet_id: application_subnet,
+                start_at_timestamp_nanos: 0,
+            })
+            .unwrap(),
+        )
+        .unwrap();
+
+    let first_node_metrics = Decode!(&response, Vec<NodeMetricsHistoryResultItem>)
+        .unwrap()
+        .remove(0)
+        .node_metrics;
+
+    let blockmaker_1_metrics = first_node_metrics
+        .iter()
+        .find(|x| x.node_id == Principal::from(blockmaker_1.clone()))
+        .unwrap()
+        .clone();
+    let blockmaker_2_metrics = first_node_metrics
+        .into_iter()
+        .find(|x| x.node_id == Principal::from(blockmaker_2.clone()))
+        .unwrap();
+
+    assert_eq!(blockmaker_1_metrics.num_blocks_proposed_total, daily_blocks);
+    assert_eq!(blockmaker_1_metrics.num_block_failures_total, 0);
+
+    assert_eq!(blockmaker_2_metrics.num_blocks_proposed_total, 0);
+    assert_eq!(blockmaker_2_metrics.num_block_failures_total, daily_blocks);
+}

packages/pocket-ic/tests/tests.rs

@@ -38,8 +38,10 @@ pub mod pocket_ic;
 pub mod state_api;
 
 use crate::state_api::state::OpOut;
-use ::pocket_ic::common::rest::{BinaryBlob, BlobId};
+use ::pocket_ic::common::rest::{BinaryBlob, BlobId, RawSubnetBlockmaker};
 use axum::async_trait;
+use candid::Principal;
+use ic_types::{NodeId, PrincipalId, SubnetId};
 use pocket_ic::PocketIc;
 use serde::Deserialize;
 
@@ -91,3 +93,28 @@ pub fn copy_dir(
     }
     Ok(())
 }
+
+#[derive(Clone, Debug)]
+pub struct SubnetBlockmaker {
+    pub subnet: SubnetId,
+    pub blockmaker: NodeId,
+    pub failed_blockmakers: Vec<NodeId>,
+}
+
+impl From<RawSubnetBlockmaker> for SubnetBlockmaker {
+    fn from(raw: RawSubnetBlockmaker) -> Self {
+        let subnet = SubnetId::from(PrincipalId::from(Principal::from(raw.subnet)));
+        let blockmaker = NodeId::from(PrincipalId::from(Principal::from(raw.blockmaker)));
+        let failed_blockmakers: Vec<NodeId> = raw
+            .failed_blockmakers
+            .into_iter()
+            .map(|node_id| NodeId::from(PrincipalId::from(Principal::from(node_id))))
+            .collect();
+
+        SubnetBlockmaker {
+            subnet,
+            blockmaker,
+            failed_blockmakers,
+        }
+    }
+}

rs/pocket_ic_server/src/lib.rs

@@ -1,5 +1,5 @@
 use crate::state_api::state::{HasStateLabel, OpOut, PocketIcError, StateLabel};
-use crate::{async_trait, copy_dir, BlobStore, OpId, Operation};
+use crate::{async_trait, copy_dir, BlobStore, OpId, Operation, SubnetBlockmaker};
 use askama::Template;
 use axum::{
     extract::State,
@@ -54,6 +54,7 @@ use ic_state_machine_tests::{
     SubmitIngressError, Subnets,
 };
 use ic_test_utilities_registry::add_subnet_list_record;
+use ic_types::batch::BlockmakerMetrics;
 use ic_types::ingress::{IngressState, IngressStatus};
 use ic_types::{
     artifact::UnvalidatedArtifactMutation,
@@ -77,7 +78,7 @@ use pocket_ic::common::rest::{
     self, BinaryBlob, BlobCompression, CanisterHttpHeader, CanisterHttpMethod, CanisterHttpRequest,
     CanisterHttpResponse, ExtendedSubnetConfigSet, MockCanisterHttpResponse, RawAddCycles,
     RawCanisterCall, RawCanisterId, RawEffectivePrincipal, RawMessageId, RawSetStableMemory,
-    SubnetInstructionConfig, SubnetKind, SubnetSpec, Topology,
+    SubnetInstructionConfig, SubnetKind, SubnetSpec, TickConfigs, Topology,
 };
 use pocket_ic::{ErrorCode, RejectCode, RejectResponse};
 use serde::{Deserialize, Serialize};
@@ -1579,14 +1580,82 @@ impl Operation for PubKey {
     }
 }
 
-#[derive(Copy, Clone, Debug)]
-pub struct Tick;
+#[derive(Clone, Debug)]
+pub struct Tick {
+    pub configs: TickConfigs,
+}
+
+impl Tick {
+    fn validate_blockmakers_per_subnet(
+        &self,
+        pic: &mut PocketIc,
+        subnets_blockmaker: &[SubnetBlockmaker],
+    ) -> Result<(), OpOut> {
+        for subnet_blockmaker in subnets_blockmaker {
+            if subnet_blockmaker
+                .failed_blockmakers
+                .contains(&subnet_blockmaker.blockmaker)
+            {
+                return Err(OpOut::Error(PocketIcError::BlockmakerContainedInFailed(
+                    subnet_blockmaker.blockmaker,
+                )));
+            }
+
+            let Some(state_machine) = pic.get_subnet_with_id(subnet_blockmaker.subnet) else {
+                return Err(OpOut::Error(PocketIcError::SubnetNotFound(
+                    subnet_blockmaker.subnet.get().0,
+                )));
+            };
+
+            let mut request_blockmakers = subnet_blockmaker.failed_blockmakers.clone();
+            request_blockmakers.push(subnet_blockmaker.blockmaker);
+            let subnet_nodes: Vec<_> = state_machine.nodes.iter().map(|n| n.node_id).collect();
+            for blockmaker in request_blockmakers {
+                if !subnet_nodes.contains(&blockmaker) {
+                    return Err(OpOut::Error(PocketIcError::BlockmakerNotFound(blockmaker)));
+                }
+            }
+        }
+        Ok(())
+    }
+}
 
 impl Operation for Tick {
     fn compute(&self, pic: &mut PocketIc) -> OpOut {
-        for subnet in pic.subnets.get_all() {
-            subnet.state_machine.execute_round();
+        let blockmakers_per_subnet = self.configs.blockmakers.as_ref().map(|cfg| {
+            cfg.blockmakers_per_subnet
+                .iter()
+                .cloned()
+                .map(SubnetBlockmaker::from)
+                .collect_vec()
+        });
+
+        if let Some(ref bm_per_subnet) = blockmakers_per_subnet {
+            if let Err(error) = self.validate_blockmakers_per_subnet(pic, bm_per_subnet) {
+                return error;
+            }
         }
+
+        let subnets = pic.subnets.subnets.read().unwrap();
+        for (subnet_id, subnet) in subnets.iter() {
+            let blockmaker_metrics = blockmakers_per_subnet.as_ref().and_then(|bm_per_subnet| {
+                bm_per_subnet
+                    .iter()
+                    .find(|bm| bm.subnet == *subnet_id)
+                    .map(|bm| BlockmakerMetrics {
+                        blockmaker: bm.blockmaker,
+                        failed_blockmakers: bm.failed_blockmakers.clone(),
+                    })
+            });
+
+            match blockmaker_metrics {
+                Some(metrics) => subnet
+                    .state_machine
+                    .execute_round_with_blockmaker_metrics(metrics),
+                None => subnet.state_machine.execute_round(),
+            }
+        }
+
         OpOut::NoOutput
     }