From 99aeabe6bb1f2d87880e466f92c84942d26aab49 Mon Sep 17 00:00:00 2001 From: ocnc <97242826+ocnc@users.noreply.github.com> Date: Wed, 4 Sep 2024 15:29:42 -0400 Subject: [PATCH] feat(beradb): storage rw (#1992) Co-authored-by: berabera --- examples/berad/pkg/storage/history.go | 16 +++++++------- examples/berad/pkg/storage/kvstore.go | 27 ++++++++++++++---------- examples/berad/pkg/storage/randao.go | 4 ++-- examples/berad/pkg/storage/registry.go | 18 ++++++++-------- examples/berad/pkg/storage/staking.go | 20 ++++++++++++++++-- examples/berad/pkg/storage/versioning.go | 12 +++++------ 6 files changed, 59 insertions(+), 38 deletions(-) diff --git a/examples/berad/pkg/storage/history.go b/examples/berad/pkg/storage/history.go index 8606144c92..b85dcf6ad1 100644 --- a/examples/berad/pkg/storage/history.go +++ b/examples/berad/pkg/storage/history.go @@ -25,7 +25,7 @@ import "github.com/berachain/beacon-kit/mod/primitives/pkg/common" // UpdateBlockRootAtIndex sets a block root in the BeaconStore. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) UpdateBlockRootAtIndex( index uint64, root common.Root, @@ -36,7 +36,7 @@ func (kv *KVStore[ // GetBlockRootAtIndex retrieves the block root from the BeaconStore. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetBlockRootAtIndex( index uint64, ) (common.Root, error) { @@ -50,7 +50,7 @@ func (kv *KVStore[ // SetLatestBlockHeader sets the latest block header in the BeaconStore. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) SetLatestBlockHeader( header BeaconBlockHeaderT, ) error { @@ -60,7 +60,7 @@ func (kv *KVStore[ // GetLatestBlockHeader retrieves the latest block header from the BeaconStore. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetLatestBlockHeader() ( BeaconBlockHeaderT, error, ) { @@ -70,7 +70,7 @@ func (kv *KVStore[ // UpdateStateRootAtIndex updates the state root at the given slot. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) UpdateStateRootAtIndex( idx uint64, stateRoot common.Root, @@ -81,7 +81,7 @@ func (kv *KVStore[ // StateRootAtIndex returns the state root at the given slot. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) StateRootAtIndex( idx uint64, ) (common.Root, error) { @@ -96,7 +96,7 @@ func (kv *KVStore[ // header from the BeaconStore. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetLatestExecutionPayloadHeader() ( ExecutionPayloadHeaderT, error, ) { @@ -113,7 +113,7 @@ func (kv *KVStore[ // the BeaconStore. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) SetLatestExecutionPayloadHeader( payloadHeader ExecutionPayloadHeaderT, ) error { diff --git a/examples/berad/pkg/storage/kvstore.go b/examples/berad/pkg/storage/kvstore.go index 96651a6da1..8aafa03caf 100644 --- a/examples/berad/pkg/storage/kvstore.go +++ b/examples/berad/pkg/storage/kvstore.go @@ -51,6 +51,7 @@ type KVStore[ ValidatorT Validator[ValidatorT], ValidatorsT ~[]ValidatorT, WithdrawalT any, + WithdrawalsT ~[]WithdrawalT, ] struct { ctx context.Context // Versioning @@ -90,7 +91,7 @@ type KVStore[ randaoMix sdkcollections.Map[uint64, []byte] // Staking // withdrawals stores a list of initiated withdrawals. - withdrawals sdkcollections.Map[uint64, WithdrawalT] + withdrawals sdkcollections.Item[WithdrawalsT] } // New creates a new instance of Store. @@ -120,17 +121,22 @@ func New[ constraints.Empty[WithdrawalT] constraints.SSZMarshallable }, + WithdrawalsT interface { + ~[]WithdrawalT + constraints.Empty[WithdrawalsT] + constraints.SSZMarshallable + }, ]( kss store.KVStoreService, payloadCodec *encoding.SSZInterfaceCodec[ExecutionPayloadHeaderT], ) *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ] { schemaBuilder := sdkcollections.NewSchemaBuilder(kss) return &KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]{ ctx: nil, genesisValidatorsRoot: sdkcollections.NewItem( @@ -216,12 +222,11 @@ func New[ keys.LatestBeaconBlockHeaderPrefixHumanReadable, encoding.SSZValueCodec[BeaconBlockHeaderT]{}, ), - withdrawals: sdkcollections.NewMap( + withdrawals: sdkcollections.NewItem( schemaBuilder, sdkcollections.NewPrefix([]byte{keys.WithdrawalsPrefix}), keys.WithdrawalsPrefixHumanReadable, - sdkcollections.Uint64Key, - encoding.SSZValueCodec[WithdrawalT]{}, + encoding.SSZValueCodec[WithdrawalsT]{}, ), } } @@ -229,10 +234,10 @@ func New[ // Copy returns a copy of the Store. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) Copy() *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ] { // TODO: Decouple the KVStore type from the Cosmos-SDK. cctx, _ := sdk.UnwrapSDKContext(kv.ctx).CacheContext() @@ -243,7 +248,7 @@ func (kv *KVStore[ // Context returns the context of the Store. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) Context() context.Context { return kv.ctx } @@ -251,12 +256,12 @@ func (kv *KVStore[ // WithContext returns a copy of the Store with the given context. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) WithContext( ctx context.Context, ) *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ] { cpy := *kv cpy.ctx = ctx diff --git a/examples/berad/pkg/storage/randao.go b/examples/berad/pkg/storage/randao.go index 9f31a2a3be..43e5a443ea 100644 --- a/examples/berad/pkg/storage/randao.go +++ b/examples/berad/pkg/storage/randao.go @@ -25,7 +25,7 @@ import "github.com/berachain/beacon-kit/mod/primitives/pkg/common" // UpdateRandaoMixAtIndex sets the current RANDAO mix in the store. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) UpdateRandaoMixAtIndex( index uint64, mix common.Bytes32, @@ -36,7 +36,7 @@ func (kv *KVStore[ // GetRandaoMixAtIndex retrieves the current RANDAO mix from the store. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetRandaoMixAtIndex( index uint64, ) (common.Bytes32, error) { diff --git a/examples/berad/pkg/storage/registry.go b/examples/berad/pkg/storage/registry.go index d84bbaf66b..ed2db2451f 100644 --- a/examples/berad/pkg/storage/registry.go +++ b/examples/berad/pkg/storage/registry.go @@ -29,7 +29,7 @@ import ( // AddValidator registers a new validator in the beacon state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) AddValidator(val ValidatorT) error { // Get the next validator index from the sequence. idx, err := kv.validatorIndex.Next(kv.ctx) @@ -44,7 +44,7 @@ func (kv *KVStore[ // UpdateValidatorAtIndex updates a validator at a specific index. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) UpdateValidatorAtIndex( index math.ValidatorIndex, val ValidatorT, @@ -55,7 +55,7 @@ func (kv *KVStore[ // ValidatorIndexByPubkey returns the validator address by index. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) ValidatorIndexByPubkey( pubkey crypto.BLSPubkey, ) (math.ValidatorIndex, error) { @@ -72,7 +72,7 @@ func (kv *KVStore[ // ValidatorIndexByCometBFTAddress returns the validator address by index. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) ValidatorIndexByCometBFTAddress( cometBFTAddress []byte, ) (math.ValidatorIndex, error) { @@ -89,7 +89,7 @@ func (kv *KVStore[ // ValidatorByIndex returns the validator address by index. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) ValidatorByIndex( index math.ValidatorIndex, ) (ValidatorT, error) { @@ -104,7 +104,7 @@ func (kv *KVStore[ // GetValidators retrieves all validators from the beacon state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetValidators() ( ValidatorsT, error, ) { @@ -140,7 +140,7 @@ func (kv *KVStore[ // GetTotalValidators returns the total number of validators. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetTotalValidators() (uint64, error) { validators, err := kv.GetValidators() if err != nil { @@ -153,7 +153,7 @@ func (kv *KVStore[ // effective balance from the beacon state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetValidatorsByEffectiveBalance() ( []ValidatorT, error, ) { @@ -190,7 +190,7 @@ func (kv *KVStore[ // TODO: this shouldn't live in KVStore func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetTotalActiveBalances( slotsPerEpoch uint64, ) (math.Gwei, error) { diff --git a/examples/berad/pkg/storage/staking.go b/examples/berad/pkg/storage/staking.go index ba32d4d3b4..56f9b01284 100644 --- a/examples/berad/pkg/storage/staking.go +++ b/examples/berad/pkg/storage/staking.go @@ -23,7 +23,7 @@ package beacondb // GetEth1DepositIndex retrieves the eth1 deposit index from the beacon state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetEth1DepositIndex() (uint64, error) { return kv.eth1DepositIndex.Get(kv.ctx) } @@ -31,9 +31,25 @@ func (kv *KVStore[ // SetEth1DepositIndex sets the eth1 deposit index in the beacon state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) SetEth1DepositIndex( index uint64, ) error { return kv.eth1DepositIndex.Set(kv.ctx, index) } + +// GetWithdrawals retrieves the pending withdrawals from the beacon state. +func (kv *KVStore[ + BeaconBlockHeaderT, ExecutionPayloadHeaderT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, +]) GetWithdrawals() (WithdrawalsT, error) { + return kv.withdrawals.Get(kv.ctx) +} + +// SetWithdrawals overwrites the pending withdrawals from the beacon state. +func (kv *KVStore[ + BeaconBlockHeaderT, ExecutionPayloadHeaderT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, +]) SetWithdrawals(withdrawals WithdrawalsT) error { + return kv.withdrawals.Set(kv.ctx, withdrawals) +} diff --git a/examples/berad/pkg/storage/versioning.go b/examples/berad/pkg/storage/versioning.go index 276272a350..5d13b06522 100644 --- a/examples/berad/pkg/storage/versioning.go +++ b/examples/berad/pkg/storage/versioning.go @@ -29,7 +29,7 @@ import ( // state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) SetGenesisValidatorsRoot( root common.Root, ) error { @@ -40,7 +40,7 @@ func (kv *KVStore[ // beacon state. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetGenesisValidatorsRoot() (common.Root, error) { bz, err := kv.genesisValidatorsRoot.Get(kv.ctx) if err != nil { @@ -52,7 +52,7 @@ func (kv *KVStore[ // GetSlot returns the current slot. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetSlot() (math.Slot, error) { slot, err := kv.slot.Get(kv.ctx) return math.Slot(slot), err @@ -61,7 +61,7 @@ func (kv *KVStore[ // SetSlot sets the current slot. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) SetSlot( slot math.Slot, ) error { @@ -71,7 +71,7 @@ func (kv *KVStore[ // SetFork sets the fork version for the given epoch. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) SetFork( fork ForkT, ) error { @@ -81,7 +81,7 @@ func (kv *KVStore[ // GetFork gets the fork version for the given epoch. func (kv *KVStore[ BeaconBlockHeaderT, ExecutionPayloadHeaderT, - ForkT, ValidatorT, ValidatorsT, WithdrawalT, + ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT, ]) GetFork() (ForkT, error) { return kv.fork.Get(kv.ctx) }