Skip to content

Commit

Permalink
Add ePBS to db
Browse files Browse the repository at this point in the history
  • Loading branch information
terencechain committed May 8, 2024
1 parent ecc2b29 commit 49d6287
Show file tree
Hide file tree
Showing 22 changed files with 2,252 additions and 19 deletions.
2 changes: 2 additions & 0 deletions beacon-chain/db/iface/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ type ReadOnlyDatabase interface {
DepositContractAddress(ctx context.Context) ([]byte, error)
// ExecutionChainData operations.
ExecutionChainData(ctx context.Context) (*ethpb.ETH1ChainData, error)
SignedExecutionPayloadEnvelopeBlind(ctx context.Context, blockRoot []byte) (*ethpb.SignedExecutionPayloadEnvelopeBlind, error)
// Fee recipients operations.
FeeRecipientByValidatorID(ctx context.Context, id primitives.ValidatorIndex) (common.Address, error)
RegistrationByValidatorID(ctx context.Context, id primitives.ValidatorIndex) (*ethpb.ValidatorRegistrationV1, error)
Expand Down Expand Up @@ -87,6 +88,7 @@ type NoHeadAccessDatabase interface {
SaveDepositContractAddress(ctx context.Context, addr common.Address) error
// SaveExecutionChainData operations.
SaveExecutionChainData(ctx context.Context, data *ethpb.ETH1ChainData) error
SaveSignedExecutionPayloadEnvelopeBlind(ctx context.Context, envelope *ethpb.SignedExecutionPayloadEnvelopeBlind) error
// Run any required database migrations.
RunMigrations(ctx context.Context) error
// Fee recipients operations.
Expand Down
3 changes: 3 additions & 0 deletions beacon-chain/db/kv/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ go_library(
"archived_point.go",
"backfill.go",
"backup.go",
"blind_payload_envelope.go",
"blocks.go",
"checkpoint.go",
"deposit_contract.go",
Expand Down Expand Up @@ -78,6 +79,7 @@ go_test(
"archived_point_test.go",
"backfill_test.go",
"backup_test.go",
"blind_payload_envelope_test.go",
"blocks_test.go",
"checkpoint_test.go",
"deposit_contract_test.go",
Expand Down Expand Up @@ -119,6 +121,7 @@ go_test(
"//testing/assert:go_default_library",
"//testing/require:go_default_library",
"//testing/util:go_default_library",
"//testing/util/random:go_default_library",
"@com_github_ethereum_go_ethereum//common:go_default_library",
"@com_github_golang_snappy//:go_default_library",
"@com_github_pkg_errors//:go_default_library",
Expand Down
45 changes: 45 additions & 0 deletions beacon-chain/db/kv/blind_payload_envelope.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package kv

import (
"context"

ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
bolt "go.etcd.io/bbolt"
"go.opencensus.io/trace"
)

// SaveSignedExecutionPayloadEnvelopeBlind saves a signed execution payload envelope blind in the database.
func (s *Store) SaveSignedExecutionPayloadEnvelopeBlind(ctx context.Context, env *ethpb.SignedExecutionPayloadEnvelopeBlind) error {
ctx, span := trace.StartSpan(ctx, "BeaconDB.SaveSignedExecutionPayloadEnvelopeBlind")
defer span.End()

enc, err := encode(ctx, env)
if err != nil {
return err
}

r := env.Message.BeaconBlockRoot
err = s.db.Update(func(tx *bolt.Tx) error {
bucket := tx.Bucket(executionPayloadEnvelopeBucket)
return bucket.Put(r, enc)
})

return err
}

// SignedExecutionPayloadEnvelopeBlind retrieves a signed execution payload envelope blind from the database.
func (s *Store) SignedExecutionPayloadEnvelopeBlind(ctx context.Context, blockRoot []byte) (*ethpb.SignedExecutionPayloadEnvelopeBlind, error) {
ctx, span := trace.StartSpan(ctx, "BeaconDB.SignedExecutionPayloadEnvelopeBlind")
defer span.End()

env := &ethpb.SignedExecutionPayloadEnvelopeBlind{}
err := s.db.View(func(tx *bolt.Tx) error {
bkt := tx.Bucket(executionPayloadEnvelopeBucket)
enc := bkt.Get(blockRoot)
if enc == nil {
return ErrNotFound
}
return decode(ctx, enc, env)
})
return env, err
}
23 changes: 23 additions & 0 deletions beacon-chain/db/kv/blind_payload_envelope_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
package kv

import (
"context"
"testing"

"github.com/prysmaticlabs/prysm/v5/testing/require"
"github.com/prysmaticlabs/prysm/v5/testing/util/random"
)

func TestStore_SignedExecutionPayloadEnvelopeBlind(t *testing.T) {
db := setupDB(t)
ctx := context.Background()
_, err := db.SignedExecutionPayloadEnvelopeBlind(ctx, []byte("test"))
require.ErrorIs(t, err, ErrNotFound)

env := random.SignedExecutionPayloadEnvelopeBlind(t)
err = db.SaveSignedExecutionPayloadEnvelopeBlind(ctx, env)
require.NoError(t, err)
got, err := db.SignedExecutionPayloadEnvelopeBlind(ctx, env.Message.BeaconBlockRoot)
require.NoError(t, err)
require.DeepEqual(t, got, env)
}
9 changes: 8 additions & 1 deletion beacon-chain/db/kv/blocks.go
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,7 @@ func prepareBlockBatch(blks []blocks.ROBlock, shouldBlind bool) ([]blockBatchEnt
if !errors.Is(err, blocks.ErrUnsupportedVersion) {
return nil, errors.Wrapf(err, "could not convert block to blinded format for root %#x", batch[i].root)
}
// Pre-deneb blocks give ErrUnsupportedVersion; use the full block already in the batch entry.
// Pre-bellatrix and ePBS blocks give ErrUnsupportedVersion; use the full block already in the batch entry.
} else {
batch[i].block = blinded
}
Expand Down Expand Up @@ -813,6 +813,11 @@ func unmarshalBlock(_ context.Context, enc []byte) (interfaces.ReadOnlySignedBea
if err := rawBlock.UnmarshalSSZ(enc[len(denebBlindKey):]); err != nil {
return nil, errors.Wrap(err, "could not unmarshal blinded Deneb block")
}
case hasEpbsKey(enc):
rawBlock = &ethpb.SignedBeaconBlockEpbs{}
if err := rawBlock.UnmarshalSSZ(enc[len(epbsKey):]); err != nil {
return nil, errors.Wrap(err, "could not unmarshal EPBS block")
}
default:
// Marshal block bytes to phase 0 beacon block.
rawBlock = &ethpb.SignedBeaconBlock{}
Expand Down Expand Up @@ -842,6 +847,8 @@ func encodeBlock(blk interfaces.ReadOnlySignedBeaconBlock) ([]byte, error) {

func keyForBlock(blk interfaces.ReadOnlySignedBeaconBlock) ([]byte, error) {
switch blk.Version() {
case version.EPBS:
return epbsKey, nil
case version.Deneb:
if blk.IsBlinded() {
return denebBlindKey, nil
Expand Down
12 changes: 12 additions & 0 deletions beacon-chain/db/kv/blocks_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"github.com/prysmaticlabs/prysm/v5/testing/assert"
"github.com/prysmaticlabs/prysm/v5/testing/require"
"github.com/prysmaticlabs/prysm/v5/testing/util"
"github.com/prysmaticlabs/prysm/v5/testing/util/random"
"google.golang.org/protobuf/proto"
)

Expand Down Expand Up @@ -125,6 +126,17 @@ var blockTests = []struct {
return blocks.NewSignedBeaconBlock(b)
},
},
{
name: "epbs",
newBlock: func(slot primitives.Slot, root []byte) (interfaces.ReadOnlySignedBeaconBlock, error) {
b := random.SignedBeaconBlock(&testing.T{})
b.Block.Slot = slot
if root != nil {
b.Block.ParentRoot = root
}
return blocks.NewSignedBeaconBlock(b)
},
},
}

func TestStore_SaveBlock_NoDuplicates(t *testing.T) {
Expand Down
2 changes: 2 additions & 0 deletions beacon-chain/db/kv/encoding.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ func isSSZStorageFormat(obj interface{}) bool {
return true
case *ethpb.VoluntaryExit:
return true
case *ethpb.SignedExecutionPayloadEnvelopeBlind:
return true
case *ethpb.ValidatorRegistrationV1:
return true
default:
Expand Down
7 changes: 7 additions & 0 deletions beacon-chain/db/kv/key.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,10 @@ func hasDenebBlindKey(enc []byte) bool {
}
return bytes.Equal(enc[:len(denebBlindKey)], denebBlindKey)
}

func hasEpbsKey(enc []byte) bool {
if len(epbsKey) >= len(enc) {
return false
}
return bytes.Equal(enc[:len(epbsKey)], epbsKey)
}
3 changes: 3 additions & 0 deletions beacon-chain/db/kv/kv.go
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,9 @@ var Buckets = [][]byte{

feeRecipientBucket,
registrationBucket,

// ePBS
executionPayloadEnvelopeBucket,
}

// KVStoreOption is a functional option that modifies a kv.Store.
Expand Down
20 changes: 11 additions & 9 deletions beacon-chain/db/kv/schema.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,16 @@ package kv
// it easy to scan for keys that have a certain shard number as a prefix and return those
// corresponding attestations.
var (
blocksBucket = []byte("blocks")
stateBucket = []byte("state")
stateSummaryBucket = []byte("state-summary")
chainMetadataBucket = []byte("chain-metadata")
checkpointBucket = []byte("check-point")
powchainBucket = []byte("powchain")
stateValidatorsBucket = []byte("state-validators")
feeRecipientBucket = []byte("fee-recipient")
registrationBucket = []byte("registration")
blocksBucket = []byte("blocks")
stateBucket = []byte("state")
stateSummaryBucket = []byte("state-summary")
chainMetadataBucket = []byte("chain-metadata")
checkpointBucket = []byte("check-point")
powchainBucket = []byte("powchain")
stateValidatorsBucket = []byte("state-validators")
feeRecipientBucket = []byte("fee-recipient")
registrationBucket = []byte("registration")
executionPayloadEnvelopeBucket = []byte("execution-payload-envelope")

// Deprecated: This bucket was migrated in PR 6461. Do not use, except for migrations.
slotsHasObjectBucket = []byte("slots-has-objects")
Expand Down Expand Up @@ -48,6 +49,7 @@ var (
saveBlindedBeaconBlocksKey = []byte("save-blinded-beacon-blocks")
denebKey = []byte("deneb")
denebBlindKey = []byte("blind-deneb")
epbsKey = []byte("epbs")

// block root included in the beacon state used by weak subjectivity initial sync
originCheckpointBlockRootKey = []byte("origin-checkpoint-block-root")
Expand Down
56 changes: 56 additions & 0 deletions beacon-chain/db/kv/state.go
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,25 @@ func (s *Store) saveStatesEfficientInternal(ctx context.Context, tx *bolt.Tx, bl
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
return err
}
case *ethpb.BeaconStateEPBS:
pbState, err := getEpbsPbState(rawType)
if err != nil {
return err
}
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
rawObj, err := pbState.MarshalSSZ()
if err != nil {
return err
}
encodedState := snappy.Encode(nil, append(epbsKey, rawObj...))
if err := bucket.Put(rt[:], encodedState); err != nil {
return err
}
pbState.Validators = valEntries
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
return err
}
default:
return errors.New("invalid state type")
}
Expand Down Expand Up @@ -385,6 +404,17 @@ func getDenebPbState(rawState interface{}) (*ethpb.BeaconStateDeneb, error) {
return pbState, nil
}

func getEpbsPbState(rawState interface{}) (*ethpb.BeaconStateEPBS, error) {
pbState, ok := rawState.(*ethpb.BeaconStateEPBS)
if !ok {
return nil, errors.New("input is not type pb.BeaconStateElectra")
}
if pbState == nil {
return nil, errors.New("nil state")
}
return pbState, nil
}

func (s *Store) storeValidatorEntriesSeparately(ctx context.Context, tx *bolt.Tx, validatorsEntries map[string]*ethpb.Validator) error {
valBkt := tx.Bucket(stateValidatorsBucket)
for hashStr, validatorEntry := range validatorsEntries {
Expand Down Expand Up @@ -534,6 +564,19 @@ func (s *Store) unmarshalState(_ context.Context, enc []byte, validatorEntries [
}

switch {
case hasEpbsKey(enc):
protoState := &ethpb.BeaconStateEPBS{}
if err := protoState.UnmarshalSSZ(enc[len(epbsKey):]); err != nil {
return nil, errors.Wrap(err, "failed to unmarshal encoding for EPBS")
}
ok, err := s.isStateValidatorMigrationOver()
if err != nil {
return nil, err
}
if ok {
protoState.Validators = validatorEntries
}
return statenative.InitializeFromProtoUnsafeEpbs(protoState)
case hasDenebKey(enc):
protoState := &ethpb.BeaconStateDeneb{}
if err := protoState.UnmarshalSSZ(enc[len(denebKey):]); err != nil {
Expand Down Expand Up @@ -667,6 +710,19 @@ func marshalState(ctx context.Context, st state.ReadOnlyBeaconState) ([]byte, er
return nil, err
}
return snappy.Encode(nil, append(denebKey, rawObj...)), nil
case *ethpb.BeaconStateEPBS:
rState, ok := st.ToProtoUnsafe().(*ethpb.BeaconStateEPBS)
if !ok {
return nil, errors.New("non valid inner state")
}
if rState == nil {
return nil, errors.New("nil state")
}
rawObj, err := rState.MarshalSSZ()
if err != nil {
return nil, err
}
return snappy.Encode(nil, append(epbsKey, rawObj...)), nil
default:
return nil, errors.New("invalid inner state")
}
Expand Down
32 changes: 31 additions & 1 deletion beacon-chain/db/kv/state_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"time"

"github.com/prysmaticlabs/prysm/v5/beacon-chain/state"
statenative "github.com/prysmaticlabs/prysm/v5/beacon-chain/state/state-native"
"github.com/prysmaticlabs/prysm/v5/config/features"
"github.com/prysmaticlabs/prysm/v5/config/params"
"github.com/prysmaticlabs/prysm/v5/consensus-types/blocks"
Expand All @@ -21,6 +22,7 @@ import (
"github.com/prysmaticlabs/prysm/v5/testing/assert"
"github.com/prysmaticlabs/prysm/v5/testing/require"
"github.com/prysmaticlabs/prysm/v5/testing/util"
"github.com/prysmaticlabs/prysm/v5/testing/util/random"
bolt "go.etcd.io/bbolt"
)

Expand Down Expand Up @@ -132,6 +134,16 @@ func TestState_CanSaveRetrieve(t *testing.T) {
},
rootSeed: 'D',
},
{
name: "epbs",
s: func() state.BeaconState {
stPb := random.BeaconState(t)
st, err := statenative.InitializeFromProtoUnsafeEpbs(stPb)
require.NoError(t, err)
return st
},
rootSeed: 'E',
},
}

db := setupDB(t)
Expand All @@ -155,7 +167,7 @@ func TestState_CanSaveRetrieve(t *testing.T) {
savedSt, err := db.State(context.Background(), root)
require.NoError(t, err)

assert.DeepSSZEqual(t, st.ToProtoUnsafe(), savedSt.ToProtoUnsafe())
require.DeepEqual(t, st.ToProtoUnsafe(), savedSt.ToProtoUnsafe())
})
}

Expand Down Expand Up @@ -1086,6 +1098,24 @@ func TestDenebState_CanDelete(t *testing.T) {
require.Equal(t, state.ReadOnlyBeaconState(nil), savedS, "Unsaved state should've been nil")
}

func TestEpbsState_CanDelete(t *testing.T) {
db := setupDB(t)

r := [32]byte{'A'}
s := random.BeaconState(t)
st, err := statenative.InitializeFromProtoUnsafeEpbs(s)
require.NoError(t, err)
require.NoError(t, st.SetSlot(100))

require.NoError(t, db.SaveState(context.Background(), st, r))
require.Equal(t, true, db.HasState(context.Background(), r))

require.NoError(t, db.DeleteState(context.Background(), r))
savedS, err := db.State(context.Background(), r)
require.NoError(t, err)
require.Equal(t, state.ReadOnlyBeaconState(nil), savedS, "Unsaved state should've been nil")
}

func TestStateDeneb_CanSaveRetrieveValidatorEntries(t *testing.T) {
db := setupDB(t)

Expand Down
2 changes: 1 addition & 1 deletion proto/engine/v1/generated.ssz.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion proto/eth/v1/generated.ssz.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit 49d6287

Please sign in to comment.