test(IC-1579): Isolate TLA tests in Rust-based (non-canister) tests (#…

…2241)

Switch to a `LocalKey`-based approach for Rust-based tests, and add an
annotation (proc macro) that scopes the `LocalKey` to the test.

Previously, we used only a global `RwLock`, which would cause traces of
one test to be picked up by other tests. In contract, a `LocalKey` is
task-local, eliminating the interference. We still use the `RwLock` as a
fallback to store traces if the `LocalKey` is not present (i.e., no
scope has been opened). This means that, first, tests that aren't
annotated still collect the traces, but they aren't checked, which is a
bit wasteful but should be OK as these aren't extremely heavy
operations. Second, and more importantly, this means that the
instrumentation will still work in canister-based tests, allowing a
separate query method to pick the traces up from the global variable.

rs/nns/governance/src/governance.rs

@@ -135,7 +135,7 @@ pub mod tla;
 #[cfg(feature = "tla")]
 pub use tla::{
     claim_neuron_desc, split_neuron_desc, tla_update_method, InstrumentationState, ToTla,
-    TLA_INSTRUMENTATION_STATE, TLA_TRACES,
+    TLA_INSTRUMENTATION_STATE, TLA_TRACES_LKEY, TLA_TRACES_MUTEX,
 };
 
 // 70 KB (for executing NNS functions that are not canister upgrades)

rs/nns/governance/src/governance/tla/mod.rs

@@ -21,7 +21,7 @@ mod store;
 
 pub use common::{account_to_tla, opt_subaccount_to_tla, subaccount_to_tla};
 use common::{function_domain_union, governance_account_id};
-pub use store::{TLA_INSTRUMENTATION_STATE, TLA_TRACES};
+pub use store::{TLA_INSTRUMENTATION_STATE, TLA_TRACES_LKEY, TLA_TRACES_MUTEX};
 
 mod split_neuron;
 pub use split_neuron::split_neuron_desc;
@@ -213,10 +213,25 @@ pub fn check_traces() {
     // improving that later, for now we introduce a hard limit on the state size, and
     // skip checking states larger than the limit
     const STATE_SIZE_LIMIT: u64 = 500;
+    fn is_under_limit(p: &ResolvedStatePair) -> bool {
+        p.start.size() < STATE_SIZE_LIMIT && p.end.size() < STATE_SIZE_LIMIT
+    }
+
+    fn print_stats(traces: &Vec<UpdateTrace>) {
+        println!("Checking {} traces with TLA/Apalache", traces.len());
+        for t in traces {
+            let total_len = t.state_pairs.len();
+            let under_limit_len = t.state_pairs.iter().filter(|p| is_under_limit(p)).count();
+            println!(
+                "TLA/Apalache checks: keeping {}/{} states for update {}",
+                under_limit_len, total_len, t.update.process_id
+            );
+        }
+    }
+
     let traces = {
-        // Introduce a scope to drop the write lock immediately, in order
-        // not to poison the lock if we panic later
-        let mut t = TLA_TRACES.write().unwrap();
+        let t = TLA_TRACES_LKEY.get();
+        let mut t = t.borrow_mut();
         std::mem::take(&mut (*t))
     };
 
@@ -230,11 +245,13 @@ pub fn check_traces() {
         panic!("bad apalache bin from 'TLA_APALACHE_BIN': '{:?}'", apalache);
     }
 
+    print_stats(&traces);
+
     let chunk_size = 20;
     let all_pairs = traces.into_iter().flat_map(|t| {
         t.state_pairs
             .into_iter()
-            .filter(|p| p.start.size() < STATE_SIZE_LIMIT && p.end.size() < STATE_SIZE_LIMIT)
+            .filter(is_under_limit)
             .map(move |p| (t.update.clone(), t.constants.clone(), p))
     });
     let chunks = all_pairs.chunks(chunk_size);
@@ -267,7 +284,7 @@ pub fn check_traces() {
                 println!("Possible divergence from the TLA model detected when interacting with the ledger!");
                 println!("If you did not expect to change the interaction between governance and the ledger, reconsider whether your change is safe. You can find additional data on the step that triggered the error below.");
                 println!("If you are confident that your change is correct, please contact the #formal-models Slack channel and describe the problem.");
-                println!("You can edit nervous_system/tla/feature_flags.bzl to disable TLA checks in the CI and get on with your business.");
+                println!("You can edit nns/governance/feature_flags.bzl to disable TLA checks in the CI and get on with your business.");
                 println!("-------------------");
                 println!("Error occured while checking the state pair:\n{:#?}\nwith constants:\n{:#?}", e.pair, e.constants);
                 println!("Apalache returned:\n{:#?}", e.apalache_error);

rs/nns/governance/src/governance/tla/store.rs

@@ -1,4 +1,5 @@
 use local_key::task_local;
+use std::cell::RefCell;
 use std::sync::RwLock;
 pub use tla_instrumentation::{InstrumentationState, UpdateTrace};
 
@@ -8,7 +9,8 @@ pub use tla_instrumentation::{InstrumentationState, UpdateTrace};
 #[cfg(feature = "tla")]
 task_local! {
     pub static TLA_INSTRUMENTATION_STATE: InstrumentationState;
+    pub static TLA_TRACES_LKEY: RefCell<Vec<UpdateTrace>>;
 }
 
 #[cfg(feature = "tla")]
-pub static TLA_TRACES: RwLock<Vec<UpdateTrace>> = RwLock::new(Vec::new());
+pub static TLA_TRACES_MUTEX: RwLock<Vec<UpdateTrace>> = RwLock::new(Vec::new());

rs/nns/governance/tests/governance.rs

@@ -131,7 +131,9 @@ use std::{
 };
 
 #[cfg(feature = "tla")]
-use ic_nns_governance::governance::tla;
+use ic_nns_governance::governance::tla::{check_traces as tla_check_traces, TLA_TRACES_LKEY};
+#[cfg(feature = "tla")]
+use tla_instrumentation_proc_macros::with_tla_trace_check;
 
 /// The 'fake' module is the old scheme for providing NNS test fixtures, aka
 /// the FakeDriver. It is being used here until the older tests have been
@@ -4650,6 +4652,7 @@ fn claim_neuron_by_memo(
 /// Tests that the controller of a neuron (the principal whose hash was used
 /// to build the subaccount) can claim a neuron just with the memo.
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_claim_neuron_by_memo_only() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let memo = 1234u64;
@@ -4669,12 +4672,10 @@ fn test_claim_neuron_by_memo_only() {
     let neuron = gov.neuron_store.with_neuron(&nid, |n| n.clone()).unwrap();
     assert_eq!(neuron.controller(), owner);
     assert_eq!(neuron.cached_neuron_stake_e8s, stake.get_e8s());
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_claim_neuron_without_minimum_stake_fails() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let memo = 1234u64;
@@ -4692,9 +4693,6 @@ fn test_claim_neuron_without_minimum_stake_fails() {
         }
         _ => panic!("Invalid response."),
     };
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 fn do_test_claim_neuron_by_memo_and_controller(owner: PrincipalId, caller: PrincipalId) {
@@ -4731,14 +4729,12 @@ fn do_test_claim_neuron_by_memo_and_controller(owner: PrincipalId, caller: Princ
     let neuron = gov.neuron_store.with_neuron(&nid, |n| n.clone()).unwrap();
     assert_eq!(neuron.controller(), owner);
     assert_eq!(neuron.cached_neuron_stake_e8s, stake.get_e8s());
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 /// Like the above, but explicitly sets the controller in the MemoAndController
 /// struct.
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_claim_neuron_memo_and_controller_by_controller() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     do_test_claim_neuron_by_memo_and_controller(owner, owner);
@@ -4747,6 +4743,7 @@ fn test_claim_neuron_memo_and_controller_by_controller() {
 /// Tests that a non-controller can claim a neuron for the controller (the
 /// principal whose id was used to build the subaccount).
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_claim_neuron_memo_and_controller_by_proxy() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let caller = *TEST_NEURON_2_OWNER_PRINCIPAL;
@@ -4755,6 +4752,7 @@ fn test_claim_neuron_memo_and_controller_by_proxy() {
 
 /// Tests that a non-controller can't claim a neuron for themselves.
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_non_controller_cant_claim_neuron_for_themselves() {
     let owner = *TEST_NEURON_1_OWNER_PRINCIPAL;
     let claimer = *TEST_NEURON_2_OWNER_PRINCIPAL;
@@ -4783,9 +4781,6 @@ fn test_non_controller_cant_claim_neuron_for_themselves() {
         CommandResponse::Error(_) => (),
         _ => panic!("Claim should have failed."),
     };
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 fn refresh_neuron_by_memo(owner: PrincipalId, caller: PrincipalId) {
@@ -4953,6 +4948,7 @@ fn test_refresh_neuron_by_subaccount_by_proxy() {
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_claim_or_refresh_neuron_does_not_overflow() {
     let (mut driver, mut gov, neuron) = create_mature_neuron(true);
     let nid = neuron.id.unwrap();
@@ -5017,9 +5013,6 @@ fn test_claim_or_refresh_neuron_does_not_overflow() {
             .stake_e8s,
         previous_stake_e8s + 100_000_000_000_000
     );
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 #[test]
@@ -5254,6 +5247,7 @@ fn test_cant_disburse_without_paying_fees() {
 /// * the list of all neurons remained unchanged
 /// * the list of accounts is unchanged
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_neuron_split_fails() {
     let from = *TEST_NEURON_1_OWNER_PRINCIPAL;
     // Compute the subaccount to which the transfer would have been made
@@ -5376,12 +5370,10 @@ fn test_neuron_split_fails() {
     assert_eq!(gov.neuron_store.len(), 1);
     //  There is still only one ledger account.
     driver.assert_num_neuron_accounts_exist(1);
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_neuron_split() {
     let from = *TEST_NEURON_1_OWNER_PRINCIPAL;
     // Compute the subaccount to which the transfer would have been made
@@ -5482,12 +5474,10 @@ fn test_neuron_split() {
     let mut expected_neuron_ids = vec![id, child_nid];
     expected_neuron_ids.sort_unstable();
     assert_eq!(neuron_ids, expected_neuron_ids);
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 #[test]
+#[cfg_attr(feature = "tla", with_tla_trace_check)]
 fn test_seed_neuron_split() {
     let from = *TEST_NEURON_1_OWNER_PRINCIPAL;
     // Compute the subaccount to which the transfer would have been made
@@ -5558,9 +5548,6 @@ fn test_seed_neuron_split() {
     assert_eq!(child_neuron.dissolve_state, parent_neuron.dissolve_state);
     assert_eq!(child_neuron.kyc_verified, true);
     assert_eq!(child_neuron.neuron_type, Some(NeuronType::Seed as i32));
-
-    #[cfg(feature = "tla")]
-    tla::check_traces();
 }
 
 // Spawn neurons has the least priority in the periodic tasks, so we need to run

rs/tla_instrumentation/tla_instrumentation/src/lib.rs

@@ -41,7 +41,7 @@ pub struct Update {
     pub post_process: fn(&mut Vec<ResolvedStatePair>) -> TlaConstantAssignment,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct UpdateTrace {
     pub update: Update,
     pub state_pairs: Vec<ResolvedStatePair>,

rs/tla_instrumentation/tla_instrumentation/tests/multiple_calls.rs

@@ -35,9 +35,10 @@ mod tla_stuff {
 
     task_local! {
         pub static TLA_INSTRUMENTATION_STATE: InstrumentationState;
+        pub static TLA_TRACES_LKEY: std::cell::RefCell<Vec<UpdateTrace>>;
     }
 
-    pub static TLA_TRACES: RwLock<Vec<UpdateTrace>> = RwLock::new(Vec::new());
+    pub static TLA_TRACES_MUTEX: RwLock<Vec<UpdateTrace>> = RwLock::new(Vec::new());
 
     pub fn tla_get_globals(c: &StructCanister) -> GlobalState {
         let mut state = GlobalState::new();
@@ -110,7 +111,9 @@ mod tla_stuff {
     }
 }
 
-use tla_stuff::{my_f_desc, CAN_NAME, PID, TLA_INSTRUMENTATION_STATE, TLA_TRACES};
+use tla_stuff::{
+    my_f_desc, CAN_NAME, PID, TLA_INSTRUMENTATION_STATE, TLA_TRACES_LKEY, TLA_TRACES_MUTEX,
+};
 
 struct StructCanister {
     pub counter: u64,
@@ -162,7 +165,7 @@ fn multiple_calls_test() {
         let canister = &mut *addr_of_mut!(GLOBAL);
         tokio_test::block_on(canister.my_method());
     }
-    let trace = &TLA_TRACES.read().unwrap()[0];
+    let trace = &TLA_TRACES_MUTEX.read().unwrap()[0];
     assert_eq!(
         trace.constants.to_map().get("MAX_COUNTER"),
         Some(&3_u64.to_string())

rs/tla_instrumentation/tla_instrumentation/tests/structs.rs

@@ -10,7 +10,7 @@ use tla_instrumentation::{
     tla_value::{TlaValue, ToTla},
     Destination, InstrumentationState,
 };
-use tla_instrumentation_proc_macros::tla_update_method;
+use tla_instrumentation_proc_macros::{tla_update_method, with_tla_trace_check};
 
 mod common;
 use common::check_tla_trace;
@@ -35,9 +35,10 @@ mod tla_stuff {
 
     task_local! {
         pub static TLA_INSTRUMENTATION_STATE: InstrumentationState;
+        pub static TLA_TRACES_LKEY: std::cell::RefCell<Vec<UpdateTrace>>;
     }
 
-    pub static TLA_TRACES: RwLock<Vec<UpdateTrace>> = RwLock::new(Vec::new());
+    pub static TLA_TRACES_MUTEX: RwLock<Vec<UpdateTrace>> = RwLock::new(Vec::new());
 
     pub fn tla_get_globals(c: &StructCanister) -> GlobalState {
         let mut state = GlobalState::new();
@@ -110,7 +111,9 @@ mod tla_stuff {
     }
 }
 
-use tla_stuff::{my_f_desc, CAN_NAME, PID, TLA_INSTRUMENTATION_STATE, TLA_TRACES};
+use tla_stuff::{
+    my_f_desc, CAN_NAME, PID, TLA_INSTRUMENTATION_STATE, TLA_TRACES_LKEY, TLA_TRACES_MUTEX,
+};
 
 struct StructCanister {
     pub counter: u64,
@@ -155,7 +158,7 @@ fn struct_test() {
         let canister = &mut *addr_of_mut!(GLOBAL);
         tokio_test::block_on(canister.my_method());
     }
-    let trace = &TLA_TRACES.read().unwrap()[0];
+    let trace = &TLA_TRACES_MUTEX.read().unwrap()[0];
     assert_eq!(
         trace.constants.to_map().get("MAX_COUNTER"),
         Some(&2_u64.to_string())
@@ -263,3 +266,18 @@ fn struct_test() {
 
     check_tla_trace(trace);
 }
+
+fn tla_check_traces() {
+    let traces = TLA_TRACES_LKEY.get();
+    let traces = traces.borrow();
+    for t in &*traces {
+        check_tla_trace(t)
+    }
+}
+
+#[test]
+#[with_tla_trace_check]
+fn annotated_test() {
+    let canister = &mut StructCanister { counter: 0 };
+    tokio_test::block_on(canister.my_method());
+}