Skip to content

Commit 76a634c

Browse files
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<RefCell<Governance>>`) while returning the raw pointer for the existing access pattern. # Why it is safe The `LocalKey<RefCell<T>>` 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<Governance>`. 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<Governance>` and `unwrap()` every time, but it seems more cumbersome.
1 parent d19a1b4 commit 76a634c

File tree

2 files changed

+35
-17
lines changed

2 files changed

+35
-17
lines changed

rs/nns/governance/canister/canister.rs

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,13 @@ const WASM_PAGES_RESERVED_FOR_UPGRADES_MEMORY: u64 = 65_536;
8181

8282
pub(crate) const LOG_PREFIX: &str = "[Governance] ";
8383

84-
// https://dfinity.atlassian.net/browse/NNS1-1050: We are not following
85-
// standard/best practices for canister globals here.
86-
//
87-
// Do not access these global variables directly. Instead, use accessor
88-
// functions, which are defined immediately after.
89-
static mut GOVERNANCE: Option<Governance> = None;
84+
thread_local! {
85+
static GOVERNANCE: RefCell<Governance> = RefCell::new(Governance::new_uninitialized(
86+
Box::new(CanisterEnv::new()),
87+
Box::new(IcpLedgerCanister::<CdkRuntime>::new(LEDGER_CANISTER_ID)),
88+
Box::new(CMCCanister::<CdkRuntime>::new()),
89+
));
90+
}
9091

9192
/*
9293
Recommendations for Using `unsafe` in the Governance canister:
@@ -135,27 +136,20 @@ are best practices for making use of the unsafe block:
135136
/// This should only be called once the global state has been initialized, which
136137
/// happens in `canister_init` or `canister_post_upgrade`.
137138
fn governance() -> &'static Governance {
138-
unsafe { GOVERNANCE.as_ref().expect("Canister not initialized!") }
139+
unsafe { &*GOVERNANCE.with(|g| g.as_ptr()) }
139140
}
140141

141142
/// Returns a mutable reference to the global state.
142143
///
143144
/// This should only be called once the global state has been initialized, which
144145
/// happens in `canister_init` or `canister_post_upgrade`.
145146
fn governance_mut() -> &'static mut Governance {
146-
unsafe { GOVERNANCE.as_mut().expect("Canister not initialized!") }
147+
unsafe { &mut *GOVERNANCE.with(|g| g.as_ptr()) }
147148
}
148149

149150
// Sets governance global state to the given object.
150151
fn set_governance(gov: Governance) {
151-
unsafe {
152-
assert!(
153-
GOVERNANCE.is_none(),
154-
"{}Trying to initialize an already-initialized governance canister!",
155-
LOG_PREFIX
156-
);
157-
GOVERNANCE = Some(gov);
158-
}
152+
GOVERNANCE.set(gov);
159153

160154
governance()
161155
.validate()

rs/nns/governance/src/governance.rs

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ use rust_decimal_macros::dec;
117117
use std::{
118118
borrow::Cow,
119119
cmp::{max, Ordering},
120-
collections::{HashMap, HashSet},
120+
collections::{BTreeMap, HashMap, HashSet},
121121
convert::{TryFrom, TryInto},
122122
fmt,
123123
future::Future,
@@ -1935,6 +1935,30 @@ fn spawn_in_canister_env(future: impl Future<Output = ()> + Sized + 'static) {
19351935
}
19361936

19371937
impl Governance {
1938+
/// Creates a new Governance instance with uninitialized fields. The canister should only have
1939+
/// such state before the state is recovered from the stable memory in post_upgrade or
1940+
/// initialized in init. In any other case, the `Governance` object should be initialized with
1941+
/// either `new` or `new_restored`.
1942+
pub fn new_uninitialized(
1943+
env: Box<dyn Environment>,
1944+
ledger: Box<dyn IcpLedger>,
1945+
cmc: Box<dyn CMC>,
1946+
) -> Self {
1947+
Self {
1948+
heap_data: HeapGovernanceData::default(),
1949+
neuron_store: NeuronStore::new(BTreeMap::new()),
1950+
env,
1951+
ledger,
1952+
cmc,
1953+
closest_proposal_deadline_timestamp_seconds: 0,
1954+
latest_gc_timestamp_seconds: 0,
1955+
latest_gc_num_proposals: 0,
1956+
neuron_data_validator: NeuronDataValidator::new(),
1957+
minting_node_provider_rewards: false,
1958+
neuron_rate_limits: NeuronRateLimits::default(),
1959+
}
1960+
}
1961+
19381962
/// Initializes Governance for the first time from init payload. When restoring after an upgrade
19391963
/// with its persisted state, `Governance::new_restored` should be called instead.
19401964
pub fn new(

0 commit comments

Comments
 (0)