From 0686e6568ed24882437f769cf77f554b6ae320c4 Mon Sep 17 00:00:00 2001 From: Jure Bajic Date: Thu, 25 May 2023 11:34:58 +0200 Subject: [PATCH 1/6] Clean flat state via range --- core/store/src/flat/storage.rs | 26 +++++++++++--------------- core/store/src/flat/store_helper.rs | 11 ++++++++--- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/core/store/src/flat/storage.rs b/core/store/src/flat/storage.rs index 0f0552a57d7..38277ba380f 100644 --- a/core/store/src/flat/storage.rs +++ b/core/store/src/flat/storage.rs @@ -307,24 +307,20 @@ impl FlatStorage { // TODO (#7327): call it just after we stopped tracking a shard. // TODO (#7327): remove FlatStateChanges. Consider custom serialization of keys to remove them by // prefix. - // TODO (#7327): support range deletions which are much faster than naive deletions. For that, we - // can delete ranges of keys like - // [ [0]+boundary_accounts(shard_id) .. [0]+boundary_accounts(shard_id+1) ), etc. - // 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(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(DBCol::FlatState, &key); - } - } + store_helper::remove_range_by_shard_uid( + &mut store_update, + guard.shard_uid, + &[DBCol::FlatState], + ); + let removed_items = guard.store.iter(DBCol::FlatState).count(); info!(target: "store", %shard_id, %removed_items, "Removing old items from flat storage"); - store_helper::remove_all_deltas(&mut store_update, guard.shard_uid); + store_helper::remove_range_by_shard_uid( + &mut store_update, + guard.shard_uid, + &[DBCol::FlatStateChanges, DBCol::FlatStateDeltaMetadata], + ); store_helper::set_flat_storage_status( &mut store_update, guard.shard_uid, diff --git a/core/store/src/flat/store_helper.rs b/core/store/src/flat/store_helper.rs index 01b1296286b..52b93297c3f 100644 --- a/core/store/src/flat/store_helper.rs +++ b/core/store/src/flat/store_helper.rs @@ -57,11 +57,16 @@ pub fn remove_delta(store_update: &mut StoreUpdate, shard_uid: ShardUId, block_h store_update.delete(DBCol::FlatStateDeltaMetadata, &key); } -pub fn remove_all_deltas(store_update: &mut StoreUpdate, shard_uid: ShardUId) { +pub fn remove_range_by_shard_uid( + store_update: &mut StoreUpdate, + shard_uid: ShardUId, + cols: &[DBCol], +) { let key_from = shard_uid.to_bytes(); let key_to = ShardUId::next_shard_prefix(&key_from); - store_update.delete_range(DBCol::FlatStateChanges, &key_from, &key_to); - store_update.delete_range(DBCol::FlatStateDeltaMetadata, &key_from, &key_to); + for col in cols { + store_update.delete_range(*col, &key_from, &key_to); + } } pub(crate) fn encode_flat_state_db_key(shard_uid: ShardUId, key: &[u8]) -> Vec { From bb49b2d41d7ac715df430dccb2107742c9c6d40f Mon Sep 17 00:00:00 2001 From: Jure Bajic Date: Thu, 25 May 2023 11:46:13 +0200 Subject: [PATCH 2/6] Remove shard_layout arg --- chain/chain/src/chain.rs | 2 +- chain/chain/src/test_utils/kv_runtime.rs | 1 - chain/chain/src/types.rs | 1 - core/store/src/flat/manager.rs | 5 ++--- core/store/src/flat/storage.rs | 4 ++-- nearcore/src/runtime/mod.rs | 6 ++---- tools/flat-storage/src/commands.rs | 2 +- 7 files changed, 8 insertions(+), 13 deletions(-) diff --git a/chain/chain/src/chain.rs b/chain/chain/src/chain.rs index 16a76655eb8..4fd699af132 100644 --- a/chain/chain/src/chain.rs +++ b/chain/chain/src/chain.rs @@ -3203,7 +3203,7 @@ impl Chain { // Before working with state parts, remove existing flat storage data. let epoch_id = self.get_block_header(&sync_hash)?.epoch_id().clone(); let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, &epoch_id)?; - self.runtime_adapter.remove_flat_storage_for_shard(shard_uid, &epoch_id)?; + self.runtime_adapter.remove_flat_storage_for_shard(shard_uid)?; let shard_state_header = self.get_state_header(shard_id, sync_hash)?; let state_root = shard_state_header.chunk_prev_state_root(); diff --git a/chain/chain/src/test_utils/kv_runtime.rs b/chain/chain/src/test_utils/kv_runtime.rs index 664f8b3bc52..09d07bf0b09 100644 --- a/chain/chain/src/test_utils/kv_runtime.rs +++ b/chain/chain/src/test_utils/kv_runtime.rs @@ -987,7 +987,6 @@ impl RuntimeAdapter for KeyValueRuntime { fn remove_flat_storage_for_shard( &self, _shard_uid: ShardUId, - _epoch_id: &EpochId, ) -> Result<(), Error> { Ok(()) } diff --git a/chain/chain/src/types.rs b/chain/chain/src/types.rs index 2a26a694247..df12325c571 100644 --- a/chain/chain/src/types.rs +++ b/chain/chain/src/types.rs @@ -313,7 +313,6 @@ pub trait RuntimeAdapter: Send + Sync { fn remove_flat_storage_for_shard( &self, shard_uid: ShardUId, - epoch_id: &EpochId, ) -> Result<(), Error>; fn set_flat_storage_for_genesis( diff --git a/core/store/src/flat/manager.rs b/core/store/src/flat/manager.rs index ccae92aabd1..2731585b007 100644 --- a/core/store/src/flat/manager.rs +++ b/core/store/src/flat/manager.rs @@ -3,7 +3,7 @@ use crate::flat::{ }; use near_primitives::errors::StorageError; use near_primitives::hash::CryptoHash; -use near_primitives::shard_layout::{ShardLayout, ShardUId}; +use near_primitives::shard_layout::ShardUId; use near_primitives::types::BlockHeight; use std::collections::HashMap; use std::sync::{Arc, Mutex}; @@ -143,14 +143,13 @@ impl FlatStorageManager { pub fn remove_flat_storage_for_shard( &self, shard_uid: ShardUId, - shard_layout: ShardLayout, ) -> Result<(), StorageError> { let mut flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR); match flat_storages.remove(&shard_uid) { None => {} Some(flat_storage) => { - flat_storage.clear_state(shard_layout)?; + flat_storage.clear_state()?; } } diff --git a/core/store/src/flat/storage.rs b/core/store/src/flat/storage.rs index 38277ba380f..c35206d8264 100644 --- a/core/store/src/flat/storage.rs +++ b/core/store/src/flat/storage.rs @@ -3,7 +3,7 @@ use std::sync::{Arc, RwLock}; use near_primitives::errors::StorageError; use near_primitives::hash::CryptoHash; -use near_primitives::shard_layout::{ShardLayout, ShardUId}; +use near_primitives::shard_layout::ShardUId; use tracing::{debug, info, warn}; use crate::flat::delta::CachedFlatStateChanges; @@ -298,7 +298,7 @@ impl FlatStorage { } /// Clears all State key-value pairs from flat storage. - pub fn clear_state(&self, shard_layout: ShardLayout) -> Result<(), StorageError> { + pub fn clear_state(&self) -> Result<(), StorageError> { let guard = self.0.write().expect(super::POISONED_LOCK_ERR); let shard_id = guard.shard_uid.shard_id(); diff --git a/nearcore/src/runtime/mod.rs b/nearcore/src/runtime/mod.rs index 3790d0bb5bd..8ff016af931 100644 --- a/nearcore/src/runtime/mod.rs +++ b/nearcore/src/runtime/mod.rs @@ -749,12 +749,10 @@ impl RuntimeAdapter for NightshadeRuntime { fn remove_flat_storage_for_shard( &self, - shard_uid: ShardUId, - epoch_id: &EpochId, + shard_uid: ShardUId ) -> Result<(), Error> { - let shard_layout = self.epoch_manager.get_shard_layout(epoch_id)?; self.flat_storage_manager - .remove_flat_storage_for_shard(shard_uid, shard_layout) + .remove_flat_storage_for_shard(shard_uid) .map_err(Error::StorageError)?; Ok(()) } diff --git a/tools/flat-storage/src/commands.rs b/tools/flat-storage/src/commands.rs index 7db0711b9ec..bb2d69924f5 100644 --- a/tools/flat-storage/src/commands.rs +++ b/tools/flat-storage/src/commands.rs @@ -175,7 +175,7 @@ impl FlatStorageCommand { let shard_uid = epoch_manager.shard_id_to_uid(reset_cmd.shard_id, &tip.epoch_id)?; rw_hot_runtime.create_flat_storage_for_shard(shard_uid); - rw_hot_runtime.remove_flat_storage_for_shard(shard_uid, &tip.epoch_id)?; + rw_hot_runtime.remove_flat_storage_for_shard(shard_uid)?; } SubCommand::Init(init_cmd) => { let (_, epoch_manager, rw_hot_runtime, rw_chain_store, rw_hot_store) = Self::get_db( From 1f962c8621f9d9d4ba0ee59280b17be599ec1fb2 Mon Sep 17 00:00:00 2001 From: Jure Bajic Date: Thu, 25 May 2023 12:00:10 +0200 Subject: [PATCH 3/6] Fix format errors --- chain/chain/src/test_utils/kv_runtime.rs | 5 +---- chain/chain/src/types.rs | 5 +---- core/store/src/flat/manager.rs | 5 +---- nearcore/src/runtime/mod.rs | 5 +---- 4 files changed, 4 insertions(+), 16 deletions(-) diff --git a/chain/chain/src/test_utils/kv_runtime.rs b/chain/chain/src/test_utils/kv_runtime.rs index 09d07bf0b09..c83b6837341 100644 --- a/chain/chain/src/test_utils/kv_runtime.rs +++ b/chain/chain/src/test_utils/kv_runtime.rs @@ -984,10 +984,7 @@ impl RuntimeAdapter for KeyValueRuntime { panic!("Flat storage state can't be created for shard {shard_uid} because KeyValueRuntime doesn't support this"); } - fn remove_flat_storage_for_shard( - &self, - _shard_uid: ShardUId, - ) -> Result<(), Error> { + fn remove_flat_storage_for_shard(&self, _shard_uid: ShardUId) -> Result<(), Error> { Ok(()) } diff --git a/chain/chain/src/types.rs b/chain/chain/src/types.rs index df12325c571..3d370893211 100644 --- a/chain/chain/src/types.rs +++ b/chain/chain/src/types.rs @@ -310,10 +310,7 @@ pub trait RuntimeAdapter: Send + Sync { /// Removes flat storage state for shard, if it exists. /// Used to clear old flat storage data from disk and memory before syncing to newer state. - fn remove_flat_storage_for_shard( - &self, - shard_uid: ShardUId, - ) -> Result<(), Error>; + fn remove_flat_storage_for_shard(&self, shard_uid: ShardUId) -> Result<(), Error>; fn set_flat_storage_for_genesis( &self, diff --git a/core/store/src/flat/manager.rs b/core/store/src/flat/manager.rs index 2731585b007..40330c11869 100644 --- a/core/store/src/flat/manager.rs +++ b/core/store/src/flat/manager.rs @@ -140,10 +140,7 @@ impl FlatStorageManager { flat_storages.get(&shard_uid).cloned() } - pub fn remove_flat_storage_for_shard( - &self, - shard_uid: ShardUId, - ) -> Result<(), StorageError> { + pub fn remove_flat_storage_for_shard(&self, shard_uid: ShardUId) -> Result<(), StorageError> { let mut flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR); match flat_storages.remove(&shard_uid) { diff --git a/nearcore/src/runtime/mod.rs b/nearcore/src/runtime/mod.rs index 8ff016af931..f5546f9b8ad 100644 --- a/nearcore/src/runtime/mod.rs +++ b/nearcore/src/runtime/mod.rs @@ -747,10 +747,7 @@ impl RuntimeAdapter for NightshadeRuntime { self.flat_storage_manager.add_flat_storage_for_shard(shard_uid, flat_storage); } - fn remove_flat_storage_for_shard( - &self, - shard_uid: ShardUId - ) -> Result<(), Error> { + fn remove_flat_storage_for_shard(&self, shard_uid: ShardUId) -> Result<(), Error> { self.flat_storage_manager .remove_flat_storage_for_shard(shard_uid) .map_err(Error::StorageError)?; From 1dc75114d0ecb70747e5d057c49be2c13f42adfa Mon Sep 17 00:00:00 2001 From: Jure Bajic Date: Thu, 25 May 2023 12:49:48 +0200 Subject: [PATCH 4/6] Resolve clippy errors --- integration-tests/src/tests/client/flat_storage.rs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/integration-tests/src/tests/client/flat_storage.rs b/integration-tests/src/tests/client/flat_storage.rs index 1a788ece71e..5e44b77dfd3 100644 --- a/integration-tests/src/tests/client/flat_storage.rs +++ b/integration-tests/src/tests/client/flat_storage.rs @@ -159,12 +159,10 @@ fn test_flat_storage_creation_sanity() { ); } - let block_hash = env.clients[0].chain.get_block_hash_by_height(START_HEIGHT - 1).unwrap(); - let epoch_id = env.clients[0].chain.epoch_manager.get_epoch_id(&block_hash).unwrap(); env.clients[0] .chain .runtime_adapter - .remove_flat_storage_for_shard(shard_uid, &epoch_id) + .remove_flat_storage_for_shard(shard_uid) .unwrap(); } @@ -247,12 +245,10 @@ fn test_flat_storage_creation_two_shards() { ); } - let block_hash = env.clients[0].chain.get_block_hash_by_height(START_HEIGHT - 1).unwrap(); - let epoch_id = env.clients[0].chain.epoch_manager.get_epoch_id(&block_hash).unwrap(); env.clients[0] .chain .runtime_adapter - .remove_flat_storage_for_shard(shard_uids[0], &epoch_id) + .remove_flat_storage_for_shard(shard_uids[0]) .unwrap(); } @@ -393,12 +389,10 @@ fn test_catchup_succeeds_even_if_no_new_blocks() { env.produce_block(0, height); } // Remove flat storage. - let block_hash = env.clients[0].chain.get_block_hash_by_height(START_HEIGHT - 1).unwrap(); - let epoch_id = env.clients[0].chain.epoch_manager.get_epoch_id(&block_hash).unwrap(); env.clients[0] .chain .runtime_adapter - .remove_flat_storage_for_shard(shard_uid, &epoch_id) + .remove_flat_storage_for_shard(shard_uid) .unwrap(); } let mut env = setup_env(&genesis, store.clone()); From 602c4658d404431c57fff70ab682b9dbb3c9e051 Mon Sep 17 00:00:00 2001 From: Jure Bajic Date: Thu, 25 May 2023 15:03:00 +0200 Subject: [PATCH 5/6] Fix format errors --- .../src/tests/client/flat_storage.rs | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/integration-tests/src/tests/client/flat_storage.rs b/integration-tests/src/tests/client/flat_storage.rs index 5e44b77dfd3..cb5ef12dc27 100644 --- a/integration-tests/src/tests/client/flat_storage.rs +++ b/integration-tests/src/tests/client/flat_storage.rs @@ -159,11 +159,7 @@ fn test_flat_storage_creation_sanity() { ); } - env.clients[0] - .chain - .runtime_adapter - .remove_flat_storage_for_shard(shard_uid) - .unwrap(); + env.clients[0].chain.runtime_adapter.remove_flat_storage_for_shard(shard_uid).unwrap(); } // Create new chain and runtime using the same store. It should produce next blocks normally, but now it should @@ -245,11 +241,7 @@ fn test_flat_storage_creation_two_shards() { ); } - env.clients[0] - .chain - .runtime_adapter - .remove_flat_storage_for_shard(shard_uids[0]) - .unwrap(); + env.clients[0].chain.runtime_adapter.remove_flat_storage_for_shard(shard_uids[0]).unwrap(); } // Check that flat storage is not ready for shard 0 but ready for shard 1. @@ -389,11 +381,7 @@ fn test_catchup_succeeds_even_if_no_new_blocks() { env.produce_block(0, height); } // Remove flat storage. - env.clients[0] - .chain - .runtime_adapter - .remove_flat_storage_for_shard(shard_uid) - .unwrap(); + env.clients[0].chain.runtime_adapter.remove_flat_storage_for_shard(shard_uid).unwrap(); } let mut env = setup_env(&genesis, store.clone()); assert!(env.clients[0].runtime_adapter.get_flat_storage_for_shard(shard_uid).is_none()); From 7aa6011974904d437a5c53c8545dc1eaeb057449 Mon Sep 17 00:00:00 2001 From: Jure Bajic Date: Fri, 26 May 2023 09:34:49 +0200 Subject: [PATCH 6/6] Address review comments --- core/store/src/flat/manager.rs | 7 ++----- core/store/src/flat/storage.rs | 24 ++++-------------------- core/store/src/flat/store_helper.rs | 19 +++++++++++-------- 3 files changed, 17 insertions(+), 33 deletions(-) diff --git a/core/store/src/flat/manager.rs b/core/store/src/flat/manager.rs index 40330c11869..859c0e3c596 100644 --- a/core/store/src/flat/manager.rs +++ b/core/store/src/flat/manager.rs @@ -143,11 +143,8 @@ impl FlatStorageManager { pub fn remove_flat_storage_for_shard(&self, shard_uid: ShardUId) -> Result<(), StorageError> { let mut flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR); - match flat_storages.remove(&shard_uid) { - None => {} - Some(flat_storage) => { - flat_storage.clear_state()?; - } + if let Some(flat_store) = flat_storages.remove(&shard_uid) { + flat_store.clear_state()?; } Ok(()) diff --git a/core/store/src/flat/storage.rs b/core/store/src/flat/storage.rs index c35206d8264..7bde0ce8202 100644 --- a/core/store/src/flat/storage.rs +++ b/core/store/src/flat/storage.rs @@ -4,11 +4,11 @@ use std::sync::{Arc, RwLock}; use near_primitives::errors::StorageError; use near_primitives::hash::CryptoHash; use near_primitives::shard_layout::ShardUId; -use tracing::{debug, info, warn}; +use tracing::{debug, warn}; use crate::flat::delta::CachedFlatStateChanges; use crate::flat::{FlatStorageReadyStatus, FlatStorageStatus}; -use crate::{DBCol, Store, StoreUpdate}; +use crate::{Store, StoreUpdate}; use super::delta::{CachedFlatStateDelta, FlatStateDelta}; use super::metrics::FlatStorageMetrics; @@ -300,27 +300,11 @@ impl FlatStorage { /// Clears all State key-value pairs from flat storage. pub fn clear_state(&self) -> Result<(), StorageError> { let guard = self.0.write().expect(super::POISONED_LOCK_ERR); - let shard_id = guard.shard_uid.shard_id(); - // Removes all items belonging to the shard one by one. - // Note that it does not work for resharding. - // TODO (#7327): call it just after we stopped tracking a shard. - // TODO (#7327): remove FlatStateChanges. Consider custom serialization of keys to remove them by - // prefix. let mut store_update = guard.store.store_update(); - store_helper::remove_range_by_shard_uid( - &mut store_update, - guard.shard_uid, - &[DBCol::FlatState], - ); - let removed_items = guard.store.iter(DBCol::FlatState).count(); - info!(target: "store", %shard_id, %removed_items, "Removing old items from flat storage"); + store_helper::remove_all_flat_state_values(&mut store_update, guard.shard_uid); - store_helper::remove_range_by_shard_uid( - &mut store_update, - guard.shard_uid, - &[DBCol::FlatStateChanges, DBCol::FlatStateDeltaMetadata], - ); + store_helper::remove_all_deltas(&mut store_update, guard.shard_uid); store_helper::set_flat_storage_status( &mut store_update, guard.shard_uid, diff --git a/core/store/src/flat/store_helper.rs b/core/store/src/flat/store_helper.rs index 52b93297c3f..389f3077f69 100644 --- a/core/store/src/flat/store_helper.rs +++ b/core/store/src/flat/store_helper.rs @@ -57,16 +57,19 @@ pub fn remove_delta(store_update: &mut StoreUpdate, shard_uid: ShardUId, block_h store_update.delete(DBCol::FlatStateDeltaMetadata, &key); } -pub fn remove_range_by_shard_uid( - store_update: &mut StoreUpdate, - shard_uid: ShardUId, - cols: &[DBCol], -) { +fn remove_range_by_shard_uid(store_update: &mut StoreUpdate, shard_uid: ShardUId, col: DBCol) { let key_from = shard_uid.to_bytes(); let key_to = ShardUId::next_shard_prefix(&key_from); - for col in cols { - store_update.delete_range(*col, &key_from, &key_to); - } + store_update.delete_range(col, &key_from, &key_to); +} + +pub fn remove_all_deltas(store_update: &mut StoreUpdate, shard_uid: ShardUId) { + remove_range_by_shard_uid(store_update, shard_uid, DBCol::FlatStateChanges); + remove_range_by_shard_uid(store_update, shard_uid, DBCol::FlatStateDeltaMetadata); +} + +pub fn remove_all_flat_state_values(store_update: &mut StoreUpdate, shard_uid: ShardUId) { + remove_range_by_shard_uid(store_update, shard_uid, DBCol::FlatState); } pub(crate) fn encode_flat_state_db_key(shard_uid: ShardUId, key: &[u8]) -> Vec {