Fix QueryStatsFeatureGate

dfinity · Jan 22, 2024 · da9fc05 · da9fc05
1 parent 826f954
commit da9fc05
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/rs/query_stats/src/payload_builder.rs b/rs/query_stats/src/payload_builder.rs
@@ -11,7 +11,7 @@ use ic_logger::{error, warn, ReplicaLogger};
 use ic_metrics::MetricsRegistry;
 use ic_replicated_state::ReplicatedState;
 use ic_types::{
-    batch::{LocalQueryStats, QueryStatsPayload, ValidationContext, ENABLE_QUERY_STATS},
+    batch::{LocalQueryStats, QueryStatsPayload, ValidationContext},
     epoch_from_height, CanisterId, Height, NodeId, NumBytes, QueryStatsEpoch,
 };
 use std::{
@@ -94,10 +94,10 @@ impl BatchPayloadBuilder for QueryStatsPayloadBuilderImpl {
             }
         }
 
-        match ENABLE_QUERY_STATS {
-            true => self.build_payload_impl(height, max_size, past_payloads, context),
-            false => vec![],
+        if !self.enabled {
+            return vec![];
         }
+        self.build_payload_impl(height, max_size, past_payloads, context)
     }
 
     fn validate_payload(
@@ -120,7 +120,7 @@ impl BatchPayloadBuilder for QueryStatsPayloadBuilderImpl {
 
         // Check whether feature is enabled and reject if it isn't.
         // NOTE: All payloads that are processed at this point are non-empty
-        if !ENABLE_QUERY_STATS {
+        if !self.enabled {
             return Err(transient_error(
                 QueryStatsTransientValidationError::Disabled,
             ));
@@ -138,10 +138,6 @@ impl QueryStatsPayloadBuilderImpl {
         past_payloads: &[PastPayload],
         context: &ValidationContext,
     ) -> Vec<u8> {
-        if !self.enabled {
-            return vec![];
-        }
-
         let Ok(current_stats) = self.current_stats.read() else {
             return vec![];
         };

diff --git a/rs/types/types/src/batch.rs b/rs/types/types/src/batch.rs
@@ -29,14 +29,13 @@ use ic_btc_types_internal::BitcoinAdapterResponse;
 #[cfg(test)]
 use ic_exhaustive_derive::ExhaustiveSet;
 use ic_ic00_types::EcdsaKeyId;
+use ic_protobuf::proxy::ProxyDecodeError;
 use serde::{Deserialize, Serialize};
 use std::{
     collections::{BTreeMap, BTreeSet},
     convert::TryInto,
 };
 
-pub const ENABLE_QUERY_STATS: bool = false;
-
 /// The `Batch` provided to Message Routing for deterministic processing.
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Batch {
@@ -118,22 +117,28 @@ pub struct BatchMessages {
     pub query_stats: Option<QueryStatsPayload>,
 }
 
+/// Error type that can occur during an `BatchPayload::into_messages` call
+#[derive(Debug)]
+pub enum IntoMessagesError {
+    IngressPayloadError(IngressPayloadError),
+    QueryStatsPayloadError(ProxyDecodeError),
+}
+
 impl BatchPayload {
     /// Extract and return the set of ingress and xnet messages in a
     /// BatchPayload.
     /// Return error if deserialization of ingress payload fails.
     #[allow(clippy::result_large_err)]
-    pub fn into_messages(self) -> Result<BatchMessages, IngressPayloadError> {
+    pub fn into_messages(self) -> Result<BatchMessages, IntoMessagesError> {
         Ok(BatchMessages {
-            signed_ingress_msgs: self.ingress.try_into()?,
+            signed_ingress_msgs: self
+                .ingress
+                .try_into()
+                .map_err(IntoMessagesError::IngressPayloadError)?,
             certified_stream_slices: self.xnet.stream_slices,
             bitcoin_adapter_responses: self.self_validating.0,
-            query_stats: match ENABLE_QUERY_STATS {
-                true => QueryStatsPayload::deserialize(&self.query_stats)
-                    .ok()
-                    .flatten(),
-                false => None,
-            },
+            query_stats: QueryStatsPayload::deserialize(&self.query_stats)
+                .map_err(IntoMessagesError::QueryStatsPayloadError)?,
         })
     }