refactor: change OnceLock to RwLock

axonweb3 · Jul 28, 2023 · f49ba86 · f49ba86
1 parent 579a8a4
commit f49ba86
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 36 deletions.
diff --git a/core/executor/src/system_contract/ckb_light_client/store.rs b/core/executor/src/system_contract/ckb_light_client/store.rs
@@ -15,15 +15,18 @@ pub struct CkbLightClientStore {
 
 impl CkbLightClientStore {
     pub fn new(root: H256) -> ProtocolResult<Self> {
-        let trie_db = match HEADER_CELL_DB.get() {
-            Some(db) => db,
-            None => return Err(SystemScriptError::TrieDbNotInit.into()),
+        let trie_db = {
+            let lock = HEADER_CELL_DB.read();
+            match lock.clone() {
+                Some(db) => db,
+                None => return Err(SystemScriptError::TrieDbNotInit.into()),
+            }
         };
 
         let trie = if root == H256::default() {
-            MPTTrie::new(Arc::clone(trie_db))
+            MPTTrie::new(Arc::clone(&trie_db))
         } else {
-            match MPTTrie::from_root(root, Arc::clone(trie_db)) {
+            match MPTTrie::from_root(root, Arc::clone(&trie_db)) {
                 Ok(m) => m,
                 Err(e) => return Err(SystemScriptError::RestoreMpt(e.to_string()).into()),
             }

diff --git a/core/executor/src/system_contract/image_cell/store.rs b/core/executor/src/system_contract/image_cell/store.rs
@@ -25,15 +25,18 @@ pub struct ImageCellStore {
 
 impl ImageCellStore {
     pub fn new(root: H256) -> ProtocolResult<Self> {
-        let trie_db = match HEADER_CELL_DB.get() {
-            Some(db) => db,
-            None => return Err(SystemScriptError::TrieDbNotInit.into()),
+        let trie_db = {
+            let lock = HEADER_CELL_DB.read();
+            match lock.clone() {
+                Some(db) => db,
+                None => return Err(SystemScriptError::TrieDbNotInit.into()),
+            }
         };
 
         let trie = if root == H256::default() {
-            MPTTrie::new(Arc::clone(trie_db))
+            MPTTrie::new(Arc::clone(&trie_db))
         } else {
-            match MPTTrie::from_root(root, Arc::clone(trie_db)) {
+            match MPTTrie::from_root(root, Arc::clone(&trie_db)) {
                 Ok(m) => m,
                 Err(e) => return Err(SystemScriptError::RestoreMpt(e.to_string()).into()),
             }

diff --git a/core/executor/src/system_contract/metadata/store.rs b/core/executor/src/system_contract/metadata/store.rs
@@ -16,20 +16,23 @@ pub struct MetadataStore {
 
 impl MetadataStore {
     pub fn new(root: H256) -> ProtocolResult<Self> {
-        let trie_db = match METADATA_DB.get() {
-            Some(db) => db,
-            None => return Err(SystemScriptError::TrieDbNotInit.into()),
+        let trie_db = {
+            let lock = METADATA_DB.read().clone();
+            match lock {
+                Some(db) => db,
+                None => return Err(SystemScriptError::TrieDbNotInit.into()),
+            }
         };
 
         let trie = if root == H256::default() {
-            let mut m = MPTTrie::new(Arc::clone(trie_db));
+            let mut m = MPTTrie::new(Arc::clone(&trie_db));
             m.insert(
                 EPOCH_SEGMENT_KEY.as_bytes(),
                 &EpochSegment::new().as_bytes(),
             )?;
             m
         } else {
-            match MPTTrie::from_root(root, Arc::clone(trie_db)) {
+            match MPTTrie::from_root(root, Arc::clone(&trie_db)) {
                 Ok(m) => m,
                 Err(e) => return Err(SystemScriptError::RestoreMpt(e.to_string()).into()),
             }

diff --git a/core/executor/src/system_contract/mod.rs b/core/executor/src/system_contract/mod.rs
@@ -19,12 +19,12 @@ pub use crate::system_contract::native_token::{
 
 use std::path::Path;
 use std::sync::Arc;
-use std::sync::OnceLock;
 
 use ckb_traits::{CellDataProvider, HeaderProvider};
 use ckb_types::core::cell::{CellProvider, CellStatus};
 use ckb_types::core::{HeaderBuilder, HeaderView};
 use ckb_types::{packed, prelude::*};
+use parking_lot::RwLock;
 
 use common_config_parser::types::ConfigRocksDB;
 use protocol::types::{Bytes, Hasher, SignedTransaction, TxResp, H160, H256};
@@ -43,14 +43,11 @@ pub const fn system_contract_address(addr: u8) -> H160 {
 const HEADER_CELL_DB_CACHE_SIZE: usize = 200;
 const METADATA_DB_CACHE_SIZE: usize = 10;
 
-/// System contract init section. It needs to initialize two databases, one for
-/// Metadata and one for CkbLightClient&ImageCell
-static HEADER_CELL_DB: OnceLock<Arc<RocksTrieDB>> = OnceLock::new();
-static METADATA_DB: OnceLock<Arc<RocksTrieDB>> = OnceLock::new();
-
 lazy_static::lazy_static! {
     pub static ref HEADER_CELL_ROOT_KEY: H256 = Hasher::digest("header_cell_mpt_root");
     pub static ref METADATA_ROOT_KEY: H256 = Hasher::digest("metadata_root");
+    pub(crate) static ref METADATA_DB: RwLock<Option<Arc<RocksTrieDB>>> = RwLock::new(None);
+    pub(crate) static ref HEADER_CELL_DB: RwLock<Option<Arc<RocksTrieDB>>> = RwLock::new(None);
 }
 
 #[macro_export]
@@ -80,6 +77,20 @@ pub trait SystemContract {
     fn after_block_hook<Adapter: ExecutorAdapter>(&self, _adapter: &mut Adapter) {}
 }
 
+pub fn swap_metadata_db(new_db: Arc<RocksTrieDB>) -> Arc<RocksTrieDB> {
+    METADATA_DB
+        .write()
+        .replace(new_db)
+        .unwrap_or_else(|| panic!("metadata db is not initialized"))
+}
+
+pub fn swap_header_cell_db(new_db: Arc<RocksTrieDB>) -> Arc<RocksTrieDB> {
+    HEADER_CELL_DB
+        .write()
+        .replace(new_db)
+        .unwrap_or_else(|| panic!("header cell db is not initialized"))
+}
+
 pub fn init<P: AsRef<Path>, Adapter: ExecutorAdapter>(
     path: P,
     config: ConfigRocksDB,
@@ -89,25 +100,22 @@ pub fn init<P: AsRef<Path>, Adapter: ExecutorAdapter>(
 
     // Init metadata db.
     let metadata_db_path = path.as_ref().join("metadata");
-    METADATA_DB.get_or_init(|| {
-        Arc::new(
+    {
+        let mut db = METADATA_DB.write();
+        db.replace(Arc::new(
             RocksTrieDB::new(metadata_db_path, config.clone(), METADATA_DB_CACHE_SIZE)
                 .expect("[system contract] metadata new rocksdb error"),
-        )
-    });
+        ));
+    }
 
-    // Init header cell db.
     let header_cell_db_path = path.as_ref().join("header_cell");
-    HEADER_CELL_DB.get_or_init(|| {
-        Arc::new(
-            RocksTrieDB::new(
-                header_cell_db_path,
-                config.clone(),
-                HEADER_CELL_DB_CACHE_SIZE,
-            )
-            .expect("[system contract] header&cell new rocksdb error"),
-        )
-    });
+    {
+        let mut db = HEADER_CELL_DB.write();
+        db.replace(Arc::new(
+            RocksTrieDB::new(header_cell_db_path, config, HEADER_CELL_DB_CACHE_SIZE)
+                .expect("[system contract] header&cell new rocksdb error"),
+        ));
+    }
 
     let current_cell_root = adapter.storage(CkbLightClientContract::ADDRESS, *HEADER_CELL_ROOT_KEY);