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)
dfinity · Jan 22, 2025 · b0fd3c2 · b0fd3c2
1 parent 7505d4e
commit b0fd3c2
Show file tree

Hide file tree

Showing 11 changed files with 123 additions and 50 deletions.
diff --git a/rs/nns/governance/canbench/canbench_results.yml b/rs/nns/governance/canbench/canbench_results.yml
@@ -1,67 +1,67 @@
 benches:
   add_neuron_active_maximum:
     total:
-      instructions: 42752805
+      instructions: 42749561
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_active_typical:
     total:
-      instructions: 2170667
+      instructions: 2170531
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_maximum:
     total:
-      instructions: 112624643
+      instructions: 112621399
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
   add_neuron_inactive_typical:
     total:
-      instructions: 8497304
+      instructions: 8497168
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   cascading_vote_all_heap:
     total:
-      instructions: 35002536
+      instructions: 34979056
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_heap_neurons_stable_index:
     total:
-      instructions: 61137575
+      instructions: 61114095
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_stable_everything:
     total:
-      instructions: 188621982
+      instructions: 188616886
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   cascading_vote_stable_neurons_with_heap_index:
     total:
-      instructions: 162353547
+      instructions: 162348451
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   centralized_following_all_stable:
     total:
-      instructions: 78268135
+      instructions: 78267527
       heap_increase: 0
       stable_memory_increase: 128
     scopes: {}
   compute_ballots_for_new_proposal_with_stable_neurons:
     total:
-      instructions: 2169168
+      instructions: 2222211
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   draw_maturity_from_neurons_fund_heap:
     total:
-      instructions: 7598030
+      instructions: 7584430
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -103,13 +103,13 @@ benches:
     scopes: {}
   list_neurons_stable:
     total:
-      instructions: 113374022
+      instructions: 113372722
       heap_increase: 5
       stable_memory_increase: 0
     scopes: {}
   list_proposals:
     total:
-      instructions: 168095717
+      instructions: 6477736
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -127,13 +127,13 @@ benches:
     scopes: {}
   neuron_data_validation_heap:
     total:
-      instructions: 406864991
+      instructions: 406848584
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   neuron_data_validation_stable:
     total:
-      instructions: 362661286
+      instructions: 362648972
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -145,19 +145,19 @@ benches:
     scopes: {}
   neuron_metrics_calculation_stable:
     total:
-      instructions: 3027095
+      instructions: 3021895
       heap_increase: 0
       stable_memory_increase: 0
     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: {}

diff --git a/rs/nns/governance/proto/ic_nns_governance/pb/v1/governance.proto b/rs/nns/governance/proto/ic_nns_governance/pb/v1/governance.proto
@@ -2212,6 +2212,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 {

diff --git a/rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs b/rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs
@@ -3086,6 +3086,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 {

diff --git a/rs/nns/governance/src/governance.rs b/rs/nns/governance/src/governance.rs
@@ -77,6 +77,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,
@@ -2051,7 +2052,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,
@@ -2063,7 +2064,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)
@@ -6247,6 +6251,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.
     ///

diff --git a/rs/nns/governance/src/governance/benches.rs b/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},
@@ -643,6 +644,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(

diff --git a/rs/nns/governance/src/pb/conversions.rs b/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 {
@@ -2896,22 +2897,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,
         }
     }
 }
@@ -2925,17 +2916,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,
         }
     }
 }

diff --git a/rs/nns/governance/src/pb/mod.rs b/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;
diff --git a/rs/nns/governance/src/pb/tests.rs b/rs/nns/governance/src/pb/tests.rs
@@ -1,47 +1,50 @@
-use super::*;
-
 use crate::pb::v1 as pb;
+
 use ic_base_types::PrincipalId;
+use ic_crypto_sha2::Sha256;
 use ic_nns_governance_api::pb::v1 as pb_api;
 
 #[test]
-fn install_code_internal_to_api() {
+fn install_code_request_to_internal() {
     let test_cases = vec![
         (
-            pb::InstallCode {
+            pb_api::InstallCodeRequest {
                 canister_id: Some(PrincipalId::new_user_test_id(1)),
                 install_mode: Some(pb::install_code::CanisterInstallMode::Install as i32),
                 skip_stopping_before_installing: None,
                 wasm_module: Some(vec![1, 2, 3]),
                 arg: Some(vec![]),
             },
-            pb_api::InstallCode {
+            pb::InstallCode {
                 canister_id: Some(PrincipalId::new_user_test_id(1)),
                 install_mode: Some(pb_api::install_code::CanisterInstallMode::Install as i32),
                 skip_stopping_before_installing: None,
+                wasm_module: Some(vec![1, 2, 3]),
+                arg: Some(vec![]),
                 wasm_module_hash: Some(Sha256::hash(&[1, 2, 3]).to_vec()),
                 arg_hash: Some(vec![]),
             },
         ),
         (
-            pb::InstallCode {
+            pb_api::InstallCodeRequest {
                 canister_id: Some(PrincipalId::new_user_test_id(1)),
                 install_mode: Some(pb::install_code::CanisterInstallMode::Upgrade as i32),
                 skip_stopping_before_installing: Some(true),
                 wasm_module: Some(vec![1, 2, 3]),
                 arg: Some(vec![4, 5, 6]),
             },
-            pb_api::InstallCode {
+            pb::InstallCode {
                 canister_id: Some(PrincipalId::new_user_test_id(1)),
                 install_mode: Some(pb_api::install_code::CanisterInstallMode::Upgrade as i32),
                 skip_stopping_before_installing: Some(true),
+                wasm_module: Some(vec![1, 2, 3]),
+                arg: Some(vec![4, 5, 6]),
                 wasm_module_hash: Some(Sha256::hash(&[1, 2, 3]).to_vec()),
                 arg_hash: Some(Sha256::hash(&[4, 5, 6]).to_vec()),
             },
         ),
     ];
-
-    for (internal, api) in test_cases {
-        assert_eq!(pb_api::InstallCode::from(internal), api);
+    for (request, internal) in test_cases {
+        assert_eq!(pb::InstallCode::from(request), internal);
     }
 }