Skip to content

Commit

Permalink
Add electra DB
Browse files Browse the repository at this point in the history
  • Loading branch information
@terencechain
terencechain committed May 9, 2024
1 parent 41edee9 commit a5711b4
Show file tree
Hide file tree
Showing 7 changed files with 314 additions and 109 deletions.
15 changes: 15 additions & 0 deletions beacon-chain/db/kv/blocks.go
View file
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
View file
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
View file
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) {
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
View file
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
View file
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 {
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

0 comments on commit a5711b4

Please sign in to comment.