From bee59dcd71a64cdad62261a49fde847c286da7ac Mon Sep 17 00:00:00 2001
From: Yin Guanhao <sopium@mysterious.site>
Date: Wed, 23 Nov 2022 15:32:38 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20use=20mem=20pool=20state=20for=20?=
 =?UTF-8?q?=E2=80=9Cget=20block=E2=80=9D=20RPCs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 crates/rpc-server/src/registry.rs  | 33 ++++++++++++++++++++++--------
 crates/store/src/mem_pool_state.rs | 15 +++++++++++++-
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/crates/rpc-server/src/registry.rs b/crates/rpc-server/src/registry.rs
index 8c189af85..d7ec04b7c 100644
--- a/crates/rpc-server/src/registry.rs
+++ b/crates/rpc-server/src/registry.rs
@@ -912,17 +912,29 @@ async fn get_block(
     }))
 }
 
+// Why do we read from `MemPoolState` instead of `Store` for these “get block”
+// RPCs:
+//
+// `MemPoolState` can fall behind `Store` (at the moment after a new block is
+// inserted but mem pool state hasn't been updated). If we read from `Store`, we
+// may get a block that is not in `MemPoolState` yet. If we then try to get
+// scripts for accounts in the block, we may get an error response, because the
+// `get_script` / `get_script_hash` RPCs use `MemPoolState`.
+//
+// Instead if we always read from `MemPoolState`, it is much less likely that we
+// get an error response when getting scripts for accounts in the new block.
+
 async fn get_block_by_number(
     Params((block_number,)): Params<(gw_jsonrpc_types::ckb_jsonrpc_types::Uint64,)>,
-    store: Data<Store>,
+    mem_pool_state: Data<Arc<MemPoolState>>,
 ) -> Result<Option<L2BlockView>> {
     let block_number = block_number.value();
-    let snap = store.get_snapshot();
-    let block_hash = match snap.get_block_hash_by_number(block_number)? {
+    let mem_store = mem_pool_state.load_mem_store();
+    let block_hash = match mem_store.get_block_hash_by_number(block_number)? {
         Some(hash) => hash,
         None => return Ok(None),
     };
-    let block_opt = snap.get_block(&block_hash)?.map(|block| {
+    let block_opt = mem_store.get_block(&block_hash)?.map(|block| {
         let block_view: L2BlockView = block.into();
         block_view
     });
@@ -931,16 +943,19 @@ async fn get_block_by_number(
 
 async fn get_block_hash(
     Params((block_number,)): Params<(gw_jsonrpc_types::ckb_jsonrpc_types::Uint64,)>,
-    store: Data<Store>,
+    mem_pool_state: Data<Arc<MemPoolState>>,
 ) -> Result<Option<JsonH256>> {
     let block_number = block_number.value();
-    let db = store.get_snapshot();
-    let hash_opt = db.get_block_hash_by_number(block_number)?.map(to_jsonh256);
+    let mem_store = mem_pool_state.load_mem_store();
+    let hash_opt = mem_store
+        .get_block_hash_by_number(block_number)?
+        .map(to_jsonh256);
     Ok(hash_opt)
 }
 
-async fn get_tip_block_hash(store: Data<Store>) -> Result<JsonH256> {
-    let tip_block_hash = store.get_last_valid_tip_block_hash()?;
+async fn get_tip_block_hash(mem_pool_state: Data<Arc<MemPoolState>>) -> Result<JsonH256> {
+    let mem_store = mem_pool_state.load_mem_store();
+    let tip_block_hash = mem_store.get_last_valid_tip_block_hash()?;
     Ok(to_jsonh256(tip_block_hash))
 }
 
diff --git a/crates/store/src/mem_pool_state.rs b/crates/store/src/mem_pool_state.rs
index 08006c81d..59ee7edd6 100644
--- a/crates/store/src/mem_pool_state.rs
+++ b/crates/store/src/mem_pool_state.rs
@@ -6,7 +6,10 @@ use std::sync::{
 use arc_swap::ArcSwap;
 use gw_types::packed::{self, BlockInfo};
 
-use crate::state::MemStateDB;
+use crate::{
+    snapshot::StoreSnapshot,
+    state::{overlay::mem_store::MemStore, MemStateDB},
+};
 
 pub const META_MEM_BLOCK_INFO: &[u8] = b"MEM_BLOCK_INFO";
 /// account SMT root
@@ -58,6 +61,16 @@ impl MemPoolState {
         self.inner.load().mem_block.clone()
     }
 
+    pub fn load_mem_store(&self) -> MemStore<StoreSnapshot> {
+        self.inner
+            .load()
+            .state_db
+            .inner_smt_tree()
+            .store()
+            .inner_store()
+            .clone()
+    }
+
     /// Load shared
     pub fn load_shared(&self) -> Shared {
         Shared::clone(&self.inner.load())