Skip to content

Commit

Permalink
refactor: Remove FlatStateColumn (#9051)
Browse files Browse the repository at this point in the history 
Addresses issue #9044
  • Loading branch information
Jure Bajic authored May 15, 2023
1 parent 0511999 commit 48087bd
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 48 deletions.
7 changes: 3 additions & 4 deletions core/store/src/flat/storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,8 @@ use near_primitives::shard_layout::{ShardLayout, ShardUId};
use tracing::{debug, info, warn};

use crate::flat::delta::CachedFlatStateChanges;
use crate::flat::store_helper::FlatStateColumn;
use crate::flat::{FlatStorageReadyStatus, FlatStorageStatus};
use crate::{Store, StoreUpdate};
use crate::{DBCol, Store, StoreUpdate};

use super::delta::{CachedFlatStateDelta, FlatStateDelta};
use super::metrics::FlatStorageMetrics;
Expand Down Expand Up @@ -301,13 +300,13 @@ impl FlatStorage {
// We should also take fixed accounts into account.
let mut store_update = guard.store.store_update();
let mut removed_items = 0;
for item in guard.store.iter(FlatStateColumn::State.to_db_col()) {
for item in guard.store.iter(DBCol::FlatState) {
let (key, _) =
item.map_err(|e| StorageError::StorageInconsistentState(e.to_string()))?;

if store_helper::key_belongs_to_shard(&key, &shard_layout, shard_id)? {
removed_items += 1;
store_update.delete(FlatStateColumn::State.to_db_col(), &key);
store_update.delete(DBCol::FlatState, &key);
}
}
info!(target: "store", %shard_id, %removed_items, "Removing old items from flat storage");
Expand Down
57 changes: 15 additions & 42 deletions core/store/src/flat/store_helper.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
use crate::flat::delta::{FlatStateChanges, KeyForFlatStateDelta};
use crate::flat::types::FlatStorageError;
use crate::{Store, StoreUpdate};
use crate::{DBCol, Store, StoreUpdate};
use near_primitives::errors::StorageError;
use near_primitives::hash::CryptoHash;
use near_primitives::shard_layout::{ShardLayout, ShardUId};
Expand All @@ -15,35 +15,14 @@ use super::types::{FlatStateValue, FlatStorageStatus};
pub const FLAT_STATE_HEAD_KEY_PREFIX: &[u8; 4] = b"HEAD";
pub const FLAT_STATE_CREATION_STATUS_KEY_PREFIX: &[u8; 6] = b"STATUS";

/// This was needed to avoid `#[cfg(feature = "protocol_feature_flat_state")]`
/// from `DBCol::FlatState*` cascading all over the code.
/// TODO (#7327): remove.
pub enum FlatStateColumn {
State,
Changes,
DeltaMetadata,
Status,
}

impl FlatStateColumn {
pub const fn to_db_col(&self) -> crate::DBCol {
match self {
FlatStateColumn::State => crate::DBCol::FlatState,
FlatStateColumn::Changes => crate::DBCol::FlatStateChanges,
FlatStateColumn::DeltaMetadata => crate::DBCol::FlatStateDeltaMetadata,
FlatStateColumn::Status => crate::DBCol::FlatStorageStatus,
}
}
}

pub fn get_delta_changes(
store: &Store,
shard_uid: ShardUId,
block_hash: CryptoHash,
) -> Result<Option<FlatStateChanges>, FlatStorageError> {
let key = KeyForFlatStateDelta { shard_uid, block_hash };
Ok(store
.get_ser::<FlatStateChanges>(FlatStateColumn::Changes.to_db_col(), &key.to_bytes())
.get_ser::<FlatStateChanges>(DBCol::FlatStateChanges, &key.to_bytes())
.map_err(|_| FlatStorageError::StorageInternalError)?)
}

Expand All @@ -52,7 +31,7 @@ pub fn get_all_deltas_metadata(
shard_uid: ShardUId,
) -> Result<Vec<FlatStateDeltaMetadata>, FlatStorageError> {
store
.iter_prefix_ser(FlatStateColumn::DeltaMetadata.to_db_col(), &shard_uid.to_bytes())
.iter_prefix_ser(DBCol::FlatStateDeltaMetadata, &shard_uid.to_bytes())
.map(|res| res.map(|(_, value)| value).map_err(|_| FlatStorageError::StorageInternalError))
.collect()
}
Expand All @@ -64,25 +43,25 @@ pub fn set_delta(
) -> Result<(), FlatStorageError> {
let key = KeyForFlatStateDelta { shard_uid, block_hash: delta.metadata.block.hash }.to_bytes();
store_update
.set_ser(FlatStateColumn::Changes.to_db_col(), &key, &delta.changes)
.set_ser(DBCol::FlatStateChanges, &key, &delta.changes)
.map_err(|_| FlatStorageError::StorageInternalError)?;
store_update
.set_ser(FlatStateColumn::DeltaMetadata.to_db_col(), &key, &delta.metadata)
.set_ser(DBCol::FlatStateDeltaMetadata, &key, &delta.metadata)
.map_err(|_| FlatStorageError::StorageInternalError)?;
Ok(())
}

pub fn remove_delta(store_update: &mut StoreUpdate, shard_uid: ShardUId, block_hash: CryptoHash) {
let key = KeyForFlatStateDelta { shard_uid, block_hash }.to_bytes();
store_update.delete(FlatStateColumn::Changes.to_db_col(), &key);
store_update.delete(FlatStateColumn::DeltaMetadata.to_db_col(), &key);
store_update.delete(DBCol::FlatStateChanges, &key);
store_update.delete(DBCol::FlatStateDeltaMetadata, &key);
}

pub fn remove_all_deltas(store_update: &mut StoreUpdate, shard_uid: ShardUId) {
let key_from = shard_uid.to_bytes();
let key_to = ShardUId::next_shard_prefix(&key_from);
store_update.delete_range(FlatStateColumn::Changes.to_db_col(), &key_from, &key_to);
store_update.delete_range(FlatStateColumn::DeltaMetadata.to_db_col(), &key_from, &key_to);
store_update.delete_range(DBCol::FlatStateChanges, &key_from, &key_to);
store_update.delete_range(DBCol::FlatStateDeltaMetadata, &key_from, &key_to);
}

fn encode_flat_state_db_key(shard_uid: ShardUId, key: &[u8]) -> Vec<u8> {
Expand Down Expand Up @@ -113,9 +92,7 @@ pub(crate) fn get_flat_state_value(
key: &[u8],
) -> Result<Option<FlatStateValue>, FlatStorageError> {
let db_key = encode_flat_state_db_key(shard_uid, key);
store
.get_ser(FlatStateColumn::State.to_db_col(), &db_key)
.map_err(|_| FlatStorageError::StorageInternalError)
store.get_ser(DBCol::FlatState, &db_key).map_err(|_| FlatStorageError::StorageInternalError)
}

// TODO(#8577): make pub(crate) once flat storage creator is moved inside `flat` module.
Expand All @@ -128,15 +105,15 @@ pub fn set_flat_state_value(
let db_key = encode_flat_state_db_key(shard_uid, &key);
match value {
Some(value) => store_update
.set_ser(FlatStateColumn::State.to_db_col(), &db_key, &value)
.set_ser(DBCol::FlatState, &db_key, &value)
.map_err(|_| FlatStorageError::StorageInternalError),
None => Ok(store_update.delete(FlatStateColumn::State.to_db_col(), &db_key)),
None => Ok(store_update.delete(DBCol::FlatState, &db_key)),
}
}

pub fn get_flat_storage_status(store: &Store, shard_uid: ShardUId) -> FlatStorageStatus {
store
.get_ser(FlatStateColumn::Status.to_db_col(), &shard_uid.to_bytes())
.get_ser(DBCol::FlatStorageStatus, &shard_uid.to_bytes())
.expect("Error reading flat head from storage")
.unwrap_or(FlatStorageStatus::Empty)
}
Expand All @@ -147,7 +124,7 @@ pub fn set_flat_storage_status(
status: FlatStorageStatus,
) {
store_update
.set_ser(FlatStateColumn::Status.to_db_col(), &shard_uid.to_bytes(), &status)
.set_ser(DBCol::FlatStorageStatus, &shard_uid.to_bytes(), &status)
.expect("Borsh should not fail")
}

Expand All @@ -165,11 +142,7 @@ pub fn iter_flat_state_entries<'a>(
to: Option<&'a Vec<u8>>,
) -> impl Iterator<Item = (Vec<u8>, Box<[u8]>)> + 'a {
store
.iter_range(
FlatStateColumn::State.to_db_col(),
from.map(|x| x.as_slice()),
to.map(|x| x.as_slice()),
)
.iter_range(DBCol::FlatState, from.map(|x| x.as_slice()), to.map(|x| x.as_slice()))
.filter_map(move |result| {
if let Ok((key, value)) = result {
// Currently all the data in flat storage is 'together' - so we have to parse the key,
Expand Down
4 changes: 2 additions & 2 deletions tools/flat-storage/src/commands.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use near_chain::{
use near_epoch_manager::{EpochManager, EpochManagerAdapter, EpochManagerHandle};
use near_primitives::{state::ValueRef, trie_key::trie_key_parsers::parse_account_id_from_raw_key};
use near_store::flat::{store_helper, FlatStateDelta, FlatStateDeltaMetadata, FlatStorageStatus};
use near_store::{Mode, NodeStorage, ShardUId, Store, StoreOpener};
use near_store::{DBCol, Mode, NodeStorage, ShardUId, Store, StoreOpener};
use nearcore::{load_config, NearConfig, NightshadeRuntime};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tqdm::tqdm;
Expand Down Expand Up @@ -122,7 +122,7 @@ impl FlatStorageCommand {
let (.., hot_store) =
Self::get_db(&opener, home_dir, &near_config, near_store::Mode::ReadOnly);
println!("DB version: {:?}", hot_store.get_db_version()?);
for item in hot_store.iter(store_helper::FlatStateColumn::Status.to_db_col()) {
for item in hot_store.iter(DBCol::FlatStorageStatus) {
let (bytes_shard_uid, status) = item?;
let shard_uid = ShardUId::try_from(bytes_shard_uid.as_ref()).unwrap();
let status = FlatStorageStatus::try_from_slice(&status)?;
Expand Down

0 comments on commit 48087bd

Please sign in to comment.