feat(governance): Log instructions of queries in Governance.

Cargo.toml

@@ -270,6 +270,7 @@ members = [
     "rs/nervous_system/proto",
     "rs/nervous_system/proto/protobuf_generator",
     "rs/nervous_system/proxied_canister_calls_tracker",
+    "rs/nervous_system/query_instruction_logger",
     "rs/nervous_system/rate_limits",
     "rs/nervous_system/root",
     "rs/nervous_system/runtime",

rs/nervous_system/query_instruction_logger/BUILD.bazel

@@ -0,0 +1,35 @@
+load("@rules_rust//rust:defs.bzl", "rust_proc_macro", "rust_test")
+
+package(default_visibility = ["//visibility:public"])
+
+DEPENDENCIES = [
+    "@crate_index//:quote",
+    "@crate_index//:syn",
+]
+
+DEV_DEPENDENCIES = [
+    "@crate_index//:candid",
+    "@crate_index//:ic-cdk",
+]
+
+LIB_SRCS = glob(
+    ["src/**/*.rs"],
+    exclude = ["**/*tests*/**"],
+)
+
+rust_proc_macro(
+    name = "query_instruction_logger",
+    srcs = LIB_SRCS,
+    crate_name = "ic_nervous_system_query_instruction_logger",
+    version = "0.0.1",
+    deps = DEPENDENCIES,
+)
+
+rust_test(
+    name = "query_instruction_logger_test",
+    srcs = glob(["tests/*.rs"]),
+    proc_macro_deps = [
+        ":query_instruction_logger",
+    ],
+    deps = DEPENDENCIES + DEV_DEPENDENCIES,
+)

rs/nervous_system/query_instruction_logger/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "ic-nervous-system-query-instruction-logger"
+version = "0.0.1"
+edition.workspace = true
+
+[lib]
+proc-macro = true
+
+[dependencies]
+quote = "1"
+syn = { workspace = true }
+
+[dev-dependencies]
+candid = { workspace = true }
+ic-cdk = { workspace = true }

rs/nervous_system/query_instruction_logger/src/lib.rs

@@ -0,0 +1,50 @@
+use proc_macro::TokenStream;
+use quote::{ToTokens, quote};
+use syn::{ItemFn, Meta, Token, parse::Parser, parse_macro_input, punctuated::Punctuated};
+
+/// Wraps the `#[ic_cdk::query]` macro and adds a statement to log the instructions used.
+#[proc_macro_attribute]
+pub fn query(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let attr = Punctuated::<Meta, Token![,]>::parse_terminated
+        .parse(attr)
+        .expect("Failed to parse attribute arguments");
+
+    let query_attr = if attr.is_empty() {
+        quote! { #[ic_cdk::query] }
+    } else {
+        let attrs = attr.iter();
+        quote! { #[ic_cdk::query(#(#attrs),*)] }
+    };
+
+    let item_fn = parse_macro_input!(item as ItemFn);
+    let fn_name = item_fn.sig.ident.to_string();
+
+    // Reconstruct the function so that we capture its execution and then log instructions
+    // at the very end before returning the result. No drop needed.
+    let attrs = &item_fn.attrs;
+    let vis = &item_fn.vis;
+    let sig = &item_fn.sig;
+    let block = &item_fn.block;
+
+    // Expand to the new function definition
+    let new_fn = quote! {
+        #(#attrs)*
+        #query_attr
+        #vis #sig {
+            let __result = {
+                #block
+            };
+
+            ic_cdk::println!(
+                "{}Instructions used by method {}: {}",
+                crate::LOG_PREFIX,
+                #fn_name,
+                ic_cdk::api::call_context_instruction_counter()
+            );
+
+            __result
+        }
+    };
+
+    TokenStream::from(new_fn.into_token_stream())
+}

rs/nervous_system/query_instruction_logger/tests/tests.rs

@@ -0,0 +1,59 @@
+use ic_nervous_system_query_instruction_logger::query;
+use std::cell::RefCell;
+
+pub const LOG_PREFIX: &str = "[TEST] ";
+
+thread_local! {
+    static WAS_CALLED: RefCell<Vec<&'static str>> = Default::default();
+}
+
+// Mocking the behavior so we don't actually invoke the real ic_cdk::println!
+// which can't easily be intercepted here. We just want to ensure it compiles
+// and basically structures the code correctly.
+mod ic_cdk {
+    pub use ::ic_cdk::query;
+
+    pub mod api {
+        pub fn call_context_instruction_counter() -> u64 {
+            super::super::WAS_CALLED
+                .with(|c| c.borrow_mut().push("call_context_instruction_counter"));
+            42
+        }
+    }
+
+    // We mock the println macro by just logging to our tracking list
+    #[macro_export]
+    macro_rules! mock_println {
+        ($fmt:expr $(, $arg:expr)* $(,)?) => {
+            $( let _ = $arg; )* // Evaluate all arguments
+            $crate::WAS_CALLED.with(|c| c.borrow_mut().push("println"));
+        }
+    }
+
+    pub use mock_println as println;
+}
+
+#[query]
+fn my_query_method() -> i32 {
+    WAS_CALLED.with(|c| c.borrow_mut().push("my_query_method"));
+    100
+}
+
+#[test]
+fn test_query_macro() {
+    let result = my_query_method();
+
+    assert_eq!(result, 100);
+
+    let observed_calls = WAS_CALLED.with(|c| c.borrow().clone());
+
+    // The macro should have injected call_context_instruction_counter and println
+    assert_eq!(
+        observed_calls,
+        vec![
+            "my_query_method",
+            "call_context_instruction_counter",
+            "println",
+        ],
+    );
+}

rs/nns/governance/BUILD.bazel

@@ -134,6 +134,7 @@ DEPENDENCIES = [
 MACRO_DEPENDENCIES = [
     # Keep sorted.
     "//rs/nervous_system/common/build_metadata",
+    "//rs/nervous_system/query_instruction_logger",
     "//rs/nns/governance/derive_self_describing",
     "@crate_index//:async-trait",
     "@crate_index//:rust_decimal_macros",

rs/nns/governance/Cargo.toml

@@ -49,6 +49,7 @@ ic-nervous-system-governance = { path = "../../nervous_system/governance" }
 ic-nervous-system-linear-map = { path = "../../nervous_system/linear_map" }
 ic-nervous-system-lock = { path = "../../nervous_system/lock" }
 ic-nervous-system-long-message = { path = "../../nervous_system/long_message" }
+ic-nervous-system-query-instruction-logger = { path = "../../nervous_system/query_instruction_logger" }
 ic-nervous-system-timer-task = { path = "../../nervous_system/timer_task" }
 ic-nervous-system-rate-limits = { path = "../../nervous_system/rate_limits" }
 ic-nervous-system-root = { path = "../../nervous_system/root" }

rs/nns/governance/canister/canister.rs

@@ -1,11 +1,12 @@
 use ic_base_types::PrincipalId;
-use ic_cdk::{heartbeat, init, post_upgrade, pre_upgrade, println, query, update};
+use ic_cdk::{heartbeat, init, post_upgrade, pre_upgrade, println, update};
 use ic_http_types::{HttpRequest, HttpResponse, HttpResponseBuilder};
 use ic_nervous_system_canisters::cmc::CMCCanister;
 use ic_nervous_system_common::{
     memory_manager_upgrade_storage::{load_protobuf, store_protobuf},
     serve_metrics,
 };
+use ic_nervous_system_query_instruction_logger::query;
 use ic_nervous_system_runtime::CdkRuntime;
 use ic_nns_common::{
     access_control::{check_caller_is_gtc, check_caller_is_ledger},