From 76a634c31dfb840da25fbe286855eb0be1818ca8 Mon Sep 17 00:00:00 2001 From: jasonz-dfinity <133917836+jasonz-dfinity@users.noreply.github.com> Date: Thu, 2 Jan 2025 11:25:07 -0800 Subject: [PATCH] refactor(nns): Switch NNS Governance global state from static to thread_local (#2844) # Why `governance()` and `governance_mut()` are bad Representing the canister global state with `thread_local!` avoids using unsafe blocks to access it. When using unsafe blocks to access it, one can easily write code with undefined behavior by retaining references across await boundary (more precisely, after an inter-canister call). # Why this PR The NNS Governance heavily depends on the accessing the global state as `static`, and there will be a lot of refactoring needed in order to get away with the current pattern. This PR is the first step towards getting rid of those bad access patterns - with this change, we can gradually migrate from using `governance()`/`governance_mut()` to using `GOVERNANCE.with()`. When all the usage of `governance()`/`governance_mut()` are replaced, we can delete them and declare victory. # What Define the global state with `thread_local!` (`LocalKey>`) while returning the raw pointer for the existing access pattern. # Why it is safe The `LocalKey>` is set once and only once during `post_upgrade` or `init`, so the pointer should remain constant, given that the canister code is single threaded. When accessing through `governance()` or `governance_mut()` one can still write code with undefined behavior, but it is the same danger as what we have now. # Why `Governance::new_uninitialized` This is needed in order for the `thread_local!` state to be `Governance` instead of `Option`. We know the "uninitialized" version won't be used except for the code in the init/post_upgrade before the "initialized" state is set. However, there doesn't seem to be a good way to express that understanding. An alternative is to still use `Option` and `unwrap()` every time, but it seems more cumbersome. --- rs/nns/governance/canister/canister.rs | 26 ++++++++++---------------- rs/nns/governance/src/governance.rs | 26 +++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/rs/nns/governance/canister/canister.rs b/rs/nns/governance/canister/canister.rs index d2e351bf410..0d343ebe194 100644 --- a/rs/nns/governance/canister/canister.rs +++ b/rs/nns/governance/canister/canister.rs @@ -81,12 +81,13 @@ const WASM_PAGES_RESERVED_FOR_UPGRADES_MEMORY: u64 = 65_536; pub(crate) const LOG_PREFIX: &str = "[Governance] "; -// https://dfinity.atlassian.net/browse/NNS1-1050: We are not following -// standard/best practices for canister globals here. -// -// Do not access these global variables directly. Instead, use accessor -// functions, which are defined immediately after. -static mut GOVERNANCE: Option = None; +thread_local! { + static GOVERNANCE: RefCell = RefCell::new(Governance::new_uninitialized( + Box::new(CanisterEnv::new()), + Box::new(IcpLedgerCanister::::new(LEDGER_CANISTER_ID)), + Box::new(CMCCanister::::new()), + )); +} /* Recommendations for Using `unsafe` in the Governance canister: @@ -135,7 +136,7 @@ are best practices for making use of the unsafe block: /// This should only be called once the global state has been initialized, which /// happens in `canister_init` or `canister_post_upgrade`. fn governance() -> &'static Governance { - unsafe { GOVERNANCE.as_ref().expect("Canister not initialized!") } + unsafe { &*GOVERNANCE.with(|g| g.as_ptr()) } } /// Returns a mutable reference to the global state. @@ -143,19 +144,12 @@ fn governance() -> &'static Governance { /// This should only be called once the global state has been initialized, which /// happens in `canister_init` or `canister_post_upgrade`. fn governance_mut() -> &'static mut Governance { - unsafe { GOVERNANCE.as_mut().expect("Canister not initialized!") } + unsafe { &mut *GOVERNANCE.with(|g| g.as_ptr()) } } // Sets governance global state to the given object. fn set_governance(gov: Governance) { - unsafe { - assert!( - GOVERNANCE.is_none(), - "{}Trying to initialize an already-initialized governance canister!", - LOG_PREFIX - ); - GOVERNANCE = Some(gov); - } + GOVERNANCE.set(gov); governance() .validate() diff --git a/rs/nns/governance/src/governance.rs b/rs/nns/governance/src/governance.rs index 28b809f7cd5..d5829c24be5 100644 --- a/rs/nns/governance/src/governance.rs +++ b/rs/nns/governance/src/governance.rs @@ -117,7 +117,7 @@ use rust_decimal_macros::dec; use std::{ borrow::Cow, cmp::{max, Ordering}, - collections::{HashMap, HashSet}, + collections::{BTreeMap, HashMap, HashSet}, convert::{TryFrom, TryInto}, fmt, future::Future, @@ -1935,6 +1935,30 @@ fn spawn_in_canister_env(future: impl Future + Sized + 'static) { } impl Governance { + /// Creates a new Governance instance with uninitialized fields. The canister should only have + /// such state before the state is recovered from the stable memory in post_upgrade or + /// initialized in init. In any other case, the `Governance` object should be initialized with + /// either `new` or `new_restored`. + pub fn new_uninitialized( + env: Box, + ledger: Box, + cmc: Box, + ) -> Self { + Self { + heap_data: HeapGovernanceData::default(), + neuron_store: NeuronStore::new(BTreeMap::new()), + env, + ledger, + cmc, + closest_proposal_deadline_timestamp_seconds: 0, + latest_gc_timestamp_seconds: 0, + latest_gc_num_proposals: 0, + neuron_data_validator: NeuronDataValidator::new(), + minting_node_provider_rewards: false, + neuron_rate_limits: NeuronRateLimits::default(), + } + } + /// Initializes Governance for the first time from init payload. When restoring after an upgrade /// with its persisted state, `Governance::new_restored` should be called instead. pub fn new(