Skip to content

Commit

Permalink
refactor(nns): Switch NNS Governance global state from static to thre…
Browse files Browse the repository at this point in the history
…ad_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.
  • Loading branch information
jasonz-dfinity authored Jan 2, 2025
1 parent d19a1b4 commit 76a634c
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 17 deletions.
26 changes: 10 additions & 16 deletions rs/nns/governance/canister/canister.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Governance> = None;
thread_local! {
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()),
));
}

/*
Recommendations for Using `unsafe` in the Governance canister:
Expand Down Expand Up @@ -135,27 +136,20 @@ 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.
///
/// 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()
Expand Down
26 changes: 25 additions & 1 deletion rs/nns/governance/src/governance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1935,6 +1935,30 @@ fn spawn_in_canister_env(future: impl Future<Output = ()> + 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<dyn Environment>,
ledger: Box<dyn IcpLedger>,
cmc: Box<dyn CMC>,
) -> 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(
Expand Down

0 comments on commit 76a634c

Please sign in to comment.