diff --git a/core/executor/src/system_contract/ckb_light_client/store.rs b/core/executor/src/system_contract/ckb_light_client/store.rs index 452664641..3ebff59cd 100644 --- 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 { - 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 index cf728f620..32e995cbc 100644 --- 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 { - 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 index d8b830382..b0980f318 100644 --- 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 { - 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 index 20ada6c11..36331eaaf 100644 --- 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> = OnceLock::new(); -static METADATA_DB: OnceLock> = 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>> = RwLock::new(None); + pub(crate) static ref HEADER_CELL_DB: RwLock>> = RwLock::new(None); } #[macro_export] @@ -80,6 +77,20 @@ pub trait SystemContract { fn after_block_hook(&self, _adapter: &mut Adapter) {} } +pub fn swap_metadata_db(new_db: Arc) -> Arc { + METADATA_DB + .write() + .replace(new_db) + .unwrap_or_else(|| panic!("metadata db is not initialized")) +} + +pub fn swap_header_cell_db(new_db: Arc) -> Arc { + HEADER_CELL_DB + .write() + .replace(new_db) + .unwrap_or_else(|| panic!("header cell db is not initialized")) +} + pub fn init, Adapter: ExecutorAdapter>( path: P, config: ConfigRocksDB, @@ -89,25 +100,22 @@ pub fn init, 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);