Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add electra DB #13975

Merged
merged 4 commits into from
May 9, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions beacon-chain/db/kv/blocks.go
Original file line number Diff line number Diff line change
Expand Up @@ -813,6 +813,16 @@ 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 hasElectraKey(enc):
rawBlock = &ethpb.SignedBeaconBlockElectra{}
if err := rawBlock.UnmarshalSSZ(enc[len(electraKey):]); err != nil {
return nil, errors.Wrap(err, "could not unmarshal Electra block")
}
case hasElectraBlindKey(enc):
rawBlock = &ethpb.SignedBlindedBeaconBlockElectra{}
if err := rawBlock.UnmarshalSSZ(enc[len(electraBlindKey):]); err != nil {
return nil, errors.Wrap(err, "could not unmarshal blinded Electra block")
}
default:
// Marshal block bytes to phase 0 beacon block.
rawBlock = &ethpb.SignedBeaconBlock{}
Expand Down Expand Up @@ -842,6 +852,11 @@ func encodeBlock(blk interfaces.ReadOnlySignedBeaconBlock) ([]byte, error) {

func keyForBlock(blk interfaces.ReadOnlySignedBeaconBlock) ([]byte, error) {
switch blk.Version() {
case version.Electra:
if blk.IsBlinded() {
return electraBlindKey, nil
}
return electraKey, nil
case version.Deneb:
if blk.IsBlinded() {
return denebBlindKey, nil
Expand Down
41 changes: 41 additions & 0 deletions beacon-chain/db/kv/blocks_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,47 @@ var blockTests = []struct {
return blocks.NewSignedBeaconBlock(b)
},
},
{
name: "electra",
newBlock: func(slot primitives.Slot, root []byte) (interfaces.ReadOnlySignedBeaconBlock, error) {
b := util.NewBeaconBlockElectra()
b.Block.Slot = slot
if root != nil {
b.Block.ParentRoot = root
b.Block.Body.Consolidations = []*ethpb.SignedConsolidation{
{
Message: &ethpb.Consolidation{
SourceIndex: 1,
TargetIndex: 2,
Epoch: 3,
},
Signature: make([]byte, 96),
},
}
}
return blocks.NewSignedBeaconBlock(b)
},
},
{
name: "electra blind",
newBlock: func(slot primitives.Slot, root []byte) (interfaces.ReadOnlySignedBeaconBlock, error) {
b := util.NewBlindedBeaconBlockElectra()
b.Message.Slot = slot
if root != nil {
b.Message.ParentRoot = root
b.Message.Body.Consolidations = []*ethpb.SignedConsolidation{
{
Message: &ethpb.Consolidation{
SourceIndex: 1,
TargetIndex: 2,
Epoch: 3,
},
Signature: make([]byte, 96),
},
}
}
return blocks.NewSignedBeaconBlock(b)
}},
}

func TestStore_SaveBlock_NoDuplicates(t *testing.T) {
Expand Down
14 changes: 14 additions & 0 deletions beacon-chain/db/kv/key.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,17 @@ func hasDenebBlindKey(enc []byte) bool {
}
return bytes.Equal(enc[:len(denebBlindKey)], denebBlindKey)
}

func hasElectraKey(enc []byte) bool {
if len(electraKey) >= len(enc) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't really need to do this, it would be fine to just call bytes.Equal.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kasey it is required or a runtime error for slice bounds out of range will happen when len(enc) is less than len(electraKey)

Copy link
Member

@prestonvanloon prestonvanloon May 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are thinking of bytes.Prefix, i think. It would be safe to call bytes.HasPrefix without this check

https://pkg.go.dev/bytes?utm_source=godoc#HasPrefix

func hasElectraKey(enc []byte) bool {
  return bytes.hasPrefix(enc, electraKey)
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right, this is basically replicating bytes.HasPrefix.

return false
}
return bytes.Equal(enc[:len(electraKey)], electraKey)
}

func hasElectraBlindKey(enc []byte) bool {
if len(electraBlindKey) >= len(enc) {
return false
}
return bytes.Equal(enc[:len(electraBlindKey)], electraBlindKey)
}
2 changes: 2 additions & 0 deletions beacon-chain/db/kv/schema.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ var (
saveBlindedBeaconBlocksKey = []byte("save-blinded-beacon-blocks")
denebKey = []byte("deneb")
denebBlindKey = []byte("blind-deneb")
electraKey = []byte("electra")
electraBlindKey = []byte("blind-electra")

// block root included in the beacon state used by weak subjectivity initial sync
originCheckpointBlockRootKey = []byte("origin-checkpoint-block-root")
Expand Down
224 changes: 116 additions & 108 deletions beacon-chain/db/kv/state.go
Original file line number Diff line number Diff line change
Expand Up @@ -229,160 +229,142 @@ func (s *Store) saveStatesEfficientInternal(ctx context.Context, tx *bolt.Tx, bl
// look at issue https://github.com/prysmaticlabs/prysm/issues/9262.
switch rawType := states[i].ToProtoUnsafe().(type) {
case *ethpb.BeaconState:
pbState, err := getPhase0PbState(rawType)
if err != nil {
return err
}
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
encodedState, err := encode(ctx, pbState)
if err != nil {
return err
}
pbState.Validators = valEntries
if err := bucket.Put(rt[:], encodedState); err != nil {
return err
}
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
if err := s.processPhase0(ctx, rawType, rt[:], bucket, valIdxBkt, validatorKeys[i]); err != nil {
return err
}
case *ethpb.BeaconStateAltair:
pbState, err := getAltairPbState(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(altairKey, rawObj...))
if err := bucket.Put(rt[:], encodedState); err != nil {
return err
}
pbState.Validators = valEntries
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
if err := s.processAltair(ctx, rawType, rt[:], bucket, valIdxBkt, validatorKeys[i]); err != nil {
return err
}
case *ethpb.BeaconStateBellatrix:
pbState, err := getBellatrixPbState(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(bellatrixKey, rawObj...))
if err := bucket.Put(rt[:], encodedState); err != nil {
return err
}
pbState.Validators = valEntries
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
if err := s.processBellatrix(ctx, rawType, rt[:], bucket, valIdxBkt, validatorKeys[i]); err != nil {
return err
}
case *ethpb.BeaconStateCapella:
pbState, err := getCapellaPbState(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(capellaKey, rawObj...))
if err := bucket.Put(rt[:], encodedState); err != nil {
return err
}
pbState.Validators = valEntries
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
if err := s.processCapella(ctx, rawType, rt[:], bucket, valIdxBkt, validatorKeys[i]); err != nil {
return err
}
case *ethpb.BeaconStateDeneb:
pbState, err := getDenebPbState(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(denebKey, rawObj...))
if err := bucket.Put(rt[:], encodedState); err != nil {
if err := s.processDeneb(ctx, rawType, rt[:], bucket, valIdxBkt, validatorKeys[i]); err != nil {
return err
}
pbState.Validators = valEntries
if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
case *ethpb.BeaconStateElectra:
if err := s.processEPBS(ctx, rawType, rt[:], bucket, valIdxBkt, validatorKeys[i]); err != nil {
return err
}
default:
return errors.New("invalid state type")
}
}
// store the validator entries separately to save space.

return s.storeValidatorEntriesSeparately(ctx, tx, validatorsEntries)
}

func getPhase0PbState(rawState interface{}) (*ethpb.BeaconState, error) {
pbState, err := statenative.ProtobufBeaconStatePhase0(rawState)
func (s *Store) processPhase0(ctx context.Context, pbState *ethpb.BeaconState, rootHash []byte, bucket, valIdxBkt *bolt.Bucket, validatorKey []byte) error {
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
encodedState, err := encode(ctx, pbState)
if err != nil {
return nil, err
return err
}
pbState.Validators = valEntries
if err := bucket.Put(rootHash, encodedState); err != nil {
return err
}
if pbState == nil {
return nil, errors.New("nil state")
if err := valIdxBkt.Put(rootHash, validatorKey); err != nil {
return err
}
return pbState, nil
return nil
}

func getAltairPbState(rawState interface{}) (*ethpb.BeaconStateAltair, error) {
pbState, err := statenative.ProtobufBeaconStateAltair(rawState)
func (s *Store) processAltair(ctx context.Context, pbState *ethpb.BeaconStateAltair, rootHash []byte, bucket, valIdxBkt *bolt.Bucket, validatorKey []byte) error {
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
rawObj, err := pbState.MarshalSSZ()
if err != nil {
return nil, err
return err
}
if pbState == nil {
return nil, errors.New("nil state")
encodedState := snappy.Encode(nil, append(altairKey, rawObj...))
if err := bucket.Put(rootHash, encodedState); err != nil {
return err
}
return pbState, nil
pbState.Validators = valEntries
if err := valIdxBkt.Put(rootHash, validatorKey); err != nil {
return err
}
return nil
}

func getBellatrixPbState(rawState interface{}) (*ethpb.BeaconStateBellatrix, error) {
pbState, err := statenative.ProtobufBeaconStateBellatrix(rawState)
func (s *Store) processBellatrix(ctx context.Context, pbState *ethpb.BeaconStateBellatrix, rootHash []byte, bucket, valIdxBkt *bolt.Bucket, validatorKey []byte) error {
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
rawObj, err := pbState.MarshalSSZ()
if err != nil {
return nil, err
return err
}
encodedState := snappy.Encode(nil, append(bellatrixKey, rawObj...))
if err := bucket.Put(rootHash, encodedState); err != nil {
return err
}
if pbState == nil {
return nil, errors.New("nil state")
pbState.Validators = valEntries
if err := valIdxBkt.Put(rootHash, validatorKey); err != nil {
return err
}
return pbState, nil
return nil
}

func getCapellaPbState(rawState interface{}) (*ethpb.BeaconStateCapella, error) {
pbState, err := statenative.ProtobufBeaconStateCapella(rawState)
func (s *Store) processCapella(ctx context.Context, pbState *ethpb.BeaconStateCapella, rootHash []byte, bucket, valIdxBkt *bolt.Bucket, validatorKey []byte) error {
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
rawObj, err := pbState.MarshalSSZ()
if err != nil {
return nil, err
return err
}
encodedState := snappy.Encode(nil, append(capellaKey, rawObj...))
if err := bucket.Put(rootHash, encodedState); err != nil {
return err
}
if pbState == nil {
return nil, errors.New("nil state")
pbState.Validators = valEntries
if err := valIdxBkt.Put(rootHash, validatorKey); err != nil {
return err
}
return pbState, nil
return nil
}

func getDenebPbState(rawState interface{}) (*ethpb.BeaconStateDeneb, error) {
pbState, err := statenative.ProtobufBeaconStateDeneb(rawState)
func (s *Store) processDeneb(ctx context.Context, pbState *ethpb.BeaconStateDeneb, rootHash []byte, bucket, valIdxBkt *bolt.Bucket, validatorKey []byte) error {
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
rawObj, err := pbState.MarshalSSZ()
if err != nil {
return nil, err
return err
}
encodedState := snappy.Encode(nil, append(denebKey, rawObj...))
if err := bucket.Put(rootHash, encodedState); err != nil {
return err
}
pbState.Validators = valEntries
if err := valIdxBkt.Put(rootHash, validatorKey); err != nil {
return err
}
return nil
}

func (s *Store) processEPBS(ctx context.Context, pbState *ethpb.BeaconStateElectra, rootHash []byte, bucket, valIdxBkt *bolt.Bucket, validatorKey []byte) error {
terencechain marked this conversation as resolved.
Show resolved Hide resolved
valEntries := pbState.Validators
pbState.Validators = make([]*ethpb.Validator, 0)
rawObj, err := pbState.MarshalSSZ()
if err != nil {
return err
}
encodedState := snappy.Encode(nil, append(electraKey, rawObj...))
if err := bucket.Put(rootHash, encodedState); err != nil {
return err
}
if pbState == nil {
return nil, errors.New("nil state")
pbState.Validators = valEntries
if err := valIdxBkt.Put(rootHash, validatorKey); err != nil {
return err
}
return pbState, nil
return nil
}

func (s *Store) storeValidatorEntriesSeparately(ctx context.Context, tx *bolt.Tx, validatorsEntries map[string]*ethpb.Validator) error {
Expand Down Expand Up @@ -534,6 +516,19 @@ func (s *Store) unmarshalState(_ context.Context, enc []byte, validatorEntries [
}

switch {
case hasElectraKey(enc):
protoState := &ethpb.BeaconStateElectra{}
if err := protoState.UnmarshalSSZ(enc[len(electraKey):]); err != nil {
return nil, errors.Wrap(err, "failed to unmarshal encoding for Electra")
}
ok, err := s.isStateValidatorMigrationOver()
if err != nil {
return nil, err
}
if ok {
protoState.Validators = validatorEntries
}
return statenative.InitializeFromProtoUnsafeElectra(protoState)
case hasDenebKey(enc):
protoState := &ethpb.BeaconStateDeneb{}
if err := protoState.UnmarshalSSZ(enc[len(denebKey):]); err != nil {
Expand Down Expand Up @@ -667,6 +662,19 @@ func marshalState(ctx context.Context, st state.ReadOnlyBeaconState) ([]byte, er
return nil, err
}
return snappy.Encode(nil, append(denebKey, rawObj...)), nil
case *ethpb.BeaconStateElectra:
rState, ok := st.ToProtoUnsafe().(*ethpb.BeaconStateElectra)
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(electraKey, rawObj...)), nil
default:
return nil, errors.New("invalid inner state")
}
Expand Down
Loading
Loading