feat(nns): Avoid recomputing wasm/arg hashes during read operations (#…

…3490)

# Why

Calculating hashes for wasm/arg actually takes quite a bit of
instructions (80/byte), which makes it the dominant part of the
list_proposals instructions. Such hashes can be calculated during write
operations easily, since the data should be immutable.

# What

* Add `wasm_module_hash` and `arg_hash` into `InstallCode` internal type
* Move hash calculation from read operation (internal -> response) to
write operation (request -> internal)
* Add backfill for existing `Installcode` proposals
* Update benchmark results (96% improvement on list_proposals)

rs/nns/governance/canbench/canbench_results.yml

@@ -1,55 +1,55 @@
 benches:
   add_neuron_active_maximum:
     total:
-      instructions: 42752805
+      instructions: 42749552
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_active_typical:
     total:
-      instructions: 2170667
+      instructions: 2170522
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_maximum:
     total:
-      instructions: 112624643
+      instructions: 112621384
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_typical:
     total:
-      instructions: 8497304
+      instructions: 8497153
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_all_heap:
     total:
-      instructions: 35002536
+      instructions: 34973512
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_heap_neurons_stable_index:
     total:
-      instructions: 61137575
+      instructions: 61108551
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_stable_everything:
     total:
-      instructions: 188621982
+      instructions: 188610325
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_stable_neurons_with_heap_index:
     total:
-      instructions: 162353547
+      instructions: 162341890
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   centralized_following_all_stable:
     total:
-      instructions: 78268135
+      instructions: 78265263
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
@@ -61,7 +61,7 @@ benches:
     scopes: {}
   draw_maturity_from_neurons_fund_heap:
     total:
-      instructions: 7598030
+      instructions: 7581730
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -79,13 +79,13 @@ benches:
     scopes: {}
   list_active_neurons_fund_neurons_stable:
     total:
-      instructions: 2750339
+      instructions: 2749439
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   list_neurons_heap:
     total:
-      instructions: 4880000
+      instructions: 4932713
       heap_increase: 9
       stable_memory_increase: 0
     scopes: {}
@@ -97,19 +97,19 @@ benches:
     scopes: {}
   list_neurons_ready_to_unstake_maturity_stable:
     total:
-      instructions: 41457102
+      instructions: 41457093
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   list_neurons_stable:
     total:
-      instructions: 113374022
+      instructions: 113659672
       heap_increase: 4
       stable_memory_increase: 0
     scopes: {}
   list_proposals:
     total:
-      instructions: 168095717
+      instructions: 6477719
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -121,19 +121,19 @@ benches:
     scopes: {}
   list_ready_to_spawn_neuron_ids_stable:
     total:
-      instructions: 41429005
+      instructions: 41428996
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_data_validation_heap:
     total:
-      instructions: 406864991
+      instructions: 406839584
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_data_validation_stable:
     total:
-      instructions: 362661286
+      instructions: 362638172
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -151,13 +151,13 @@ benches:
     scopes: {}
   range_neurons_performance:
     total:
-      instructions: 56448740
+      instructions: 56447340
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   single_vote_all_stable:
     total:
-      instructions: 2805838
+      instructions: 2803117
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}

rs/nns/governance/proto/ic_nns_governance/pb/v1/governance.proto

@@ -2130,6 +2130,10 @@ message InstallCode {
   optional bytes arg = 4;
   // Whether to skip stopping the canister before installing. Optional. Default is false.
   optional bool skip_stopping_before_installing = 5;
+  // The hash of the wasm module to install. Calculated from `wasm_module` when proposal is created.
+  optional bytes wasm_module_hash = 6;
+  // The hash of the arg to pass to the canister. Calculated from `arg` when proposal is created.
+  optional bytes arg_hash = 7;
 }
 
 message StopOrStartCanister {

rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs

@@ -2835,6 +2835,12 @@ pub struct InstallCode {
     /// Whether to skip stopping the canister before installing. Optional. Default is false.
     #[prost(bool, optional, tag = "5")]
     pub skip_stopping_before_installing: ::core::option::Option<bool>,
+    /// The hash of the wasm module to install. Calculated from `wasm_module` when proposal is created.
+    #[prost(bytes = "vec", optional, tag = "6")]
+    pub wasm_module_hash: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
+    /// The hash of the arg to pass to the canister. Calculated from `arg` when proposal is created.
+    #[prost(bytes = "vec", optional, tag = "7")]
+    pub arg_hash: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
 }
 /// Nested message and enum types in `InstallCode`.
 pub mod install_code {

rs/nns/governance/src/governance.rs

@@ -74,6 +74,7 @@ use ic_base_types::{CanisterId, PrincipalId};
 use ic_cdk::println;
 #[cfg(target_arch = "wasm32")]
 use ic_cdk::spawn;
+use ic_crypto_sha2::Sha256;
 use ic_nervous_system_common::{
     cmc::CMC, ledger, ledger::IcpLedger, NervousSystemError, ONE_DAY_SECONDS, ONE_MONTH_SECONDS,
     ONE_YEAR_SECONDS,
@@ -1901,7 +1902,7 @@ impl Governance {
             env.seed_rng(rng_seed);
         }
 
-        Self {
+        let mut governance = Self {
             heap_data: heap_governance_proto,
             neuron_store: NeuronStore::new_restored((heap_neurons, topic_followee_map)),
             env,
@@ -1913,7 +1914,10 @@ impl Governance {
             neuron_data_validator: NeuronDataValidator::new(),
             minting_node_provider_rewards: false,
             neuron_rate_limits: NeuronRateLimits::default(),
-        }
+        };
+        // TODO: Remove after the backfill has been run once.
+        governance.backfill_install_code_hashes();
+        governance
     }
 
     /// After calling this method, the proto and neuron_store (the heap neurons at least)
@@ -6105,6 +6109,38 @@ impl Governance {
         backfill_some_voting_power_refreshed_timestamps(&mut self.neuron_store, begin, carry_on)
     }
 
+    pub fn backfill_install_code_hashes(&mut self) {
+        for proposal in self.heap_data.proposals.values_mut() {
+            let Some(proposal) = proposal.proposal.as_mut() else {
+                continue;
+            };
+            let Some(Action::InstallCode(install_code)) = proposal.action.as_mut() else {
+                continue;
+            };
+            if install_code.wasm_module_hash.is_none() {
+                if let Some(wasm_module) = install_code.wasm_module.as_ref() {
+                    install_code.wasm_module_hash = Some(Sha256::hash(wasm_module).to_vec());
+                }
+            }
+            if install_code.arg_hash.is_none() {
+                let arg_hash = match install_code.arg.as_ref() {
+                    Some(arg) => {
+                        // We could calculate the hash of an empty arg, but it would be confusing for the
+                        // proposal reviewers, since the arg_hash is the only thing they can see, and it would
+                        // not be obvious that the arg is empty.
+                        if arg.is_empty() {
+                            Some(vec![])
+                        } else {
+                            Some(Sha256::hash(arg).to_vec())
+                        }
+                    }
+                    None => Some(vec![]),
+                };
+                install_code.arg_hash = arg_hash;
+            }
+        }
+    }
+
     /// Creates a new neuron or refreshes the stake of an existing
     /// neuron from a ledger account.
     ///

rs/nns/governance/src/governance/benches.rs

@@ -22,6 +22,7 @@ use crate::{
 use canbench_rs::{bench, bench_fn, BenchResult};
 use futures::FutureExt;
 use ic_base_types::PrincipalId;
+use ic_crypto_sha2::Sha256;
 use ic_nervous_system_proto::pb::v1::Image;
 use ic_nns_common::{
     pb::v1::{NeuronId as NeuronIdProto, ProposalId},
@@ -644,6 +645,8 @@ fn list_proposals_benchmark() -> BenchResult {
             wasm_module: Some(vec![0u8; 1 << 20]), // 1 MiB
             arg: Some(vec![0u8; 1 << 20]),         // 1 MiB
             install_mode: Some(CanisterInstallMode::Install as i32),
+            wasm_module_hash: Some(Sha256::hash(&vec![0u8; 1 << 20]).to_vec()),
+            arg_hash: Some(Sha256::hash(&vec![0u8; 1 << 20]).to_vec()),
             skip_stopping_before_installing: None,
         }),
         Action::CreateServiceNervousSystem(

rs/nns/governance/src/pb/conversions.rs

@@ -1,4 +1,5 @@
 use crate::pb::v1 as pb;
+use ic_crypto_sha2::Sha256;
 use ic_nns_governance_api::pb::v1 as pb_api;
 
 impl From<pb::NodeProvider> for pb_api::NodeProvider {
@@ -2541,22 +2542,12 @@ impl From<pb_api::create_service_nervous_system::governance_parameters::VotingRe
 
 impl From<pb::InstallCode> for pb_api::InstallCode {
     fn from(item: pb::InstallCode) -> Self {
-        let wasm_module_hash = item
-            .wasm_module
-            .map(|wasm_module| super::calculate_hash(&wasm_module).to_vec());
-        let arg = item.arg.unwrap_or_default();
-        let arg_hash = if arg.is_empty() {
-            Some(vec![])
-        } else {
-            Some(super::calculate_hash(&arg).to_vec())
-        };
-
         Self {
             canister_id: item.canister_id,
             install_mode: item.install_mode,
             skip_stopping_before_installing: item.skip_stopping_before_installing,
-            wasm_module_hash,
-            arg_hash,
+            wasm_module_hash: item.wasm_module_hash,
+            arg_hash: item.arg_hash,
         }
     }
 }
@@ -2570,17 +2561,39 @@ impl From<pb_api::InstallCode> for pb::InstallCode {
             // canister_init.
             wasm_module: None,
             arg: None,
+            wasm_module_hash: item.wasm_module_hash,
+            arg_hash: item.arg_hash,
         }
     }
 }
 impl From<pb_api::InstallCodeRequest> for pb::InstallCode {
     fn from(item: pb_api::InstallCodeRequest) -> Self {
+        let wasm_module_hash = item
+            .wasm_module
+            .as_ref()
+            .map(|wasm_module| Sha256::hash(wasm_module).to_vec());
+        let arg_hash = match item.arg.as_ref() {
+            Some(arg) => {
+                // We could calculate the hash of an empty arg, but it would be confusing for the
+                // proposal reviewers, since the arg_hash is the only thing they can see, and it would
+                // not be obvious that the arg is empty.
+                if arg.is_empty() {
+                    Some(vec![])
+                } else {
+                    Some(Sha256::hash(arg).to_vec())
+                }
+            }
+            None => Some(vec![]),
+        };
+
         Self {
             canister_id: item.canister_id,
             install_mode: item.install_mode,
             wasm_module: item.wasm_module,
             arg: item.arg,
             skip_stopping_before_installing: item.skip_stopping_before_installing,
+            wasm_module_hash,
+            arg_hash,
         }
     }
 }

rs/nns/governance/src/pb/mod.rs

@@ -1,5 +1,4 @@
 use crate::pb::v1::ArchivedMonthlyNodeProviderRewards;
-use ic_crypto_sha2::Sha256;
 use ic_stable_structures::{storable::Bound, Storable};
 use prost::Message;
 use std::borrow::Cow;
@@ -26,12 +25,5 @@ impl Storable for ArchivedMonthlyNodeProviderRewards {
     const BOUND: Bound = Bound::Unbounded;
 }
 
-/// Calculates the SHA256 hash of the given bytes.
-fn calculate_hash(bytes: &[u8]) -> [u8; 32] {
-    let mut wasm_sha = Sha256::new();
-    wasm_sha.write(bytes);
-    wasm_sha.finish()
-}
-
 #[cfg(test)]
 mod tests;