refactor(nns): Make Governance environment parameters compatible with…

… timer usages (#4161)

# What

Move the injected dependencies in Governance to be Arc instead of Box

# Why
To allow for unit tests to continue testing the overall state of
Governance in relation to timer execution so existing tests continue to
work.

rs/nervous_system/canisters/src/cmc.rs

@@ -28,7 +28,7 @@ impl<Rt: Runtime> Default for CMCCanister<Rt> {
 #[async_trait]
 impl<Rt: Runtime + Send + Sync> CMC for CMCCanister<Rt> {
     /// Returns the maturity_modulation from the CMC in basis points.
-    async fn neuron_maturity_modulation(&mut self) -> Result<i32, String> {
+    async fn neuron_maturity_modulation(&self) -> Result<i32, String> {
         let result: Result<(Result<i32, String>,), (i32, String)> =
             Rt::call_with_cleanup(self.canister_id, "neuron_maturity_modulation", ()).await;
         match result {

rs/nervous_system/common/src/cmc.rs

@@ -6,7 +6,7 @@ use mockall::automock;
 #[async_trait]
 pub trait CMC: Send + Sync {
     /// Returns the current neuron maturity modulation.
-    async fn neuron_maturity_modulation(&mut self) -> Result<i32, String>;
+    async fn neuron_maturity_modulation(&self) -> Result<i32, String>;
 }
 
 pub struct FakeCmc {}
@@ -25,7 +25,7 @@ impl Default for FakeCmc {
 
 #[async_trait]
 impl CMC for FakeCmc {
-    async fn neuron_maturity_modulation(&mut self) -> Result<i32, String> {
+    async fn neuron_maturity_modulation(&self) -> Result<i32, String> {
         Ok(0)
     }
 }

rs/nns/governance/benches/scale.rs

@@ -19,6 +19,7 @@ use futures::future::FutureExt;
 use ic_base_types::{CanisterId, PrincipalId};
 use ic_nervous_system_common::{cmc::FakeCmc, ledger::IcpLedger, NervousSystemError};
 use ic_nns_common::pb::v1::NeuronId;
+use ic_nns_governance::governance::RandomnessGenerator;
 use ic_nns_governance::{
     governance::{Environment, Governance, HeapGrowthPotential, RngError},
     pb::v1::{
@@ -28,6 +29,7 @@ use ic_nns_governance::{
 };
 use icp_ledger::{AccountIdentifier, Subaccount, Tokens};
 use std::convert::TryFrom;
+use std::sync::Arc;
 
 criterion_group! {
     name = benches;
@@ -37,17 +39,15 @@ criterion_group! {
 
 criterion_main!(benches);
 
-/// Mock of the required interface of `Governance`.
-struct MockEnvironment {
-    secs: u64,
-}
+struct FakeRandomness {}
 
-#[async_trait]
-impl Environment for MockEnvironment {
-    fn now(&self) -> u64 {
-        self.secs
+impl FakeRandomness {
+    fn new() -> Self {
+        Self {}
     }
+}
 
+impl RandomnessGenerator for FakeRandomness {
     fn random_u64(&mut self) -> Result<u64, RngError> {
         todo!()
     }
@@ -63,6 +63,18 @@ impl Environment for MockEnvironment {
     fn get_rng_seed(&self) -> Option<[u8; 32]> {
         todo!()
     }
+}
+
+/// Mock of the required interface of `Governance`.
+struct MockEnvironment {
+    secs: u64,
+}
+
+#[async_trait]
+impl Environment for MockEnvironment {
+    fn now(&self) -> u64 {
+        self.secs
+    }
 
     fn execute_nns_function(
         &self,
@@ -145,9 +157,10 @@ fn linear_20k(c: &mut Criterion) {
     let secs = 1;
     let mut gov = Governance::new(
         fixture_for_scale(20_000, true),
-        Box::new(MockEnvironment { secs }),
-        Box::new(MockLedger {}),
-        Box::new(FakeCmc::new()),
+        Arc::new(MockEnvironment { secs }),
+        Arc::new(MockLedger {}),
+        Arc::new(FakeCmc::new()),
+        Box::new(FakeRandomness::new()),
     );
     c.bench_function("linear 20k", |b| {
         b.iter(|| make_and_process_proposal(&mut gov))
@@ -158,9 +171,10 @@ fn tree_20k(c: &mut Criterion) {
     let secs = 1;
     let mut gov = Governance::new(
         fixture_for_scale(20_000, false),
-        Box::new(MockEnvironment { secs }),
-        Box::new(MockLedger {}),
-        Box::new(FakeCmc::new()),
+        Arc::new(MockEnvironment { secs }),
+        Arc::new(MockLedger {}),
+        Arc::new(FakeCmc::new()),
+        Box::new(FakeRandomness::new()),
     );
     c.bench_function("tree 20k", |b| {
         b.iter(|| make_and_process_proposal(&mut gov))
@@ -171,9 +185,10 @@ fn linear_200k(c: &mut Criterion) {
     let secs = 1;
     let mut gov = Governance::new(
         fixture_for_scale(200_000, true),
-        Box::new(MockEnvironment { secs }),
-        Box::new(MockLedger {}),
-        Box::new(FakeCmc::new()),
+        Arc::new(MockEnvironment { secs }),
+        Arc::new(MockLedger {}),
+        Arc::new(FakeCmc::new()),
+        Box::new(FakeRandomness::new()),
     );
     c.bench_function("linear 200k", |b| {
         b.iter(|| make_and_process_proposal(&mut gov))
@@ -184,9 +199,10 @@ fn tree_200k(c: &mut Criterion) {
     let secs = 1;
     let mut gov = Governance::new(
         fixture_for_scale(200_000, false),
-        Box::new(MockEnvironment { secs }),
-        Box::new(MockLedger {}),
-        Box::new(FakeCmc::new()),
+        Arc::new(MockEnvironment { secs }),
+        Arc::new(MockLedger {}),
+        Arc::new(FakeCmc::new()),
+        Box::new(FakeRandomness::new()),
     );
     c.bench_function("tree 200k", |b| {
         b.iter(|| make_and_process_proposal(&mut gov))

rs/nns/governance/canister/canister.rs

@@ -55,11 +55,12 @@ use ic_nns_governance_api::test_api::TimeWarp;
 use prost::Message;
 use rand::{RngCore, SeedableRng};
 use rand_chacha::ChaCha20Rng;
+use std::sync::Arc;
 use std::{boxed::Box, ops::Bound, time::Duration};
 
 #[cfg(not(feature = "tla"))]
 use ic_nervous_system_canisters::ledger::IcpLedgerCanister;
-use ic_nns_governance::canister_state::with_governance;
+use ic_nns_governance::canister_state::{with_governance, CanisterRandomnessGenerator};
 
 #[cfg(feature = "tla")]
 mod tla_ledger;
@@ -230,9 +231,10 @@ fn canister_init_(init_payload: ApiGovernanceProto) {
     let governance_proto = InternalGovernanceProto::from(init_payload);
     set_governance(Governance::new(
         governance_proto,
-        Box::new(CanisterEnv::new()),
-        Box::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
-        Box::new(CMCCanister::<CdkRuntime>::new()),
+        Arc::new(CanisterEnv::new()),
+        Arc::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
+        Arc::new(CMCCanister::<CdkRuntime>::new()),
+        Box::new(CanisterRandomnessGenerator::new()),
     ));
 }
 
@@ -274,9 +276,10 @@ fn canister_post_upgrade() {
     schedule_timers();
     set_governance(Governance::new_restored(
         restored_state,
-        Box::new(CanisterEnv::new()),
-        Box::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
-        Box::new(CMCCanister::<CdkRuntime>::new()),
+        Arc::new(CanisterEnv::new()),
+        Arc::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
+        Arc::new(CMCCanister::<CdkRuntime>::new()),
+        Box::new(CanisterRandomnessGenerator::new()),
     ));
 
     validate_stable_storage();

rs/nns/governance/src/canister_state.rs

@@ -1,5 +1,7 @@
 use crate::decoder_config;
-use crate::governance::{Environment, Governance, HeapGrowthPotential, RngError};
+use crate::governance::{
+    Environment, Governance, HeapGrowthPotential, RandomnessGenerator, RngError,
+};
 use async_trait::async_trait;
 use candid::{Decode, Encode};
 use ic_base_types::CanisterId;
@@ -21,12 +23,16 @@ use rand::{RngCore, SeedableRng};
 use rand_chacha::ChaCha20Rng;
 use std::cell::RefCell;
 use std::str::FromStr;
+use std::sync::Arc;
+#[cfg(any(test, feature = "test"))]
+use std::sync::RwLock;
 
 thread_local! {
     pub(crate) static GOVERNANCE: RefCell<Governance> = RefCell::new(Governance::new_uninitialized(
-        Box::new(CanisterEnv::new()),
-        Box::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
-        Box::new(CMCCanister::<CdkRuntime>::new()),
+        Arc::new(CanisterEnv::new()),
+        Arc::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
+        Arc::new(CMCCanister::<CdkRuntime>::new()),
+        Box::new(CanisterRandomnessGenerator::new()),
     ));
 }
 /*
@@ -106,29 +112,31 @@ pub fn set_governance(gov: Governance) {
 
 #[derive(Default)]
 pub struct CanisterEnv {
-    rng: Option<ChaCha20Rng>,
-    time_warp: crate::governance::TimeWarp,
+    #[cfg(any(test, feature = "test"))]
+    time_warp: RwLock<crate::governance::TimeWarp>,
 }
 
 impl CanisterEnv {
     pub fn new() -> Self {
         CanisterEnv {
-            rng: None,
-            time_warp: crate::governance::TimeWarp { delta_s: 0 },
+            #[cfg(any(test, feature = "test"))]
+            time_warp: RwLock::new(crate::governance::TimeWarp { delta_s: 0 }),
         }
     }
 }
 
-#[async_trait]
-impl Environment for CanisterEnv {
-    fn now(&self) -> u64 {
-        self.time_warp.apply(now_seconds())
-    }
+#[derive(Default)]
+pub struct CanisterRandomnessGenerator {
+    rng: Option<ChaCha20Rng>,
+}
 
-    fn set_time_warp(&mut self, new_time_warp: crate::governance::TimeWarp) {
-        self.time_warp = new_time_warp;
+impl CanisterRandomnessGenerator {
+    pub fn new() -> Self {
+        CanisterRandomnessGenerator { rng: None }
     }
+}
 
+impl RandomnessGenerator for CanisterRandomnessGenerator {
     fn random_u64(&mut self) -> Result<u64, RngError> {
         match self.rng.as_mut() {
             Some(rand) => Ok(rand.next_u64()),
@@ -154,6 +162,24 @@ impl Environment for CanisterEnv {
     fn get_rng_seed(&self) -> Option<[u8; 32]> {
         self.rng.as_ref().map(|rng| rng.get_seed())
     }
+}
+
+#[async_trait]
+impl Environment for CanisterEnv {
+    #[cfg(any(test, feature = "test"))]
+    fn now(&self) -> u64 {
+        self.time_warp.read().unwrap().apply(now_seconds())
+    }
+
+    #[cfg(not(any(test, feature = "test")))]
+    fn now(&self) -> u64 {
+        now_seconds()
+    }
+
+    #[cfg(any(test, feature = "test"))]
+    fn set_time_warp(&self, new_time_warp: crate::governance::TimeWarp) {
+        *self.time_warp.write().unwrap() = new_time_warp;
+    }
 
     fn execute_nns_function(
         &self,
@@ -420,7 +446,7 @@ mod tests {
     use candid::{Decode, Encode};
     #[test]
     fn test_set_time_warp() {
-        let mut environment = CanisterEnv::new();
+        let environment = CanisterEnv::new();
 
         let start = environment.now();
         environment.set_time_warp(crate::governance::TimeWarp { delta_s: 1_000 });