Simplify State lifecycle

rs/backend/src/accounts_store.rs

@@ -1,6 +1,7 @@
 //! User accounts and transactions.
+use crate::accounts_store::schema::accounts_in_unbounded_stable_btree_map::AccountsDbAsUnboundedStableBTreeMap;
 use crate::multi_part_transactions_processor::MultiPartTransactionsProcessor;
-use crate::state::StableState;
+use crate::state::{partitions::PartitionType, with_partitions, StableState};
 use crate::stats::Stats;
 use candid::CandidType;
 use dfn_candid::Candid;
@@ -17,7 +18,6 @@ use std::collections::{BTreeMap, HashMap, VecDeque};
 use std::fmt;
 use std::ops::RangeBounds;
 
-pub mod constructors;
 pub mod histogram;
 pub mod schema;
 use schema::{
@@ -45,7 +45,6 @@ const MAX_IMPORTED_TOKENS: i32 = 20;
 pub struct AccountsStore {
     // TODO(NNS1-720): Use AccountIdentifier directly as the key for this HashMap
     accounts_db: schema::proxy::AccountsDbAsProxy,
-
     accounts_db_stats: AccountsDbStats,
 }
 
@@ -92,7 +91,7 @@ impl AccountsDbTrait for AccountsStore {
     }
 }
 
-#[derive(Default, CandidType, Deserialize, Debug, Eq, PartialEq)]
+#[derive(Default, CandidType, Deserialize, Debug, Eq, PartialEq, Clone)]
 pub struct AccountsDbStats {
     pub sub_accounts_count: u64,
     pub hardware_wallet_accounts_count: u64,
@@ -318,6 +317,34 @@ pub enum DetachCanisterResponse {
 }
 
 impl AccountsStore {
+    /// Creates a new `AccountsStore`. Should be called in `init`.
+    pub fn new() -> Self {
+        let accounts_partition = with_partitions(|partitions| partitions.get(PartitionType::Accounts.memory_id()));
+        let accounts_db = AccountsDbAsProxy::from(AccountsDb::UnboundedStableBTreeMap(
+            AccountsDbAsUnboundedStableBTreeMap::new(accounts_partition),
+        ));
+        let accounts_db_stats = AccountsDbStats::default();
+
+        Self {
+            accounts_db,
+            accounts_db_stats,
+        }
+    }
+
+    /// Recovers the state from stable memory. Should be called in `post_upgrade`.
+    pub fn new_restored(serializable_state: AccountsStoreSerializableState) -> Self {
+        let AccountsStoreSerializableState { accounts_db_stats } = serializable_state;
+        let accounts_partition = with_partitions(|partitions| partitions.get(PartitionType::Accounts.memory_id()));
+        let accounts_db = AccountsDbAsProxy::from(AccountsDb::UnboundedStableBTreeMap(
+            AccountsDbAsUnboundedStableBTreeMap::load(accounts_partition),
+        ));
+
+        Self {
+            accounts_db,
+            accounts_db_stats,
+        }
+    }
+
     /// Determines whether a migration is being performed.
     #[must_use]
     #[allow(dead_code)]
@@ -699,13 +726,63 @@ impl AccountsStore {
     }
 }
 
-impl StableState for AccountsStore {
+/// The component(s) of the `AccountsStore` that are serialized and de-serialized during upgrades.
+pub struct AccountsStoreSerializableState {
+    accounts_db_stats: AccountsDbStats,
+}
+
+impl From<AccountsStore> for AccountsStoreSerializableState {
+    fn from(accounts_store: AccountsStore) -> Self {
+        let AccountsStore {
+            // Field(s) persisted by serialization/deserialization through upgrades:
+            accounts_db_stats,
+
+            // Field(s) not persisted by serialization/deserialization through upgrades:
+            accounts_db: _,
+        } = accounts_store;
+
+        Self { accounts_db_stats }
+    }
+}
+
+#[allow(clippy::zero_sized_map_values)]
+type SerializableStateForEncoding = (
+    BTreeMap<Vec<u8>, candid::Empty>,
+    HashMap<AccountIdentifier, AccountWrapper>,
+    HashMap<(AccountIdentifier, AccountIdentifier), candid::Empty>,
+    VecDeque<candid::Empty>,
+    HashMap<AccountIdentifier, candid::Empty>,
+    Option<BlockIndex>,
+    MultiPartTransactionsProcessor,
+    u64,
+    u64,
+    Option<AccountsDbStats>,
+);
+type SerializableStateForDecoding = (
+    candid::Reserved,
+    // TODO: Change to candid:Reserved and remove AccountWrapper after
+    // we've deployed to mainnet. If we do it now, decoding will break
+    // because of the skip quota. The decoder will think there is an
+    // attack with a lot of unnecessary data in a request.
+    HashMap<AccountIdentifier, AccountWrapper>,
+    candid::Reserved,
+    candid::Reserved,
+    candid::Reserved,
+    candid::Reserved,
+    candid::Reserved,
+    candid::Reserved,
+    candid::Reserved,
+    Option<AccountsDbStats>,
+);
+
+impl StableState for AccountsStoreSerializableState {
     fn encode(&self) -> Vec<u8> {
-        // Accounts are now in stable structures and no longer in a simple map
-        // on the heap. So we don't need to encode them here.
-        let empty_accounts = BTreeMap::<Vec<u8>, candid::Empty>::new();
-        Candid((
-            empty_accounts,
+        let AccountsStoreSerializableState { accounts_db_stats } = self;
+
+        let candid_object: SerializableStateForEncoding = (
+            // Accounts are now in stable structures and no longer in a simple map
+            // on the heap. So we don't need to encode them here.
+            BTreeMap::<Vec<u8>, candid::Empty>::new(),
             // hardware_wallets_and_sub_accounts is unused but we need to encode
             // it for backwards compatibility.
             // TODO: Change AccountWrapper to candid::Empty after we've
@@ -734,10 +811,9 @@ impl StableState for AccountsStore {
             // neurons_topped_up_count is unused but we need to encode
             // it for backwards compatibility.
             0u64,
-            Some(&self.accounts_db_stats),
-        ))
-        .into_bytes()
-        .unwrap()
+            Some(accounts_db_stats.clone()),
+        );
+        Candid(candid_object).into_bytes().unwrap()
     }
 
     fn decode(bytes: Vec<u8>) -> Result<Self, String> {
@@ -757,34 +833,13 @@ impl StableState for AccountsStore {
             _last_ledger_sync_timestamp_nanos,
             _neurons_topped_up_count,
             accounts_db_stats_maybe,
-        ): (
-            candid::Reserved,
-            // TODO: Change to candid:Reserved and remove AccountWrapper after
-            // we've deployed to mainnet. If we do it now, decoding will break
-            // because of the skip quota. The decoder will think there is an
-            // attack with a lot of unnecessary data in a request.
-            HashMap<AccountIdentifier, AccountWrapper>,
-            candid::Reserved,
-            candid::Reserved,
-            candid::Reserved,
-            candid::Reserved,
-            candid::Reserved,
-            candid::Reserved,
-            candid::Reserved,
-            Option<AccountsDbStats>,
-        ) = Candid::from_bytes(bytes).map(|c| c.0)?;
+        ): SerializableStateForDecoding = Candid::from_bytes(bytes).map(|c| c.0)?;
 
         let Some(accounts_db_stats) = accounts_db_stats_maybe else {
             return Err("Accounts DB stats should be present since the stable structures migration.".to_string());
         };
 
-        Ok(AccountsStore {
-            // Because the stable structures migration is finished, accounts_db
-            // will be replaced with an AccountsDbAsUnboundedStableBTreeMap in
-            // State::from(Partitions) so it doesn't matter what we set here.
-            accounts_db: AccountsDbAsProxy::default(),
-            accounts_db_stats,
-        })
+        Ok(Self { accounts_db_stats })
     }
 }
 

rs/backend/src/assets.rs

@@ -65,6 +65,7 @@ impl ContentEncoding {
 const LABEL_ASSETS: &[u8] = b"http_assets";
 
 #[derive(Default, Debug, Eq, PartialEq)]
+#[cfg_attr(test, derive(Clone))]
 pub struct AssetHashes(RbTree<Vec<u8>, Hash>);
 
 impl From<&Assets> for AssetHashes {
@@ -115,6 +116,7 @@ impl Asset {
 }
 
 #[derive(Default, CandidType, Deserialize, PartialEq, Eq, Debug)]
+#[cfg_attr(test, derive(Clone))]
 pub struct Assets(HashMap<String, Asset>);
 
 impl Assets {