feat: RUN-887: Query Cache: Initial composite query support

rs/execution_environment/src/metrics.rs

@@ -142,6 +142,9 @@ pub(crate) struct QueryHandlerMetrics {
     pub query_critical_error: IntCounter,
     /// The total number of tracked System API calls invoked during the query execution.
     pub query_system_api_calls: IntCounterVec,
+    /// The number of canisters evaluated and executed at least once
+    /// during the call graph evaluation.
+    pub query_evaluated_canisters: Histogram,
 }
 
 impl QueryHandlerMetrics {
@@ -252,6 +255,11 @@ impl QueryHandlerMetrics {
                         during the query execution",
                 &["system_api_call_counter"],
             ),
+            query_evaluated_canisters: metrics_registry.histogram(
+                "execution_query_evaluated_canisters",
+                "The number of canisters evaluated and executed at least once during the call graph evaluation",
+                vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 10.0],
+            ),
         }
     }
 }

rs/execution_environment/src/query_handler.rs

@@ -279,11 +279,13 @@ impl QueryHandler for InternalHttpQueryHandler {
             &measurement_scope,
         );
         context.observe_system_api_calls(&self.metrics.query_system_api_calls);
+        context.observe_evaluated_canisters(&self.metrics.query_evaluated_canisters);
 
         // Add the query execution result to the query cache (if the query caching is enabled).
         if self.config.query_caching == FlagStatus::Enabled {
             if let (Some(key), Some(env)) = (cache_entry_key, cache_entry_env) {
                 let call_counters = context.system_api_call_counters();
+                let _evaluated_ids = context.evaluated_canister_ids();
                 self.query_cache.push(key, env, &result, call_counters);
             }
         }

rs/execution_environment/src/query_handler/query_cache/tests.rs

@@ -820,6 +820,30 @@ fn query_cache_metrics_system_api_calls_work_on_composite_query() {
     );
 }
 
+#[test]
+fn query_cache_metrics_evaluated_canisters_work() {
+    let mut test = builder_with_query_caching().build();
+    let a_id = test.universal_canister().unwrap();
+    let b_id = test.universal_canister().unwrap();
+
+    let a = wasm()
+        .composite_query(
+            b_id,
+            call_args().on_reply(
+                wasm().composite_query(b_id, call_args().on_reply(wasm().reply_data(&[42]))),
+            ),
+        )
+        .build();
+    test.non_replicated_query(a_id, "composite_query", a)
+        .unwrap();
+
+    let m = &query_handler(&test).metrics;
+
+    // Two canisters reported once.
+    assert_eq!(1, m.query_evaluated_canisters.get_sample_count());
+    assert_eq!(2.0, m.query_evaluated_canisters.get_sample_sum());
+}
+
 #[test]
 fn query_cache_composite_queries_return_the_same_result() {
     let mut test = builder_with_query_caching().build();
@@ -889,8 +913,7 @@ fn query_cache_nested_queries_never_get_cached() {
     // The query has no time or balance dependencies...
     let q = wasm()
         // ...but there is a nested query.
-        .composite_query(b_id, call_args())
-        .reply_data(&[42])
+        .composite_query(b_id, call_args().on_reply(wasm().reply_data(&[42])))
         .build();
 
     // Run the query for the first time.

rs/execution_environment/src/query_handler/query_call_graph.rs

@@ -145,7 +145,7 @@ pub(super) fn evaluate_query_call_graph(
         }
     }
 
-    // Each iteraton of the loop above either pushes an entry onto the call
+    // Each iteration of the loop above either pushes an entry onto the call
     // stack or sets the callee result. At this point the call stack is empty,
     // so the callee result must have been set and `unwrap` is safe here.
     callee_result.unwrap()

rs/execution_environment/src/query_handler/query_context.rs

@@ -40,8 +40,13 @@ use ic_types::{
     methods::{FuncRef, WasmClosure},
     NumSlices,
 };
-use prometheus::{IntCounter, IntCounterVec};
-use std::{collections::VecDeque, sync::Arc, time::Duration, time::Instant};
+use prometheus::{Histogram, IntCounter, IntCounterVec};
+use std::{
+    collections::{BTreeSet, VecDeque},
+    sync::Arc,
+    time::Duration,
+    time::Instant,
+};
 
 /// The response of a query. If the query originated from a user, then it
 /// contains either `UserResponse` or `UserError`. If the query originated from
@@ -105,6 +110,9 @@ pub(super) struct QueryContext<'a> {
     local_query_execution_stats: Option<&'a QueryStatsCollector>,
     /// How many times each tracked System API call was invoked during the query execution.
     system_api_call_counters: SystemApiCallCounters,
+    /// A set of canister IDs evaluated and executed at least once in this query context.
+    /// The information is used by the query cache for composite queries.
+    evaluated_canister_ids: BTreeSet<CanisterId>,
 }
 
 impl<'a> QueryContext<'a> {
@@ -154,6 +162,7 @@ impl<'a> QueryContext<'a> {
             query_critical_error,
             local_query_execution_stats,
             system_api_call_counters: SystemApiCallCounters::default(),
+            evaluated_canister_ids: BTreeSet::default(),
         }
     }
 
@@ -407,8 +416,8 @@ impl<'a> QueryContext<'a> {
                 &mut self.round_limits,
                 self.query_critical_error,
             );
-        self.system_api_call_counters
-            .saturating_add(system_api_call_counters);
+        self.add_system_api_call_counters(system_api_call_counters);
+        self.insert_evaluated_canister_id(canister.canister_id());
         let instructions_executed = instruction_limit - instructions_left;
 
         let ingress_payload_size = method_payload.len();
@@ -451,6 +460,17 @@ impl<'a> QueryContext<'a> {
         (canister, result)
     }
 
+    /// Add up System API call counters.
+    fn add_system_api_call_counters(&mut self, system_api_call_counters: SystemApiCallCounters) {
+        self.system_api_call_counters
+            .saturating_add(system_api_call_counters);
+    }
+
+    /// Add a canister ID into a set of actually executed canisters.
+    fn insert_evaluated_canister_id(&mut self, canister_id: CanisterId) {
+        self.evaluated_canister_ids.insert(canister_id);
+    }
+
     fn finish(
         &self,
         canister: &mut CanisterState,
@@ -483,6 +503,11 @@ impl<'a> QueryContext<'a> {
             .inc_by(self.system_api_call_counters.time as u64);
     }
 
+    /// Observe the number evaluated canisters in the corresponding metrics.
+    pub(super) fn observe_evaluated_canisters(&mut self, query_evaluated_canisters: &Histogram) {
+        query_evaluated_canisters.observe(self.evaluated_canister_ids.len() as f64)
+    }
+
     fn execute_callback(
         &mut self,
         mut canister: CanisterState,
@@ -589,8 +614,8 @@ impl<'a> QueryContext<'a> {
             call_context.time(),
         );
 
-        self.system_api_call_counters
-            .saturating_add(output.system_api_call_counters);
+        self.add_system_api_call_counters(output.system_api_call_counters);
+        // There is no need to update evaluated canister ids, as it's a callback.
         let canister_current_memory_usage = canister.memory_usage();
         let canister_current_message_memory_usage = canister.message_memory_usage();
         canister.execution_state = Some(output_execution_state);
@@ -691,8 +716,8 @@ impl<'a> QueryContext<'a> {
                 time,
             );
 
-        self.system_api_call_counters
-            .saturating_add(cleanup_output.system_api_call_counters);
+        self.add_system_api_call_counters(cleanup_output.system_api_call_counters);
+        // There is no need to update evaluated canister ids, as it's a cleanup.
         canister.execution_state = Some(output_execution_state);
         match cleanup_output.wasm_result {
             Ok(_) => {
@@ -1019,4 +1044,9 @@ impl<'a> QueryContext<'a> {
     pub fn system_api_call_counters(&self) -> &SystemApiCallCounters {
         &self.system_api_call_counters
     }
+
+    /// Return a list of actually executed canister IDs.
+    pub fn evaluated_canister_ids(&self) -> &BTreeSet<CanisterId> {
+        &self.evaluated_canister_ids
+    }
 }