Skip to content

Commit

Permalink
Hide get_key_from_shard_uid_and_hash
Browse files Browse the repository at this point in the history
  • Loading branch information
@staffik
staffik committed Oct 18, 2024
1 parent 2a9cc51 commit 0f66de0
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 46 deletions.
68 changes: 40 additions & 28 deletions core/store/src/adapter/trie_store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::io;
use std::num::NonZero;
use std::sync::Arc;

use borsh::BorshDeserialize;
use near_primitives::errors::{MissingTrieValueContext, StorageError};
use near_primitives::hash::CryptoHash;
use near_primitives::shard_layout::{get_block_shard_uid, ShardUId};
Expand All @@ -11,10 +12,6 @@ use crate::{DBCol, KeyForStateChanges, Store, StoreUpdate, TrieChanges, STATE_SN

use super::{StoreAdapter, StoreUpdateAdapter, StoreUpdateHolder};

/// Accesses to the State column should use the mapped ShardUId (either itself or ancestor in the resharding tree),
/// according to the State mapping strategy introduced in Resharding V3.
pub struct MappedShardUId(ShardUId);

#[derive(Clone)]
pub struct TrieStoreAdapter {
store: Store,
Expand All @@ -35,11 +32,32 @@ impl TrieStoreAdapter {
TrieStoreUpdateAdapter { store_update: StoreUpdateHolder::Owned(self.store.store_update()) }
}

/// Constructs db key to be used to access the State column.
/// First, it consults the `StateShardUIdMapping` column to map the `shard_uid` prefix
/// to its ancestor in the resharding tree (according to Resharding V3)
/// or map to itself if the mapping does not exist.
///
/// Please note that the mapped shard uid is read from db each time which may seem slow.
/// In practice the `StateShardUIdMapping` is very small and should always be stored in the RocksDB cache.
/// The deserialization of ShardUId is also very cheap.
fn get_key_from_shard_uid_and_hash(&self, shard_uid: ShardUId, hash: &CryptoHash) -> [u8; 40] {
let mapped_shard_uid = self
.store
.get_ser::<ShardUId>(DBCol::StateShardUIdMapping, &shard_uid.to_bytes())
.expect("get_key_from_shard_uid_and_hash() failed")
.unwrap_or(shard_uid);
let mut key = [0; 40];
key[0..8].copy_from_slice(&mapped_shard_uid.to_bytes());
key[8..].copy_from_slice(hash.as_ref());
key
}

/// Replaces shard_uid prefix with a mapped value according to mapping strategy in Resharding V3.
/// For this, it does extra read from `DBCol::StateShardUIdMapping`.
///
/// For more details, see `get_key_from_shard_uid_and_hash()` docs.
pub fn get(&self, shard_uid: ShardUId, hash: &CryptoHash) -> Result<Arc<[u8]>, StorageError> {
let mapped_shard_uid = get_mapped_shard_uid(&self.store, shard_uid);
let key = get_key_from_shard_uid_and_hash(mapped_shard_uid, hash);
let key = self.get_key_from_shard_uid_and_hash(shard_uid, hash);
let val = self
.store
.get(DBCol::State, key.as_ref())
Expand All @@ -48,6 +66,15 @@ impl TrieStoreAdapter {
Ok(val.into())
}

pub fn get_ser<T: BorshDeserialize>(
&self,
shard_uid: ShardUId,
hash: &CryptoHash,
) -> Result<T, StorageError> {
let bytes = self.get(shard_uid, hash)?;
T::try_from_slice(&bytes).map_err(|e| StorageError::StorageInconsistentState(e.to_string()))
}

pub fn get_state_snapshot_hash(&self) -> Result<CryptoHash, StorageError> {
let val = self
.store
Expand Down Expand Up @@ -91,20 +118,22 @@ impl<'a> TrieStoreUpdateAdapter<'a> {
Self { store_update: StoreUpdateHolder::Reference(store_update) }
}

fn get_key_from_shard_uid_and_hash(&self, shard_uid: ShardUId, hash: &CryptoHash) -> [u8; 40] {
self.store_update.store.trie_store().get_key_from_shard_uid_and_hash(shard_uid, hash)
}

pub fn decrement_refcount_by(
&mut self,
shard_uid: ShardUId,
hash: &CryptoHash,
decrement: NonZero<u32>,
) {
let mapped_shard_uid = get_mapped_shard_uid(&self.store_update().store, shard_uid);
let key = get_key_from_shard_uid_and_hash(mapped_shard_uid, hash);
let key = self.get_key_from_shard_uid_and_hash(shard_uid, hash);
self.store_update.decrement_refcount_by(DBCol::State, key.as_ref(), decrement);
}

pub fn decrement_refcount(&mut self, shard_uid: ShardUId, hash: &CryptoHash) {
let mapped_shard_uid = get_mapped_shard_uid(&self.store_update().store, shard_uid);
let key = get_key_from_shard_uid_and_hash(mapped_shard_uid, hash);
let key = self.get_key_from_shard_uid_and_hash(shard_uid, hash);
self.store_update.decrement_refcount(DBCol::State, key.as_ref());
}

Expand All @@ -115,8 +144,7 @@ impl<'a> TrieStoreUpdateAdapter<'a> {
data: &[u8],
increment: NonZero<u32>,
) {
let mapped_shard_uid = get_mapped_shard_uid(&self.store_update().store, shard_uid);
let key = get_key_from_shard_uid_and_hash(mapped_shard_uid, hash);
let key = self.get_key_from_shard_uid_and_hash(shard_uid, hash);
self.store_update.increment_refcount_by(DBCol::State, key.as_ref(), data, increment);
}

Expand Down Expand Up @@ -166,22 +194,6 @@ impl<'a> TrieStoreUpdateAdapter<'a> {
}
}

/// Reads shard_uid mapping for given shard.
/// If the mapping does not exist, it means that `shard_uid` maps to itself
pub fn get_mapped_shard_uid(store: &Store, shard_uid: ShardUId) -> MappedShardUId {
let mapped_shard_uid = store
.get_ser::<ShardUId>(DBCol::StateShardUIdMapping, &shard_uid.to_bytes())
.expect("get_mapped_shard_uid() failed");
MappedShardUId(mapped_shard_uid.unwrap_or(shard_uid))
}

pub fn get_key_from_shard_uid_and_hash(shard_uid: MappedShardUId, hash: &CryptoHash) -> [u8; 40] {
let mut key = [0; 40];
key[0..8].copy_from_slice(&shard_uid.0.to_bytes());
key[8..].copy_from_slice(hash.as_ref());
key
}

#[cfg(test)]
mod tests {
use assert_matches::assert_matches;
Expand Down
32 changes: 14 additions & 18 deletions nearcore/src/entity_debug.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use near_primitives::views::{
BlockHeaderView, BlockView, ChunkView, ExecutionOutcomeView, ReceiptView, SignedTransactionView,
};
use near_store::adapter::flat_store::encode_flat_state_db_key;
use near_store::adapter::trie_store::{get_key_from_shard_uid_and_hash, get_mapped_shard_uid};
use near_store::adapter::StoreAdapter;
use near_store::db::GENESIS_CONGESTION_INFO_KEY;
use near_store::flat::delta::KeyForFlatStateDelta;
use near_store::flat::{FlatStateChanges, FlatStateDeltaMetadata, FlatStorageStatus};
Expand Down Expand Up @@ -248,11 +248,10 @@ impl EntityDebugHandlerImpl {
Ok(serialize_entity(&ExecutionOutcomeView::from(outcome.outcome)))
}
EntityQuery::RawTrieNodeByHash { trie_node_hash, shard_uid } => {
let mapped_shard_uid = get_mapped_shard_uid(&store, shard_uid);
let key = get_key_from_shard_uid_and_hash(mapped_shard_uid, &trie_node_hash);
let node = store
.get_ser::<RawTrieNodeWithSize>(DBCol::State, &key)?
.ok_or_else(|| anyhow!("Trie node not found"))?;
.trie_store()
.get_ser::<RawTrieNodeWithSize>(shard_uid, &trie_node_hash)
.map_err(|e| anyhow!("Trie node not found: {e}"))?;
Ok(serialize_raw_trie_node(node))
}
EntityQuery::RawTrieRootByChunkHash { chunk_hash } => {
Expand All @@ -268,21 +267,18 @@ impl EntityDebugHandlerImpl {
.shard_uids()
.nth(shard_index)
.ok_or_else(|| anyhow!("Shard {} not found", chunk.shard_id()))?;
let mapped_shard_uid = get_mapped_shard_uid(&store, shard_uid);
let key =
get_key_from_shard_uid_and_hash(mapped_shard_uid, &chunk.prev_state_root());
let node = store
.get_ser::<RawTrieNodeWithSize>(DBCol::State, &key)?
.ok_or_else(|| anyhow!("State root not found"))?;
.trie_store()
.get_ser::<RawTrieNodeWithSize>(shard_uid, &chunk.prev_state_root())
.map_err(|e| anyhow!("State root not found: {e}"))?;
Ok(serialize_raw_trie_node(node))
}
EntityQuery::RawTrieValueByHash { trie_value_hash, shard_uid } => {
let mapped_shard_uid = get_mapped_shard_uid(&store, shard_uid);
let key = get_key_from_shard_uid_and_hash(mapped_shard_uid, &trie_value_hash);
let value = store
.get(DBCol::State, &key)?
.ok_or_else(|| anyhow!("Trie value not found"))?;
Ok(serialize_entity(&hex::encode(value.as_slice())))
.trie_store()
.get(shard_uid, &trie_value_hash)
.map_err(|e| anyhow!("Trie value not found: {e}"))?;
Ok(serialize_entity(&hex::encode(value)))
}
EntityQuery::ReceiptById { receipt_id } => {
let receipt = store
Expand Down Expand Up @@ -457,11 +453,11 @@ impl EntityDebugHandlerImpl {
state: FlatStateValue,
shard_uid: ShardUId,
) -> anyhow::Result<Vec<u8>> {
let mapped_shard_uid = get_mapped_shard_uid(&store, shard_uid);
Ok(match state {
FlatStateValue::Ref(value) => store
.get(DBCol::State, &get_key_from_shard_uid_and_hash(mapped_shard_uid, &value.hash))?
.ok_or_else(|| anyhow!("ValueRef could not be dereferenced"))?
.trie_store()
.get(shard_uid, &value.hash)
.map_err(|e| anyhow!("ValueRef could not be dereferenced: {e}"))?
.to_vec(),
FlatStateValue::Inlined(data) => data,
})
Expand Down

0 comments on commit 0f66de0

Please sign in to comment.