feat(IC-1579): Governance: TLA Codelink for refresh_neuron (dfinity#3547)

andrew-lee-work · oggy-dfin · web-flow · commit 3afdf372a91d · 2025-02-04T16:22:23.000Z
Instrument refresh_neuron and its tests for TLA Code Link, add TLA model
for refresh neuron.

---------

Co-authored-by: oggy-dfin &lt;89794951+oggy-dfin@users.noreply.github.com&gt;
diff --git a/rs/nns/governance/src/governance.rs b/rs/nns/governance/src/governance.rs
@@ -154,8 +154,9 @@ use crate::storage::with_voting_state_machines_mut;
 #[cfg(feature = "tla")]
 pub use tla::{
     tla_update_method, InstrumentationState, ToTla, CLAIM_NEURON_DESC, DISBURSE_NEURON_DESC,
-    DISBURSE_TO_NEURON_DESC, MERGE_NEURONS_DESC, SPAWN_NEURONS_DESC, SPAWN_NEURON_DESC,
-    SPLIT_NEURON_DESC, TLA_INSTRUMENTATION_STATE, TLA_TRACES_LKEY, TLA_TRACES_MUTEX,
+    DISBURSE_TO_NEURON_DESC, MERGE_NEURONS_DESC, REFRESH_NEURON_DESC, SPAWN_NEURONS_DESC,
+    SPAWN_NEURON_DESC, SPLIT_NEURON_DESC, TLA_INSTRUMENTATION_STATE, TLA_TRACES_LKEY,
+    TLA_TRACES_MUTEX,
 };
 
 // 70 KB (for executing NNS functions that are not canister upgrades)
@@ -6077,6 +6078,7 @@ impl Governance {
     }
 
     /// Refreshes the stake of a given neuron by checking it's account.
+    #[cfg_attr(feature = "tla", tla_update_method(REFRESH_NEURON_DESC.clone()))]
     async fn refresh_neuron(
         &mut self,
         nid: NeuronId,
@@ -6099,6 +6101,7 @@ impl Governance {
         )?;
 
         // Get the balance of the neuron from the ledger canister.
+        tla_log_locals! { neuron_id: nid.id };
         let balance = self.ledger.account_balance(account).await?;
         let min_stake = self.economics().neuron_minimum_stake_e8s;
         if balance.get_e8s() < min_stake {
diff --git a/rs/nns/governance/src/governance/tla/mod.rs b/rs/nns/governance/src/governance/tla/mod.rs
@@ -28,6 +28,7 @@ mod claim_neuron;
 mod disburse_neuron;
 mod disburse_to_neuron;
 mod merge_neurons;
+mod refresh_neuron;
 mod spawn_neuron;
 mod spawn_neurons;
 mod split_neuron;
@@ -36,6 +37,7 @@ pub use claim_neuron::CLAIM_NEURON_DESC;
 pub use disburse_neuron::DISBURSE_NEURON_DESC;
 pub use disburse_to_neuron::DISBURSE_TO_NEURON_DESC;
 pub use merge_neurons::MERGE_NEURONS_DESC;
+pub use refresh_neuron::REFRESH_NEURON_DESC;
 pub use spawn_neuron::SPAWN_NEURON_DESC;
 pub use spawn_neurons::SPAWN_NEURONS_DESC;
 pub use split_neuron::SPLIT_NEURON_DESC;
diff --git a/rs/nns/governance/src/governance/tla/refresh_neuron.rs b/rs/nns/governance/src/governance/tla/refresh_neuron.rs
@@ -0,0 +1,27 @@
+use lazy_static::lazy_static;
+use tla_instrumentation::{Label, TlaConstantAssignment, ToTla, Update, VarAssignment};
+
+use super::{extract_common_constants, post_process_trace};
+
+lazy_static! {
+    pub static ref REFRESH_NEURON_DESC: Update = {
+        const PID: &str = "Refresh_Neuron";
+        let default_locals = VarAssignment::new().add("neuron_id", 0_u64.to_tla_value());
+
+        Update {
+            default_start_locals: default_locals.clone(),
+            default_end_locals: default_locals,
+            start_label: Label::new("RefreshNeuron1"),
+            end_label: Label::new("Done"),
+            process_id: PID.to_string(),
+            canister_name: "governance".to_string(),
+            post_process: |trace| {
+                let constants = TlaConstantAssignment {
+                    constants: extract_common_constants(PID, trace).into_iter().collect(),
+                };
+                post_process_trace(trace);
+                constants
+            },
+        }
+    };
+}
diff --git a/rs/nns/governance/tests/governance.rs b/rs/nns/governance/tests/governance.rs
@@ -5039,13 +5039,15 @@ fn refresh_neuron_by_memo(owner: PrincipalId, caller: PrincipalId) {
 
 /// Tests that a neuron can be refreshed by memo by it's controller.
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_refresh_neuron_by_memo_by_controller() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     refresh_neuron_by_memo(owner, owner);
 }
 
 /// Tests that a neuron can be refreshed by memo by proxy.
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_refresh_neuron_by_memo_by_proxy() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let caller = *TEST_NEURON_2_OWNER_PRINCIPAL;
@@ -5121,12 +5123,14 @@ fn refresh_neuron_by_id_or_subaccount(
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_refresh_neuron_by_id_by_controller() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     refresh_neuron_by_id_or_subaccount(owner, owner, RefreshBy::NeuronId);
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_refresh_neuron_by_id_by_proxy() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let caller = *TEST_NEURON_1_OWNER_PRINCIPAL;
@@ -5136,12 +5140,14 @@ fn test_refresh_neuron_by_id_by_proxy() {
 /// Tests that a neuron can be refreshed by subaccount, and that anyone can do
 /// it.
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_refresh_neuron_by_subaccount_by_controller() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     refresh_neuron_by_id_or_subaccount(owner, owner, RefreshBy::Subaccount);
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_refresh_neuron_by_subaccount_by_proxy() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let caller = *TEST_NEURON_1_OWNER_PRINCIPAL;
diff --git a/rs/nns/governance/tla/Refresh_Neuron.tla b/rs/nns/governance/tla/Refresh_Neuron.tla
@@ -0,0 +1,155 @@
+------------ MODULE Refresh_Neuron ------------
+EXTENDS TLC, Sequences, Naturals, FiniteSets, Variants
+
+CONSTANT
+    FRESH_NEURON_ID(_)
+
+CONSTANTS
+    Governance_Account_Ids,
+    Neuron_Ids
+
+CONSTANTS
+    Refresh_Neuron_Process_Ids
+
+CONSTANTS
+    \* Minimum stake a neuron can have
+    MIN_STAKE,
+    \* The transfer fee charged by the ledger canister
+    TRANSACTION_FEE
+
+DUMMY_ACCOUNT == ""
+
+\* @type: (a -> b, Set(a)) => a -> b;
+Remove_Arguments(f, S) == [ x \in (DOMAIN f \ S) |-> f[x]]
+
+request(caller, request_args) == [caller |-> caller, method_and_args |-> request_args]
+account_balance(account) == Variant("AccountBalance", [account |-> account])
+
+
+(* --algorithm Governance_Ledger_Refresh_Neuron {
+
+variables
+
+    neuron \in [{} -> {}];
+    \* Used to decide whether we should refresh or claim a neuron
+    neuron_id_by_account \in [{} -> {}];
+    \* The set of currently locked neurons
+    locks = {};
+    \* The queue of messages sent from the governance canister to the ledger canister
+    governance_to_ledger = <<>>;
+    ledger_to_governance = {};
+    spawning_neurons = FALSE;
+
+macro refresh_neuron_reset_local_vars() {
+    neuron_id := 0;
+}
+
+
+\* Modified from formal-models/tla/governance-ledger
+\* A Refresh_Neuron process simulates a call to refresh_neuron
+process ( Refresh_Neuron \in Refresh_Neuron_Process_Ids )
+    variable
+        \* There are two ways that the user can invoke a neuron refresh:
+        \* 1. by specifying an account ID
+        \* 2. by specifying an existing neuron ID
+        \* We only model the second option; the second should follow from the invariant that
+        \* \A nid aid : neuron_id_by_account[aid] = nid <=> neuron[nid].account = aid
+
+        \* The neuron_id is an argument; we let it be chosen non-deteministically
+        neuron_id = 0;
+    {
+    RefreshNeuron1:
+        either {
+            \* Simulate calls that just fail early and don't change the state.
+            \* Not so useful for model checking, but needed to follow the code traces.
+            goto Done;
+        } or {
+        with(nid \in DOMAIN(neuron) \ locks) {
+            neuron_id := nid;
+            locks := locks \union {neuron_id};
+            governance_to_ledger := Append(governance_to_ledger, request(self, account_balance(neuron[nid].account)));
+        };
+        };
+    WaitForBalanceQuery:
+        \* Note that the "with" construct implicitly awaits until the set of values to draw from is non-empty
+        with(answer \in { resp \in ledger_to_governance : resp.caller = self }) {
+            ledger_to_governance := ledger_to_governance \ {answer};
+            if(answer.response /= Variant("Fail", UNIT)) {
+                with (b = VariantGetOrElse("BalanceQueryOk", answer.response, 0)) {
+                    if(b >= MIN_STAKE) {
+                        neuron := [neuron EXCEPT ![neuron_id].cached_stake = b ]
+                    };
+                };
+            };
+            locks := locks \ {neuron_id};
+        };
+        refresh_neuron_reset_local_vars();
+    };
+
+}
+*)
+\* BEGIN TRANSLATION (chksum(pcal) = "e3951dde" /\ chksum(tla) = "3003d9f1")
+VARIABLES neuron, neuron_id_by_account, locks, governance_to_ledger, 
+          ledger_to_governance, spawning_neurons, pc, neuron_id
+
+vars == << neuron, neuron_id_by_account, locks, governance_to_ledger, 
+           ledger_to_governance, spawning_neurons, pc, neuron_id >>
+
+ProcSet == (Refresh_Neuron_Process_Ids)
+
+Init == (* Global variables *)
+        /\ neuron \in [{} -> {}]
+        /\ neuron_id_by_account \in [{} -> {}]
+        /\ locks = {}
+        /\ governance_to_ledger = <<>>
+        /\ ledger_to_governance = {}
+        /\ spawning_neurons = FALSE
+        (* Process Refresh_Neuron *)
+        /\ neuron_id = [self \in Refresh_Neuron_Process_Ids |-> 0]
+        /\ pc = [self \in ProcSet |-> "RefreshNeuron1"]
+
+RefreshNeuron1(self) == /\ pc[self] = "RefreshNeuron1"
+                        /\ \/ /\ pc' = [pc EXCEPT ![self] = "Done"]
+                              /\ UNCHANGED <<locks, governance_to_ledger, neuron_id>>
+                           \/ /\ \E nid \in DOMAIN(neuron) \ locks:
+                                   /\ neuron_id' = [neuron_id EXCEPT ![self] = nid]
+                                   /\ locks' = (locks \union {neuron_id'[self]})
+                                   /\ governance_to_ledger' = Append(governance_to_ledger, request(self, account_balance(neuron[nid].account)))
+                              /\ pc' = [pc EXCEPT ![self] = "WaitForBalanceQuery"]
+                        /\ UNCHANGED << neuron, neuron_id_by_account, 
+                                        ledger_to_governance, spawning_neurons >>
+
+WaitForBalanceQuery(self) == /\ pc[self] = "WaitForBalanceQuery"
+                             /\ \E answer \in { resp \in ledger_to_governance : resp.caller = self }:
+                                  /\ ledger_to_governance' = ledger_to_governance \ {answer}
+                                  /\ IF answer.response /= Variant("Fail", UNIT)
+                                        THEN /\ LET b == VariantGetOrElse("BalanceQueryOk", answer.response, 0) IN
+                                                  IF b >= MIN_STAKE
+                                                     THEN /\ neuron' = [neuron EXCEPT ![neuron_id[self]].cached_stake = b ]
+                                                     ELSE /\ TRUE
+                                                          /\ UNCHANGED neuron
+                                        ELSE /\ TRUE
+                                             /\ UNCHANGED neuron
+                                  /\ locks' = locks \ {neuron_id[self]}
+                             /\ neuron_id' = [neuron_id EXCEPT ![self] = 0]
+                             /\ pc' = [pc EXCEPT ![self] = "Done"]
+                             /\ UNCHANGED << neuron_id_by_account, 
+                                             governance_to_ledger, 
+                                             spawning_neurons >>
+
+Refresh_Neuron(self) == RefreshNeuron1(self) \/ WaitForBalanceQuery(self)
+
+(* Allow infinite stuttering to prevent deadlock on termination. *)
+Terminating == /\ \A self \in ProcSet: pc[self] = "Done"
+               /\ UNCHANGED vars
+
+Next == (\E self \in Refresh_Neuron_Process_Ids: Refresh_Neuron(self))
+           \/ Terminating
+
+Spec == Init /\ [][Next]_vars
+
+Termination == <>(\A self \in ProcSet: pc[self] = "Done")
+
+\* END TRANSLATION
+
+====
diff --git a/rs/nns/governance/tla/Refresh_Neuron_Apalache.tla b/rs/nns/governance/tla/Refresh_Neuron_Apalache.tla
@@ -0,0 +1,67 @@
+---- MODULE Refresh_Neuron_Apalache ----
+
+
+EXTENDS TLC, Variants
+
+(*
+@typeAlias: proc = Str;
+@typeAlias: account = Str;
+@typeAlias: neuronId = Int;
+@typeAlias: methodCall = Transfer({ from: $account, to: $account, amount: Int, fee: Int}) | AccountBalance({ account: $account });
+@typeAlias: methodResponse = Fail(UNIT) | TransferOk(UNIT) | BalanceQueryOk(Int);
+*)
+_type_alias_dummy == TRUE
+
+\* CODE_LINK_INSERT_CONSTANTS
+
+(*
+CONSTANTS
+    \* @type: Set($account);
+    Account_Ids,
+    \* @type: Set($account);
+    Governance_Account_Ids,
+    \* @type: Set($neuronId);
+    Neuron_Ids
+
+CONSTANTS
+    \* @type: Set($proc);
+    Refresh_Neuron_Process_Ids
+
+CONSTANTS
+    \* Minimum stake a neuron can have
+    \* @type: Int;
+    MIN_STAKE,
+    \* The transfer fee charged by the ledger canister
+    \* @type: Int;
+    TRANSACTION_FEE
+*)
+
+VARIABLES
+    \* @type: $neuronId -> {cached_stake: Int, account: $account, maturity: Int, fees: Int};
+    neuron,
+    \* @type: $account -> $neuronId;
+    neuron_id_by_account,
+    \* @type: Set($neuronId);
+    locks,
+    \* @type: Seq({caller : $proc, method_and_args: $methodCall });
+    governance_to_ledger,
+    \* @type: Set({caller: $proc, response: $methodResponse });
+    ledger_to_governance,
+    \* @type: $proc -> Str;
+    pc,
+    \* @type: $proc -> Int;
+    neuron_id,
+    \* Not used by this model, but it's a global variable used by spawn_neurons, so
+    \* it's the easiest to just add it to all the other models
+    \* @type: Bool;
+    spawning_neurons
+
+\* Not used in this model. Consider removing (TODO).
+\* @type: Set($neuronId) => $neuronId;
+FRESH_NEURON_ID(existing_neurons) == CHOOSE nid \in (Neuron_Ids \ existing_neurons): TRUE
+
+MOD == INSTANCE Refresh_Neuron
+
+Next == [MOD!Next]_MOD!vars
+
+====
