Skip to content

Commit d230eaf

Browse files
pugachAGpugachAG
authored
feat: introduce flat storage status (near#8689)
Part of near#8577 This PR introduces `FlatStorageStatus`: * it generalises `FlatStorageCreationStatus` * `FlatStorageCreationStatus::DontCreate` is moved to `::Disabled` status * flat head is now stored as a part of `::Ready` status. * `::Empty` status is added to explicitly express that flat storage does not exist * we no longer need `FlatStateMisc` column, so it is removed nayduck run: https://nayduck.near.org/#/run/2905
1 parent 229cf47 commit d230eaf

File tree

15 files changed

+384
-332
lines changed

15 files changed

+384
-332
lines changed

‎chain/chain/src/chain.rs

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,8 @@ use near_primitives::views::{
7070
LightClientBlockView, SignedTransactionView,
7171
};
7272
use near_store::flat::{
73-
store_helper, FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata,
74-
FlatStorageCreationStatus, FlatStorageError,
73+
store_helper, FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata, FlatStorageError,
74+
FlatStorageReadyStatus, FlatStorageStatus,
7575
};
7676
use near_store::StorageError;
7777
use near_store::{DBCol, ShardTries, StoreUpdate, WrappedTrieChanges};
@@ -633,6 +633,7 @@ impl Chain {
633633
if cfg!(feature = "protocol_feature_flat_state") {
634634
let tmp_store_update = runtime_adapter.set_flat_storage_for_genesis(
635635
genesis.hash(),
636+
genesis.header().height(),
636637
genesis.header().epoch_id(),
637638
)?;
638639
store_update.merge(tmp_store_update);
@@ -3178,23 +3179,32 @@ impl Chain {
31783179
// flat storage. Now we can set flat head to hash of this block and create flat storage.
31793180
// TODO (#7327): ensure that no flat storage work is done for `KeyValueRuntime`.
31803181
if cfg!(feature = "protocol_feature_flat_state") {
3181-
let mut store_update = self.runtime_adapter.store().store_update();
3182-
store_helper::set_flat_head(&mut store_update, shard_id, block_hash);
3183-
store_update.commit()?;
3184-
}
3185-
3186-
if self.runtime_adapter.get_flat_storage_creation_status(shard_id)
3187-
== FlatStorageCreationStatus::Ready
3188-
{
31893182
// If block_hash is equal to default - this means that we're all the way back at genesis.
31903183
// So we don't have to add the storage state for shard in such case.
31913184
// TODO(8438) - add additional test scenarios for this case.
31923185
if *block_hash != CryptoHash::default() {
31933186
let block_header = self.get_block_header(block_hash)?;
31943187
let epoch_id = block_header.epoch_id();
31953188
let shard_uid = self.runtime_adapter.shard_id_to_uid(shard_id, epoch_id)?;
3196-
3197-
self.runtime_adapter.create_flat_storage_for_shard(shard_uid);
3189+
if !matches!(
3190+
self.runtime_adapter.get_flat_storage_status(shard_uid),
3191+
FlatStorageStatus::Disabled
3192+
) {
3193+
let mut store_update = self.runtime_adapter.store().store_update();
3194+
store_helper::set_flat_storage_status(
3195+
&mut store_update,
3196+
shard_uid,
3197+
FlatStorageStatus::Ready(FlatStorageReadyStatus {
3198+
flat_head: near_store::flat::BlockInfo {
3199+
hash: *block_hash,
3200+
prev_hash: *block_header.prev_hash(),
3201+
height: block_header.height(),
3202+
},
3203+
}),
3204+
);
3205+
store_update.commit()?;
3206+
self.runtime_adapter.create_flat_storage_for_shard(shard_uid);
3207+
}
31983208
}
31993209
}
32003210

‎chain/chain/src/flat_storage_creator.rs

Lines changed: 56 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ use near_primitives::state::ValueRef;
1919
use near_primitives::state_part::PartId;
2020
use near_primitives::types::{AccountId, BlockHeight, StateRoot};
2121
use near_store::flat::{
22-
store_helper, FetchingStateStatus, FlatStateChanges, FlatStorageCreationStatus,
23-
NUM_PARTS_IN_ONE_STEP, STATE_PART_MEMORY_LIMIT,
22+
store_helper, BlockInfo, FetchingStateStatus, FlatStateChanges, FlatStorageCreationStatus,
23+
FlatStorageReadyStatus, FlatStorageStatus, NUM_PARTS_IN_ONE_STEP, STATE_PART_MEMORY_LIMIT,
2424
};
2525
use near_store::{Store, FLAT_STORAGE_HEAD_HEIGHT};
2626
use near_store::{Trie, TrieDBStorage, TrieTraversalItem};
@@ -165,10 +165,19 @@ impl FlatStorageShardCreator {
165165
) -> Result<bool, Error> {
166166
let shard_id = self.shard_uid.shard_id();
167167
let current_status =
168-
store_helper::get_flat_storage_creation_status(chain_store.store(), shard_id);
168+
store_helper::get_flat_storage_status(chain_store.store(), self.shard_uid);
169169
self.metrics.status.set((&current_status).into());
170170
match &current_status {
171-
FlatStorageCreationStatus::SavingDeltas => {
171+
FlatStorageStatus::Empty => {
172+
let mut store_update = chain_store.store().store_update();
173+
store_helper::set_flat_storage_status(
174+
&mut store_update,
175+
self.shard_uid,
176+
FlatStorageStatus::Creation(FlatStorageCreationStatus::SavingDeltas),
177+
);
178+
store_update.commit()?;
179+
}
180+
FlatStorageStatus::Creation(FlatStorageCreationStatus::SavingDeltas) => {
172181
let final_head = chain_store.final_head()?;
173182
let final_height = final_head.height;
174183

@@ -220,15 +229,19 @@ impl FlatStorageShardCreator {
220229

221230
let mut store_update = chain_store.store().store_update();
222231
self.metrics.flat_head_height.set(final_head.height as i64);
223-
store_helper::set_flat_storage_creation_status(
232+
store_helper::set_flat_storage_status(
224233
&mut store_update,
225-
shard_id,
226-
FlatStorageCreationStatus::FetchingState(status),
234+
self.shard_uid,
235+
FlatStorageStatus::Creation(FlatStorageCreationStatus::FetchingState(
236+
status,
237+
)),
227238
);
228239
store_update.commit()?;
229240
}
230241
}
231-
FlatStorageCreationStatus::FetchingState(fetching_state_status) => {
242+
FlatStorageStatus::Creation(FlatStorageCreationStatus::FetchingState(
243+
fetching_state_status,
244+
)) => {
232245
let store = self.runtime_adapter.store().clone();
233246
let block_hash = fetching_state_status.block_hash;
234247
let start_part_id = fetching_state_status.part_id;
@@ -297,26 +310,30 @@ impl FlatStorageShardCreator {
297310
num_parts,
298311
};
299312
debug!(target: "chain", %shard_id, %block_hash, ?new_status);
300-
store_helper::set_flat_storage_creation_status(
313+
store_helper::set_flat_storage_status(
301314
&mut store_update,
302-
shard_id,
303-
FlatStorageCreationStatus::FetchingState(new_status),
315+
self.shard_uid,
316+
FlatStorageStatus::Creation(
317+
FlatStorageCreationStatus::FetchingState(new_status),
318+
),
304319
);
305320
} else {
306321
// If all parts were fetched, we can start catchup.
307322
info!(target: "chain", %shard_id, %block_hash, "Finished fetching state");
308323
self.metrics.remaining_state_parts.set(0);
309-
store_helper::set_flat_storage_creation_status(
324+
store_helper::set_flat_storage_status(
310325
&mut store_update,
311-
shard_id,
312-
FlatStorageCreationStatus::CatchingUp(block_hash),
326+
self.shard_uid,
327+
FlatStorageStatus::Creation(FlatStorageCreationStatus::CatchingUp(
328+
block_hash,
329+
)),
313330
);
314331
}
315332
store_update.commit()?;
316333
}
317334
}
318335
}
319-
FlatStorageCreationStatus::CatchingUp(old_flat_head) => {
336+
FlatStorageStatus::Creation(FlatStorageCreationStatus::CatchingUp(old_flat_head)) => {
320337
let store = self.runtime_adapter.store();
321338
let mut flat_head = *old_flat_head;
322339
let chain_final_head = chain_store.final_head()?;
@@ -346,37 +363,46 @@ impl FlatStorageShardCreator {
346363
let epoch_id = self.runtime_adapter.get_epoch_id(&flat_head)?;
347364
let shard_uid = self.runtime_adapter.shard_id_to_uid(shard_id, &epoch_id)?;
348365
let old_height = chain_store.get_block_height(&old_flat_head).unwrap();
349-
let height = chain_store.get_block_height(&flat_head).unwrap();
366+
let flat_head_block_header = chain_store.get_block_header(&flat_head).unwrap();
367+
let height = flat_head_block_header.height();
350368
debug!(target: "chain", %shard_id, %old_flat_head, %old_height, %flat_head, %height, "Catching up flat head");
351369
self.metrics.flat_head_height.set(height as i64);
352370
let mut store_update = self.runtime_adapter.store().store_update();
353371
merged_changes.apply_to_flat_state(&mut store_update, shard_uid);
354372
if flat_head == chain_final_head.last_block_hash {
355373
// If we reached chain final head, we can finish catchup and finally create flat storage.
356-
store_helper::remove_flat_storage_creation_status(
374+
store_helper::set_flat_storage_status(
357375
&mut store_update,
358-
shard_id,
376+
self.shard_uid,
377+
FlatStorageStatus::Ready(FlatStorageReadyStatus {
378+
flat_head: BlockInfo {
379+
hash: flat_head,
380+
prev_hash: *flat_head_block_header.prev_hash(),
381+
height,
382+
},
383+
}),
359384
);
360-
store_helper::set_flat_head(&mut store_update, shard_id, &flat_head);
361385
store_update.commit()?;
362386
self.runtime_adapter.create_flat_storage_for_shard(shard_uid);
363387
info!(target: "chain", %shard_id, %flat_head, %height, "Flat storage creation done");
364388
} else {
365-
store_helper::set_flat_storage_creation_status(
389+
store_helper::set_flat_storage_status(
366390
&mut store_update,
367-
shard_id,
368-
FlatStorageCreationStatus::CatchingUp(flat_head),
391+
self.shard_uid,
392+
FlatStorageStatus::Creation(FlatStorageCreationStatus::CatchingUp(
393+
flat_head,
394+
)),
369395
);
370396
store_update.commit()?;
371397
}
372398
}
373399
}
374-
FlatStorageCreationStatus::Ready => {}
375-
FlatStorageCreationStatus::DontCreate => {
376-
panic!("We initiated flat storage creation for shard {shard_id} but according to flat storage state status in db it cannot be created");
400+
FlatStorageStatus::Ready(_) => return Ok(true),
401+
FlatStorageStatus::Disabled => {
402+
panic!("initiated flat storage creation for shard {shard_id} while it is disabled");
377403
}
378404
};
379-
Ok(current_status == FlatStorageCreationStatus::Ready)
405+
Ok(false)
380406
}
381407
}
382408

@@ -402,15 +428,14 @@ impl FlatStorageCreator {
402428
let mut creation_needed = false;
403429
for shard_id in 0..num_shards {
404430
if runtime_adapter.cares_about_shard(me, &chain_head.prev_block_hash, shard_id, true) {
405-
let status = runtime_adapter.get_flat_storage_creation_status(shard_id);
406431
let shard_uid = runtime_adapter.shard_id_to_uid(shard_id, &chain_head.epoch_id)?;
432+
let status = runtime_adapter.get_flat_storage_status(shard_uid);
407433

408434
match status {
409-
FlatStorageCreationStatus::Ready => {
435+
FlatStorageStatus::Ready(_) => {
410436
runtime_adapter.create_flat_storage_for_shard(shard_uid);
411437
}
412-
FlatStorageCreationStatus::DontCreate => {}
413-
_ => {
438+
FlatStorageStatus::Empty | FlatStorageStatus::Creation(_) => {
414439
creation_needed = true;
415440
shard_creators.insert(
416441
shard_uid,
@@ -421,6 +446,7 @@ impl FlatStorageCreator {
421446
),
422447
);
423448
}
449+
FlatStorageStatus::Disabled => {}
424450
}
425451
}
426452
}

‎chain/chain/src/store.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2438,7 +2438,7 @@ impl<'a> ChainStoreUpdate<'a> {
24382438
DBCol::FlatState
24392439
| DBCol::FlatStateChanges
24402440
| DBCol::FlatStateDeltaMetadata
2441-
| DBCol::FlatStateMisc => {
2441+
| DBCol::FlatStorageStatus => {
24422442
unreachable!();
24432443
}
24442444
}

‎chain/chain/src/test_utils/kv_runtime.rs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ use crate::{BlockHeader, RuntimeWithEpochManagerAdapter};
5252

5353
use near_primitives::epoch_manager::ShardConfig;
5454

55-
use near_store::flat::{FlatStorage, FlatStorageCreationStatus};
55+
use near_store::flat::{FlatStorage, FlatStorageStatus};
5656

5757
use super::ValidatorSchedule;
5858

@@ -833,8 +833,8 @@ impl RuntimeAdapter for KeyValueRuntime {
833833
None
834834
}
835835

836-
fn get_flat_storage_creation_status(&self, _shard_id: ShardId) -> FlatStorageCreationStatus {
837-
FlatStorageCreationStatus::DontCreate
836+
fn get_flat_storage_status(&self, _shard_uid: ShardUId) -> FlatStorageStatus {
837+
FlatStorageStatus::Disabled
838838
}
839839

840840
fn create_flat_storage_for_shard(&self, shard_uid: ShardUId) {
@@ -852,6 +852,7 @@ impl RuntimeAdapter for KeyValueRuntime {
852852
fn set_flat_storage_for_genesis(
853853
&self,
854854
_genesis_block: &CryptoHash,
855+
_genesis_block_height: BlockHeight,
855856
_genesis_epoch_id: &EpochId,
856857
) -> Result<StoreUpdate, Error> {
857858
Ok(self.store.store_update())

‎chain/chain/src/types.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ use near_primitives::version::{
3131
MIN_PROTOCOL_VERSION_NEP_92_FIX,
3232
};
3333
use near_primitives::views::{QueryRequest, QueryResponse};
34-
use near_store::flat::{FlatStorage, FlatStorageCreationStatus};
34+
use near_store::flat::{FlatStorage, FlatStorageStatus};
3535
use near_store::{PartialStorage, ShardTries, Store, StoreUpdate, Trie, WrappedTrieChanges};
3636

3737
pub use near_epoch_manager::EpochManagerAdapter;
@@ -299,8 +299,7 @@ pub trait RuntimeAdapter: Send + Sync {
299299

300300
fn get_flat_storage_for_shard(&self, shard_id: ShardId) -> Option<FlatStorage>;
301301

302-
/// Gets status of flat storage state background creation.
303-
fn get_flat_storage_creation_status(&self, shard_id: ShardId) -> FlatStorageCreationStatus;
302+
fn get_flat_storage_status(&self, shard_uid: ShardUId) -> FlatStorageStatus;
304303

305304
/// Creates flat storage state for given shard, assuming that all flat storage data
306305
/// is already stored in DB.
@@ -318,6 +317,7 @@ pub trait RuntimeAdapter: Send + Sync {
318317
fn set_flat_storage_for_genesis(
319318
&self,
320319
genesis_block: &CryptoHash,
320+
genesis_block_height: BlockHeight,
321321
genesis_epoch_id: &EpochId,
322322
) -> Result<StoreUpdate, Error>;
323323

‎core/store/src/columns.rs

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -258,27 +258,25 @@ pub enum DBCol {
258258
/// *Column type*: ExecutionOutcomeWithProof
259259
TransactionResultForBlock,
260260
/// Flat state contents. Used to get `ValueRef` by trie key faster than doing a trie lookup.
261-
/// - *Rows*: trie key (Vec<u8>)
261+
/// - *Rows*: `shard_uid` + trie key (Vec<u8>)
262262
/// - *Column type*: ValueRef
263263
#[cfg(feature = "protocol_feature_flat_state")]
264264
FlatState,
265265
/// Changes for flat state delta. Stores how flat state should be updated for the given shard and block.
266-
/// - *Rows*: `KeyForFlatStateDelta { shard_id, block_hash }`
266+
/// - *Rows*: `KeyForFlatStateDelta { shard_uid, block_hash }`
267267
/// - *Column type*: `FlatStateChanges`
268268
#[cfg(feature = "protocol_feature_flat_state")]
269269
FlatStateChanges,
270270
/// Metadata for flat state delta.
271-
/// - *Rows*: `KeyForFlatStateDelta { shard_id, block_hash }`
271+
/// - *Rows*: `KeyForFlatStateDelta { shard_uid, block_hash }`
272272
/// - *Column type*: `FlatStateDeltaMetadata`
273273
#[cfg(feature = "protocol_feature_flat_state")]
274274
FlatStateDeltaMetadata,
275-
/// Miscellaneous data for flat state. Stores intermediate flat storage creation statuses and flat
276-
/// state heads for each shard.
277-
/// - *Rows*: Unique key prefix (e.g. `FLAT_STATE_HEAD_KEY_PREFIX`) + ShardId
278-
/// - *Column type*: FetchingStateStatus || flat storage catchup status (bool) || flat storage head (CryptoHash)
279-
// TODO (#7327): use only during testing, come up with proper format.
275+
/// Flat storage status for the corresponding shard.
276+
/// - *Rows*: `shard_uid`
277+
/// - *Column type*: `FlatStorageStatus`
280278
#[cfg(feature = "protocol_feature_flat_state")]
281-
FlatStateMisc,
279+
FlatStorageStatus,
282280
}
283281

284282
/// Defines different logical parts of a db key.
@@ -483,7 +481,7 @@ impl DBCol {
483481
#[cfg(feature = "protocol_feature_flat_state")]
484482
DBCol::FlatStateDeltaMetadata => &[DBKeyType::ShardId, DBKeyType::BlockHash],
485483
#[cfg(feature = "protocol_feature_flat_state")]
486-
DBCol::FlatStateMisc => &[DBKeyType::ShardId],
484+
DBCol::FlatStorageStatus => &[DBKeyType::ShardUId],
487485
}
488486
}
489487
}

‎core/store/src/flat/manager.rs

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1-
use crate::flat::{store_helper, POISONED_LOCK_ERR};
1+
use crate::flat::{
2+
store_helper, BlockInfo, FlatStorageReadyStatus, FlatStorageStatus, POISONED_LOCK_ERR,
3+
};
24
use near_primitives::errors::StorageError;
35
use near_primitives::hash::CryptoHash;
4-
use near_primitives::shard_layout::ShardLayout;
5-
use near_primitives::types::ShardId;
6+
use near_primitives::shard_layout::{ShardLayout, ShardUId};
7+
use near_primitives::types::{BlockHeight, ShardId};
68
use std::collections::HashMap;
79
use std::sync::{Arc, Mutex};
810
use tracing::debug;
@@ -44,12 +46,20 @@ impl FlatStorageManager {
4446
pub fn set_flat_storage_for_genesis(
4547
&self,
4648
store_update: &mut StoreUpdate,
47-
shard_id: ShardId,
49+
shard_uid: ShardUId,
4850
genesis_block: &CryptoHash,
51+
genesis_height: BlockHeight,
4952
) {
53+
let shard_id = shard_uid.shard_id();
5054
let flat_storages = self.0.flat_storages.lock().expect(POISONED_LOCK_ERR);
5155
assert!(!flat_storages.contains_key(&shard_id));
52-
store_helper::set_flat_head(store_update, shard_id, genesis_block);
56+
store_helper::set_flat_storage_status(
57+
store_update,
58+
shard_uid,
59+
FlatStorageStatus::Ready(FlatStorageReadyStatus {
60+
flat_head: BlockInfo::genesis(*genesis_block, genesis_height),
61+
}),
62+
);
5363
}
5464

5565
/// Add a flat storage state for shard `shard_id`. The function also checks that

‎core/store/src/flat/mod.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,10 @@ pub use chunk_view::FlatStorageChunkView;
3636
pub use delta::{FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata};
3737
pub use manager::FlatStorageManager;
3838
pub use storage::FlatStorage;
39-
pub use types::{BlockInfo, FetchingStateStatus, FlatStorageCreationStatus, FlatStorageError};
39+
pub use types::{
40+
BlockInfo, FetchingStateStatus, FlatStorageCreationStatus, FlatStorageError,
41+
FlatStorageReadyStatus, FlatStorageStatus,
42+
};
4043

4144
pub(crate) const POISONED_LOCK_ERR: &str = "The lock was poisoned.";
4245

0 commit comments

Comments
 (0)