[Draft] Memtrie integration

chain/chain/src/chain.rs

@@ -4387,6 +4387,44 @@ pub fn collect_receipts_from_response(
     )
 }
 
+pub fn do_load_memtrie(
+    runtime_adapter: Arc<dyn RuntimeAdapter>,
+    shard_uid: ShardUId,
+) -> Result<(), Error> {
+    runtime_adapter.unload_trie_on_catchup(&shard_uid);
+    runtime_adapter
+        .load_mem_trie_on_catchup(&shard_uid)
+        .map_err(|err| near_chain_primitives::Error::StorageError(err))
+}
+
+#[derive(actix::Message)]
+#[rtype(result = "()")]
+pub struct LoadMemtrieRequest {
+    pub runtime_adapter: Arc<dyn RuntimeAdapter>,
+    pub shard_uid: ShardUId,
+    pub sync_hash: CryptoHash,
+}
+
+// Skip `runtime_adapter`, because it's a complex object that has complex logic
+// and many fields.
+impl Debug for LoadMemtrieRequest {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LoadMemtrieRequest")
+            .field("runtime_adapter", &"<not shown>")
+            .field("shard_uid", &self.shard_uid)
+            .field("sync_hash", &self.sync_hash)
+            .finish()
+    }
+}
+
+#[derive(actix::Message, Debug)]
+#[rtype(result = "()")]
+pub struct LoadMemtrieResponse {
+    pub load_memtrie_result: Result<(), near_chain_primitives::error::Error>,
+    pub shard_id: ShardId,
+    pub sync_hash: CryptoHash,
+}
+
 #[derive(actix::Message)]
 #[rtype(result = "()")]
 pub struct ApplyStatePartsRequest {

chain/chain/src/runtime/mod.rs

@@ -1226,6 +1226,14 @@ impl RuntimeAdapter for NightshadeRuntime {
     fn load_mem_tries_on_startup(&self, shard_uids: &[ShardUId]) -> Result<(), StorageError> {
         self.tries.load_mem_tries_for_enabled_shards(shard_uids)
     }
+
+    fn load_mem_trie_on_catchup(&self, shard_uid: &ShardUId) -> Result<(), StorageError> {
+        self.tries.load_mem_trie_for_shard(shard_uid)
+    }
+
+    fn unload_trie_on_catchup(&self, shard_uid: &ShardUId) {
+        self.tries.delete_trie_for_shard(*shard_uid, &mut self.tries.store_update());
+    }
 }
 
 impl node_runtime::adapter::ViewRuntimeAdapter for NightshadeRuntime {

chain/chain/src/test_utils/kv_runtime.rs

@@ -1456,4 +1456,10 @@ impl RuntimeAdapter for KeyValueRuntime {
     fn load_mem_tries_on_startup(&self, _shard_uids: &[ShardUId]) -> Result<(), StorageError> {
         Ok(())
     }
+
+    fn load_mem_trie_on_catchup(&self, _shard_uid: &ShardUId) -> Result<(), StorageError> {
+        Ok(())
+    }
+
+    fn unload_trie_on_catchup(&self, _shard_uid: &ShardUId) {}
 }

chain/chain/src/types.rs

@@ -515,6 +515,12 @@ pub trait RuntimeAdapter: Send + Sync {
     /// but which exact shards to load depends on configuration. This may only be called when flat
     /// storage is ready.
     fn load_mem_tries_on_startup(&self, shard_uids: &[ShardUId]) -> Result<(), StorageError>;
+
+    /// Loads in-memory trie upon catchup. This may only be called when flat storage is ready.
+    fn load_mem_trie_on_catchup(&self, shard_uid: &ShardUId) -> Result<(), StorageError>;
+
+    /// Unloads trie upon catchup. Call this before flat storage is cleared.
+    fn unload_trie_on_catchup(&self, shard_uid: &ShardUId);
 }
 
 /// The last known / checked height and time when we have processed it.

chain/client-primitives/src/types.rs

@@ -97,6 +97,8 @@ pub enum ShardSyncStatus {
     StateDownloadComplete,
     ReshardingScheduling,
     ReshardingApplying,
+    MemtrieLoadScheduling,
+    MemtrieLoadApplying,
     StateSyncDone,
 }
 
@@ -110,7 +112,9 @@ impl ShardSyncStatus {
             ShardSyncStatus::StateDownloadComplete => 4,
             ShardSyncStatus::ReshardingScheduling => 5,
             ShardSyncStatus::ReshardingApplying => 6,
-            ShardSyncStatus::StateSyncDone => 7,
+            ShardSyncStatus::MemtrieLoadScheduling => 7,
+            ShardSyncStatus::MemtrieLoadApplying => 8,
+            ShardSyncStatus::StateSyncDone => 9,
         }
     }
 }
@@ -134,6 +138,8 @@ impl ToString for ShardSyncStatus {
             ShardSyncStatus::StateDownloadComplete => "download complete".to_string(),
             ShardSyncStatus::ReshardingScheduling => "resharding scheduling".to_string(),
             ShardSyncStatus::ReshardingApplying => "resharding applying".to_string(),
+            ShardSyncStatus::MemtrieLoadScheduling => "memtrie load scheduling".to_string(),
+            ShardSyncStatus::MemtrieLoadApplying => "memtrie load applying".to_string(),
             ShardSyncStatus::StateSyncDone => "done".to_string(),
         }
     }

chain/client/src/client.rs

@@ -20,9 +20,9 @@ use near_async::futures::{AsyncComputationSpawner, FutureSpawner};
 use near_async::messaging::IntoSender;
 use near_async::messaging::{CanSend, Sender};
 use near_async::time::{Clock, Duration, Instant};
-use near_chain::chain::VerifyBlockHashAndSignatureResult;
 use near_chain::chain::{
     ApplyStatePartsRequest, BlockCatchUpRequest, BlockMissingChunks, BlocksCatchUpState,
+    LoadMemtrieRequest, VerifyBlockHashAndSignatureResult,
 };
 use near_chain::flat_storage_creator::FlatStorageCreator;
 use near_chain::orphan::OrphanMissingChunks;
@@ -2324,6 +2324,7 @@ impl Client {
         state_parts_task_scheduler: &Sender<ApplyStatePartsRequest>,
         block_catch_up_task_scheduler: &Sender<BlockCatchUpRequest>,
         resharding_scheduler: &Sender<ReshardingRequest>,
+        load_memtrie_scheduler: &Sender<LoadMemtrieRequest>,
         apply_chunks_done_callback: DoneApplyChunkCallback,
         state_parts_future_spawner: &dyn FutureSpawner,
     ) -> Result<(), Error> {
@@ -2397,6 +2398,7 @@ impl Client {
                 tracking_shards,
                 state_parts_task_scheduler,
                 resharding_scheduler,
+                load_memtrie_scheduler,
                 state_parts_future_spawner,
                 use_colour,
                 self.runtime_adapter.clone(),

chain/client/src/client_actions.rs

@@ -21,6 +21,7 @@ use near_async::time::{Duration, Instant};
 use near_async::{MultiSend, MultiSendMessage, MultiSenderFrom};
 use near_chain::chain::{
     ApplyStatePartsRequest, ApplyStatePartsResponse, BlockCatchUpRequest, BlockCatchUpResponse,
+    LoadMemtrieRequest, LoadMemtrieResponse,
 };
 use near_chain::resharding::{ReshardingRequest, ReshardingResponse};
 use near_chain::test_utils::format_hash;
@@ -91,6 +92,7 @@ pub struct SyncJobsSenderForClient {
     pub apply_state_parts: Sender<ApplyStatePartsRequest>,
     pub block_catch_up: Sender<BlockCatchUpRequest>,
     pub resharding: Sender<ReshardingRequest>,
+    pub load_memtrie: Sender<LoadMemtrieRequest>,
 }
 
 pub struct ClientActions {
@@ -1395,6 +1397,7 @@ impl ClientActions {
                 &self.sync_jobs_sender.apply_state_parts,
                 &self.sync_jobs_sender.block_catch_up,
                 &self.sync_jobs_sender.resharding,
+                &self.sync_jobs_sender.load_memtrie,
                 self.get_apply_chunks_done_callback(),
                 self.state_parts_future_spawner.as_ref(),
             ) {
@@ -1629,6 +1632,7 @@ impl ClientActions {
                         shards_to_sync,
                         &self.sync_jobs_sender.apply_state_parts,
                         &self.sync_jobs_sender.resharding,
+                        &self.sync_jobs_sender.load_memtrie,
                         self.state_parts_future_spawner.as_ref(),
                         use_colour,
                         self.client.runtime_adapter.clone(),
@@ -1790,6 +1794,21 @@ impl ClientActionHandler<ReshardingResponse> for ClientActions {
     }
 }
 
+impl ClientActionHandler<LoadMemtrieResponse> for ClientActions {
+    type Result = ();
+
+    #[perf]
+    fn handle(&mut self, msg: LoadMemtrieResponse) -> Self::Result {
+        tracing::debug!(target: "client", ?msg);
+        if let Some((sync, _, _)) = self.client.catchup_state_syncs.get_mut(&msg.sync_hash) {
+            // We are doing catchup
+            sync.set_load_memtrie_result(msg.shard_id, msg.load_memtrie_result);
+        } else {
+            self.client.state_sync.set_load_memtrie_result(msg.shard_id, msg.load_memtrie_result);
+        }
+    }
+}
+
 impl ClientActionHandler<ShardsManagerResponse> for ClientActions {
     type Result = ();