feat(sns): Data migration that retrofits Swap.{direct_participation_icp_e8s, neurons_fund_participation_icp_e8s} + golden upgrade tests

rs/sns/integration_tests/BUILD.bazel

@@ -189,6 +189,7 @@ rust_ic_test_suite_with_extra_srcs(
         ["src/*.rs"],
         exclude = [
             "src/lib.rs",
+            "src/golden_state_swap_upgrade_twice.rs",
         ],
     ),
     aliases = ALIASES,
@@ -204,3 +205,57 @@ rust_ic_test_suite_with_extra_srcs(
     tags = ["cpu:8"],
     deps = DEPENDENCIES_WITH_TEST_FEATURES + TEST_DEV_DEPENDENCIES,
 )
+
+# To run this test,
+#
+#     bazel \
+#         test \
+#         --test_env=SSH_AUTH_SOCK \
+#         --test_timeout=43200 \
+#         //rs/sns/integration_tests:golden_state_swap_upgrade_twice
+#
+# The unusual things in this command are:
+#  - `--test_env=SSH_AUTH_SOCK`: This causes the SSH_AUTH_SOCK environment variable to be
+#    "forwarded" from your shell to the sandbox where the test is run. This authorizes the test
+#    to download the test data.
+#  - `--test_timeout=43200`: This sets the test timeout to 12 hours (more than currently required).
+#
+# Additionally, the following flags are recommended (but not required):
+#
+# --test_output=streamed
+# --test_arg=--nocapture
+#
+# These let you watch the progress of the test, rather than only being able to see the output only
+# at the end.
+#
+# See the .bazelrc for more configuration information.
+rust_ic_test_suite_with_extra_srcs(
+    name = "golden_state_swap_upgrade_twice",
+    # This uses on the order of 50 GB of disk space.
+    # Therefore, size = "large" is not large enough.
+    size = "enormous",
+    srcs = [
+        "src/golden_state_swap_upgrade_twice.rs",
+    ],
+    aliases = ALIASES,
+    args = [
+        "--test-threads",
+        "7",
+    ],
+    crate_features = ["test"],
+    data = DATA_DEPS + [
+        "@mainnet_sns-swap-canister//file",
+    ],
+    env = dict(ENV.items() + [
+        ("MAINNET_SNS_SWAP_CANISTER_WASM_PATH", "$(rootpath @mainnet_sns-swap-canister//file)"),
+    ]),
+    extra_srcs = [],
+    proc_macro_deps = MACRO_DEPENDENCIES + MACRO_DEV_DEPENDENCIES,
+    tags = [
+        "cpu:8",
+        "manual",  # TODO: Enable on CI
+        "no-sandbox",  # such that the test can access the file $SSH_AUTH_SOCK.
+        "requires-network",  # Because mainnet state is downloaded (and used).
+    ],
+    deps = DEPENDENCIES_WITH_TEST_FEATURES + TEST_DEV_DEPENDENCIES + ["//rs/nns/test_utils/golden_nns_state"],
+)

rs/sns/integration_tests/Cargo.toml

@@ -38,6 +38,7 @@ ic-registry-subnet-type = { path = "../../registry/subnet_type" }
 ic-sns-governance = { path = "../governance" }
 ic-sns-init = { path = "../init" }
 ic-sns-root = { path = "../root" }
+ic-nns-test-utils-golden-nns-state = { path = "../../nns/test_utils/golden_nns_state" }
 ic-universal-canister = { path = "../../universal_canister/lib" }
 icrc-ledger-types = { path = "../../../packages/icrc-ledger-types" }
 maplit = "1.0.2"

rs/sns/integration_tests/src/golden_state_swap_upgrade_twice.rs

@@ -0,0 +1,180 @@
+use candid::{Decode, Encode};
+use ic_nns_test_utils::sns_wasm::{
+    build_swap_sns_wasm, create_modified_sns_wasm, ensure_sns_wasm_gzipped,
+};
+use ic_sns_swap::pb::v1::{DerivedState, GetStateRequest, GetStateResponse, Swap};
+use ic_sns_wasm::pb::v1::SnsWasm;
+use ic_state_machine_tests::StateMachine;
+use ic_types::{CanisterId, PrincipalId};
+use pretty_assertions::assert_eq;
+use std::str::FromStr;
+
+// TODO[NNS1-3386]: Remove this function once all existing Swaps are upgraded.
+fn redact_unavailable_swap_fields(swap_state: &mut GetStateResponse) {
+    // The following fields were added to the swap state later than some of the Swap canisters'
+    // last upgrade. These fields will become available after those canisters are upgraded.
+    //
+    // Why is it okay to redact these fields in this test? This test accompanies a data migration
+    // that sets these fields for Swaps that don't yet have it in post_upgrade.
+    {
+        let swap = swap_state.swap.clone().unwrap();
+        swap_state.swap = Some(Swap {
+            timers: None,
+            direct_participation_icp_e8s: None,
+            neurons_fund_participation_icp_e8s: None,
+            ..swap
+        });
+    }
+
+    // The following fields were added to the derived state later than some of the Swap canisters'
+    // last upgrade. These fields will become available after those canisters are upgraded.
+    //
+    // Why is it okay to redact these fields in this test? As the name suggests, these fields are
+    // part of Swap's *derived* state, i.e., they are not stored in canister memory but recomputed
+    // upon request. Therefore, the only reason they might not have reasonable values is when
+    // the Swap canister's *persisted* state (`swap_state.swap`) too incomplete to compute them.
+    {
+        let derived = swap_state.derived.unwrap();
+        swap_state.derived = Some(DerivedState {
+            direct_participant_count: None,
+            cf_participant_count: None,
+            cf_neuron_count: None,
+            direct_participation_icp_e8s: None,
+            neurons_fund_participation_icp_e8s: None,
+            ..derived
+        });
+    }
+}
+
+fn get_state(
+    state_machine: &StateMachine,
+    swap_canister_id: CanisterId,
+    sns_name: &str,
+) -> GetStateResponse {
+    let args = Encode!(&GetStateRequest {}).unwrap();
+    let state_before_upgrade = state_machine
+        .execute_ingress(swap_canister_id, "get_state", args)
+        .unwrap_or_else(|err| {
+            panic!(
+                "Unable to get state of {}'s Swap canister: {}",
+                sns_name, err,
+            )
+        });
+    Decode!(&state_before_upgrade.bytes(), GetStateResponse).unwrap()
+}
+
+fn upgrade_swap_to_tip_of_master(
+    state_machine: &StateMachine,
+    swap_canister_id: CanisterId,
+    swap_wasm: SnsWasm,
+    sns_name: &str,
+) {
+    let swap_upgrade_arg = Encode!().unwrap();
+
+    state_machine
+        .upgrade_canister(swap_canister_id, swap_wasm.wasm, swap_upgrade_arg)
+        .unwrap_or_else(|err| panic!("Cannot upgrade {}'s Swap canister: {}", sns_name, err));
+}
+
+/// Returns the pre-upgrade and post-upgrade states of the Swap.
+fn run_upgrade_for_swap(
+    state_machine: &StateMachine,
+    swap_canister_id: CanisterId,
+    swap_wasm: SnsWasm,
+    sns_name: &str,
+) -> (GetStateResponse, GetStateResponse) {
+    let swap_pre_state = get_state(state_machine, swap_canister_id, sns_name);
+
+    upgrade_swap_to_tip_of_master(state_machine, swap_canister_id, swap_wasm, sns_name);
+
+    let swap_post_state = get_state(state_machine, swap_canister_id, sns_name);
+
+    (swap_pre_state, swap_post_state)
+}
+
+fn run_test_for_swap(state_machine: &StateMachine, swap_canister_id: &str, sns_name: &str) {
+    let swap_canister_id =
+        CanisterId::unchecked_from_principal(PrincipalId::from_str(swap_canister_id).unwrap());
+
+    let swap_wasm_1 = ensure_sns_wasm_gzipped(build_swap_sns_wasm());
+    let swap_wasm_2 = create_modified_sns_wasm(&swap_wasm_1, Some(42));
+    assert_ne!(swap_wasm_1, swap_wasm_2);
+
+    // Experiment I: Upgrade from golden version to the tip of this branch.
+    {
+        let (mut swap_pre_state, mut swap_post_state) =
+            run_upgrade_for_swap(state_machine, swap_canister_id, swap_wasm_1, sns_name);
+
+        // Some fields need to be redacted as they were introduced after some Swaps were created.
+        redact_unavailable_swap_fields(&mut swap_post_state);
+
+        // Since some SNSs do have (some of) the new fields, we need to redact the same set of
+        // fields from the pre-state, too.
+        redact_unavailable_swap_fields(&mut swap_pre_state);
+
+        // Otherwise, the states before and after the migration should match.
+        assert_eq!(
+            swap_pre_state, swap_post_state,
+            "Experiment I: Swap state mismatch detected for {} ",
+            sns_name
+        );
+    }
+
+    // Experiment II: Upgrade again to test the pre-upgrade hook.
+    {
+        let (swap_pre_state, swap_post_state) =
+            run_upgrade_for_swap(state_machine, swap_canister_id, swap_wasm_2, sns_name);
+
+        // Nothing to redact in this case; we've just upgraded a recent version to its modified version.
+
+        assert_eq!(
+            swap_pre_state, swap_post_state,
+            "Experiment II: Swap state mismatch detected for {}",
+            sns_name
+        );
+    }
+}
+
+#[test]
+fn golden_state_swap_upgrade_twice() {
+    let snses_under_test = [
+        ("vuqiy-liaaa-aaaaq-aabiq-cai", "BOOM DAO"),
+        ("iuhw5-siaaa-aaaaq-aadoq-cai", "CYCLES-TRANSFER-STATION"),
+        ("uc3qt-6yaaa-aaaaq-aabnq-cai", "Catalyze"),
+        ("n223b-vqaaa-aaaaq-aadsa-cai", "DOGMI"),
+        ("xhply-dqaaa-aaaaq-aabga-cai", "DecideAI DAO"),
+        ("zcdfx-6iaaa-aaaaq-aaagq-cai", "Dragginz"),
+        ("grlys-pqaaa-aaaaq-aacoa-cai", "ELNA AI"),
+        ("bcl3g-3aaaa-aaaaq-aac5a-cai", "EstateDAO"),
+        ("t7z6p-ryaaa-aaaaq-aab7q-cai", "Gold DAO"),
+        ("4f5dx-pyaaa-aaaaq-aaa3q-cai", "ICGhost"),
+        ("habgn-xyaaa-aaaaq-aaclq-cai", "ICLighthouse DAO"),
+        ("lwslc-cyaaa-aaaaq-aadfq-cai", "ICPCC DAO LLC"),
+        ("ch7an-giaaa-aaaaq-aacwq-cai", "ICPSwap"),
+        ("c424i-4qaaa-aaaaq-aacua-cai", "ICPanda DAO"),
+        ("mzwsh-biaaa-aaaaq-aaduq-cai", "ICVC"),
+        ("mlqf6-nyaaa-aaaaq-aadxq-cai", "Juno Build"),
+        ("7sppf-6aaaa-aaaaq-aaata-cai", "Kinic"),
+        ("khyv5-2qaaa-aaaaq-aadaa-cai", "MORA DAO"),
+        ("kv6ce-waaaa-aaaaq-aadda-cai", "Motoko"),
+        ("f25or-jiaaa-aaaaq-aaceq-cai", "Neutrinite"),
+        ("q2nfe-mqaaa-aaaaq-aabua-cai", "Nuance"),
+        ("jxl73-gqaaa-aaaaq-aadia-cai", "ORIGYN"),
+        ("2hx64-daaaa-aaaaq-aaana-cai", "OpenChat"),
+        ("dkred-jaaaa-aaaaq-aacra-cai", "OpenFPL"),
+        ("qils5-aaaaa-aaaaq-aabxa-cai", "SONIC"),
+        ("rmg5p-zaaaa-aaaaq-aabra-cai", "Seers"),
+        ("hshru-3iaaa-aaaaq-aaciq-cai", "Sneed"),
+        ("ezrhx-5qaaa-aaaaq-aacca-cai", "TRAX"),
+        ("ipcky-iqaaa-aaaaq-aadma-cai", "WaterNeuron"),
+        ("6eexo-lqaaa-aaaaq-aaawa-cai", "YRAL"),
+        ("a2cof-vaaaa-aaaaq-aacza-cai", "Yuku DAO"),
+    ];
+
+    let state_machine =
+        ic_nns_test_utils_golden_nns_state::new_state_machine_with_golden_sns_state_or_panic();
+
+    for (swap_canister_id, sns_name) in snses_under_test {
+        run_test_for_swap(&state_machine, swap_canister_id, sns_name);
+    }
+}

rs/sns/integration_tests/src/lib.rs

@@ -57,3 +57,6 @@ mod http_request;
 
 #[cfg(test)]
 mod timers;
+
+#[cfg(test)]
+mod golden_state_swap_upgrade_twice;

rs/sns/swap/canister/canister.rs

@@ -454,6 +454,9 @@ fn canister_post_upgrade() {
     });
 
     init_timers();
+
+    // TODO[NNS1-3386]: Remove once all Swaps are migrated to have these fields populated.
+    swap_mut().migrate_state();
 }
 
 /// Serve an HttpRequest made to this canister

rs/sns/swap/src/swap.rs

@@ -427,8 +427,8 @@ impl Swap {
             purge_old_tickets_next_principal: Some(FIRST_PRINCIPAL_BYTES.to_vec()),
             already_tried_to_auto_finalize: Some(false),
             auto_finalize_swap_response: None,
-            direct_participation_icp_e8s: None,
-            neurons_fund_participation_icp_e8s: None,
+            direct_participation_icp_e8s: Some(0),
+            neurons_fund_participation_icp_e8s: Some(0),
             timers: None,
         };
         if init.validate_swap_init_for_one_proposal_flow().is_ok() {
@@ -1013,6 +1013,32 @@ impl Swap {
     // --- state modifying methods ---------------------------------------------
     //
 
+    // TODO[NNS1-3386]: Remove this function.
+    pub fn migrate_state(&mut self) {
+        if self.direct_participation_icp_e8s.is_none() {
+            let direct_participation_icp_e8s =
+                self.buyers
+                    .values()
+                    .fold(0_u64, |sum_icp_e8s, buyer_state| {
+                        let amount_icp_e8s = buyer_state.amount_icp_e8s();
+                        sum_icp_e8s.saturating_add(amount_icp_e8s)
+                    });
+            self.direct_participation_icp_e8s = Some(direct_participation_icp_e8s);
+        }
+
+        if self.neurons_fund_participation_icp_e8s.is_none() {
+            let neurons_fund_participation_icp_e8s =
+                self.cf_participants
+                    .iter()
+                    .fold(0_u64, |sum_icp_e8s, neurons_fund_participant| {
+                        let participant_total_icp_e8s =
+                            neurons_fund_participant.participant_total_icp_e8s();
+                        sum_icp_e8s.saturating_add(participant_total_icp_e8s)
+                    });
+            self.neurons_fund_participation_icp_e8s = Some(neurons_fund_participation_icp_e8s);
+        }
+    }
+
     /// Runs those tasks that should be run periodically.
     ///
     /// The argument 'now_fn' is a function that returns the current time for bookkeeping