feat(CON-1242): QueryStats aggregation that is robust against malicious nodes

Author: Sawchord

rs/ingress_manager/src/ingress_selector.rs

@@ -1517,6 +1517,11 @@ mod tests {
             None,
             Some(
                 ReplicatedStateBuilder::default()
+                    .with_node_ids(
+                        (1..=SMALL_APP_SUBNET_MAX_SIZE as u64)
+                            .map(node_test_id)
+                            .collect(),
+                    )
                     .with_canister(
                         CanisterStateBuilder::default()
                             .with_canister_id(canister_test_id(0))

rs/interfaces/src/query_stats.rs

@@ -10,13 +10,13 @@ pub enum QueryStatsPermanentValidationError {
     InvalidNodeId { expected: NodeId, reported: NodeId },
     /// The epoch is lower than the aggregated height in the state manager
     EpochAlreadyAggregated {
-        expected: QueryStatsEpoch,
-        reported: QueryStatsEpoch,
+        highest_aggregated_epoch: QueryStatsEpoch,
+        payload_epoch: QueryStatsEpoch,
     },
     /// The epoch is higher than the certified height would allow for
     EpochTooHigh {
-        expected: QueryStatsEpoch,
-        reported: QueryStatsEpoch,
+        max_valid_epoch: QueryStatsEpoch,
+        payload_epoch: QueryStatsEpoch,
     },
     /// Stats for a [`CanisterId`] have been send twice
     DuplicateCanisterId(CanisterId),

rs/messaging/src/message_routing.rs

@@ -564,7 +564,6 @@ impl BatchProcessorImpl {
             stream_builder,
             log.clone(),
             metrics.clone(),
-            hypervisor_config.query_stats_epoch_length,
         ));
 
         Self {

rs/messaging/src/scheduling/valid_set_rule/test.rs

@@ -19,7 +19,7 @@ use ic_test_utilities_state::{
     MockIngressHistory, ReplicatedStateBuilder,
 };
 use ic_test_utilities_types::{
-    ids::{canister_test_id, message_test_id, subnet_test_id, user_test_id},
+    ids::{canister_test_id, message_test_id, node_test_id, subnet_test_id, user_test_id},
     messages::SignedIngressBuilder,
 };
 use ic_types::{
@@ -464,6 +464,11 @@ fn canister_on_application_subnet_charges_for_ingress() {
     let own_subnet_type = SubnetType::Application;
     let own_subnet_id = subnet_test_id(0);
     let mut state = ReplicatedStateBuilder::new()
+        .with_node_ids(
+            (1..=SMALL_APP_SUBNET_MAX_SIZE as u64)
+                .map(node_test_id)
+                .collect(),
+        )
         .with_subnet_type(own_subnet_type)
         .with_canister(
             CanisterStateBuilder::new()

rs/messaging/src/state_machine.rs

@@ -36,7 +36,6 @@ pub(crate) struct StateMachineImpl {
     stream_builder: Box<dyn StreamBuilder>,
     log: ReplicaLogger,
     metrics: MessageRoutingMetrics,
-    query_stats_epoch_length: u64,
 }
 
 impl StateMachineImpl {
@@ -46,15 +45,13 @@ impl StateMachineImpl {
         stream_builder: Box<dyn StreamBuilder>,
         log: ReplicaLogger,
         metrics: MessageRoutingMetrics,
-        query_stats_epoch_length: u64,
     ) -> Self {
         Self {
             scheduler,
             demux,
             stream_builder,
             log,
             metrics,
-            query_stats_epoch_length,
         }
     }
 
@@ -86,9 +83,7 @@ impl StateMachine for StateMachineImpl {
             deliver_query_stats(
                 query_stats,
                 &mut state,
-                batch.batch_number,
                 &self.log,
-                self.query_stats_epoch_length,
                 &self.metrics.query_stats_metrics,
             );
         }

rs/messaging/src/state_machine/tests.rs

@@ -156,7 +156,6 @@ fn state_machine_populates_network_topology() {
             fixture.stream_builder,
             log,
             fixture.metrics,
-            ic_config::execution_environment::QUERY_STATS_EPOCH_LENGTH,
         ));
 
         assert_ne!(
@@ -191,7 +190,6 @@ fn test_delivered_batch(provided_batch: Batch) -> ReplicatedState {
             fixture.stream_builder,
             log,
             fixture.metrics,
-            ic_config::execution_environment::QUERY_STATS_EPOCH_LENGTH,
         ));
 
         state_machine.execute_round(
@@ -292,7 +290,6 @@ fn test_batch_time_impl(
             fixture.stream_builder,
             log,
             fixture.metrics,
-            ic_config::execution_environment::QUERY_STATS_EPOCH_LENGTH,
         );
 
         assert_eq!(

rs/protobuf/def/state/stats/v1/stats.proto

@@ -8,12 +8,13 @@ message Stats {
 }
 
 message QueryStats {
-  uint64 epoch = 1;
+  uint64 highest_aggregated_epoch = 1;
   repeated QueryStatsInner query_stats = 2;
 }
 
 message QueryStatsInner {
   types.v1.NodeId proposer = 1;
+  uint64 epoch = 7;
   types.v1.CanisterId canister = 2;
   uint32 num_calls = 3;
   uint64 num_instructions = 4;

rs/protobuf/src/gen/state/state.stats.v1.rs

@@ -8,7 +8,7 @@ pub struct Stats {
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct QueryStats {
     #[prost(uint64, tag = "1")]
-    pub epoch: u64,
+    pub highest_aggregated_epoch: u64,
     #[prost(message, repeated, tag = "2")]
     pub query_stats: ::prost::alloc::vec::Vec<QueryStatsInner>,
 }
@@ -17,6 +17,8 @@ pub struct QueryStats {
 pub struct QueryStatsInner {
     #[prost(message, optional, tag = "1")]
     pub proposer: ::core::option::Option<super::super::super::types::v1::NodeId>,
+    #[prost(uint64, tag = "7")]
+    pub epoch: u64,
     #[prost(message, optional, tag = "2")]
     pub canister: ::core::option::Option<super::super::super::types::v1::CanisterId>,
     #[prost(uint32, tag = "3")]

rs/query_stats/src/metrics.rs

@@ -1,5 +1,7 @@
 use ic_metrics::{buckets::decimal_buckets, MetricsRegistry};
-use prometheus::{HistogramVec, IntGauge};
+use prometheus::{HistogramVec, IntCounter, IntGauge};
+
+pub(crate) const CRITICAL_ERROR_AGGREGATION_FAILURE: &str = "query_stats_aggregator_failure";
 
 /// Metrics for query stats collector
 ///
@@ -77,6 +79,10 @@ pub struct QueryStatsAggregatorMetrics {
     /// This is lower than the epoch for which we collect stats, as there is
     /// a delay for propagating local query stats via consensus blocks.
     pub query_stats_aggregator_current_epoch: IntGauge,
+    /// The number of records stored in the unaggregateed state
+    pub query_stats_aggregator_num_records: IntGauge,
+    /// Critical error occuring in aggregator
+    pub query_stats_critical_error_aggregator_failure: IntCounter,
 }
 
 impl QueryStatsAggregatorMetrics {
@@ -86,6 +92,12 @@ impl QueryStatsAggregatorMetrics {
                 "query_stats_aggregator_current_epoch",
                 "Current epoch of the query stats aggregator",
             ),
+            query_stats_aggregator_num_records: metrics_registry.int_gauge(
+                "query_stats_aggregator_num_records",
+                "The number of records stored in the unaggregateed state",
+            ),
+            query_stats_critical_error_aggregator_failure: metrics_registry
+                .error_counter(CRITICAL_ERROR_AGGREGATION_FAILURE),
         }
     }
 }

rs/query_stats/src/payload_builder.rs

@@ -1,4 +1,6 @@
-use crate::metrics::QueryStatsPayloadBuilderMetrics;
+use crate::{
+    metrics::QueryStatsPayloadBuilderMetrics, state_machine::get_stats_for_node_id_and_epoch,
+};
 use crossbeam_channel::{Receiver, TryRecvError};
 use ic_interfaces::{
     batch_payload::{BatchPayloadBuilder, PastPayload, ProposalContext},
@@ -237,8 +239,8 @@ impl QueryStatsPayloadBuilderImpl {
         if payload.epoch > max_valid_epoch {
             return Err(permanent_error(
                 QueryStatsPermanentValidationError::EpochTooHigh {
-                    expected: max_valid_epoch,
-                    reported: payload.epoch,
+                    max_valid_epoch,
+                    payload_epoch: payload.epoch,
                 },
             ));
         }
@@ -321,33 +323,36 @@ impl QueryStatsPayloadBuilderImpl {
         let mut previous_ids = BTreeSet::<CanisterId>::new();
 
         // Check that the epoch we are requesting has not been aggregated yet
-        if epoch < state_stats.epoch.unwrap_or(0.into()) {
+        // If there is no `highest_aggregated_epoch` in the state, we have not aggregated
+        // any epochs, therefore we unwrap to `false`
+        if state_stats
+            .highest_aggregated_epoch
+            .map(|highest_aggregated_epoch| epoch <= highest_aggregated_epoch)
+            .unwrap_or(false)
+        {
             warn!(
                 every_n_seconds => 5,
                 self.log,
                 "QueryStats: requesting previous_ids for epoch {:?} that is below aggregated epoch {:?}",
                 epoch,
-                state_stats.epoch
+                state_stats.highest_aggregated_epoch
             );
 
             return Err(permanent_error(
                 QueryStatsPermanentValidationError::EpochAlreadyAggregated {
-                    expected: state_stats.epoch.unwrap_or(0.into()),
-                    reported: epoch,
+                    highest_aggregated_epoch: state_stats
+                        .highest_aggregated_epoch
+                        .unwrap_or(0.into()),
+                    payload_epoch: epoch,
                 },
             ));
         }
 
-        // Check the certified state for stats already sent
-        // Skip if certified state is not the same epoch
-        if state_stats.epoch == Some(epoch) {
-            previous_ids.extend(
-                state_stats
-                    .stats
-                    .iter()
-                    .filter(|(_, stat_map)| stat_map.contains_key(&node_id))
-                    .map(|(canister_id, _)| canister_id),
-            );
+        // Check the certified state for stats that we have already sent
+        if let Some(state_stats) = get_stats_for_node_id_and_epoch(state_stats, &node_id, &epoch)
+            .map(|record| record.iter().map(|(canister_id, _)| canister_id))
+        {
+            previous_ids.extend(state_stats);
         }
 
         // Check past payloads for stats already sent
@@ -475,7 +480,12 @@ mod tests {
     #[test]
     fn past_payload_test() {
         let test_stats = test_epoch_stats(0, 500);
-        let state = test_state(epoch_stats_for_state(&test_stats, 0..200, node_test_id(1)));
+        let state = test_state(epoch_stats_for_state(
+            &test_stats,
+            0..200,
+            node_test_id(1),
+            None,
+        ));
         let payload_builder = setup_payload_builder_impl(state, test_stats.clone());
         let validation_context = test_validation_context();
         let proposal_context = test_proposal_context(&validation_context);
@@ -517,8 +527,8 @@ mod tests {
     fn node_id_check_test() {
         let test_stats = test_epoch_stats(0, 500);
 
-        let stats1 = epoch_stats_for_state(&test_stats, 0..100, node_test_id(1));
-        let stats2 = epoch_stats_for_state(&test_stats, 100..200, node_test_id(2));
+        let stats1 = epoch_stats_for_state(&test_stats, 0..100, node_test_id(1), None);
+        let stats2 = epoch_stats_for_state(&test_stats, 100..200, node_test_id(2), None);
         let stats = merge_raw_query_stats(stats1, stats2);
         let state = test_state(stats);
 
@@ -594,7 +604,12 @@ mod tests {
     #[test]
     fn epoch_too_low_test() {
         let test_stats = test_epoch_stats(1234, 100);
-        let state = test_state(epoch_stats_for_state(&test_stats, 0..100, node_test_id(1)));
+        let state = test_state(epoch_stats_for_state(
+            &test_stats,
+            0..100,
+            node_test_id(1),
+            Some(1234),
+        ));
         let payload_builder = setup_payload_builder_impl(state, test_stats);
         let validation_context = test_validation_context();
         let proposal_context = test_proposal_context(&validation_context);
@@ -616,12 +631,12 @@ mod tests {
             Err(ValidationError::Permanent(
                 PayloadPermanentError::QueryStatsPayloadValidationError(
                     QueryStatsPermanentValidationError::EpochAlreadyAggregated {
-                        expected,
-                        reported,
+                        highest_aggregated_epoch,
+                        payload_epoch,
                     },
                 ),
-            )) if expected == QueryStatsEpoch::new(1234) && reported == QueryStatsEpoch::new(0) => {
-            }
+            )) if highest_aggregated_epoch == QueryStatsEpoch::new(1234)
+                && payload_epoch == QueryStatsEpoch::new(0) => {}
             Err(err) => panic!(
                 "QueryStatsPayload had epoch too low, yet instead got error {:?}",
                 err
@@ -655,7 +670,10 @@ mod tests {
         match validation_result {
             Err(ValidationError::Permanent(
                 PayloadPermanentError::QueryStatsPayloadValidationError(
-                    QueryStatsPermanentValidationError::EpochTooHigh { expected, reported },
+                    QueryStatsPermanentValidationError::EpochTooHigh {
+                        max_valid_epoch: expected,
+                        payload_epoch: reported,
+                    },
                 ),
             )) if expected == QueryStatsEpoch::new(0) && reported == QueryStatsEpoch::new(1234) => {
             }
@@ -676,7 +694,12 @@ mod tests {
     #[test]
     fn duplicate_id_test() {
         let test_stats = test_epoch_stats(0, 4);
-        let state = test_state(epoch_stats_for_state(&test_stats, 0..1, node_test_id(1)));
+        let state = test_state(epoch_stats_for_state(
+            &test_stats,
+            0..1,
+            node_test_id(1),
+            None,
+        ));
         let payload_builder = setup_payload_builder_impl(state, test_stats.clone());
         let validation_context = test_validation_context();
         let proposal_context = test_proposal_context(&validation_context);
@@ -800,23 +823,34 @@ mod tests {
         query_stats: &LocalQueryStats,
         range: Range<usize>,
         node: NodeId,
+        highest_aggregated_epoch: Option<u64>,
     ) -> RawQueryStats {
+        let stats = vec![(
+            node,
+            vec![(
+                query_stats.epoch,
+                query_stats.stats[range]
+                    .iter()
+                    .map(|stat| (stat.canister_id, stat.stats.clone()))
+                    .collect(),
+            )]
+            .into_iter()
+            .collect(),
+        )]
+        .into_iter()
+        .collect();
+
         RawQueryStats {
-            epoch: Some(query_stats.epoch),
-            stats: query_stats.stats[range]
-                .iter()
-                .map(|stat| {
-                    (
-                        stat.canister_id,
-                        [(node, stat.stats.clone())].into_iter().collect(),
-                    )
-                })
-                .collect(),
+            highest_aggregated_epoch: highest_aggregated_epoch.map(QueryStatsEpoch::new),
+            stats,
         }
     }
 
     fn merge_raw_query_stats(mut stats1: RawQueryStats, stats2: RawQueryStats) -> RawQueryStats {
-        assert_eq!(stats1.epoch, stats2.epoch);
+        assert_eq!(
+            stats1.highest_aggregated_epoch,
+            stats2.highest_aggregated_epoch
+        );
 
         for (canister_id, stat2) in stats2.stats {
             stats1