From 3e36eacd1e1c756a5d38eeb79de297b4eebbf9b9 Mon Sep 17 00:00:00 2001 From: terencechain Date: Fri, 3 Jun 2022 08:13:04 -0700 Subject: [PATCH] Add registration DB methods (#10812) --- beacon-chain/db/kv/blocks.go | 56 ++++++++++++++++++++++- beacon-chain/db/kv/blocks_test.go | 74 +++++++++++++++++++++++++++++++ beacon-chain/db/kv/encoding.go | 2 + beacon-chain/db/kv/kv.go | 1 + beacon-chain/db/kv/schema.go | 1 + proto/prysm/v1alpha1/BUILD.bazel | 1 + 6 files changed, 134 insertions(+), 1 deletion(-) diff --git a/beacon-chain/db/kv/blocks.go b/beacon-chain/db/kv/blocks.go index 46a0ca462..bfdd632ec 100644 --- a/beacon-chain/db/kv/blocks.go +++ b/beacon-chain/db/kv/blocks.go @@ -469,8 +469,20 @@ func (s *Store) FeeRecipientByValidatorID(ctx context.Context, id types.Validato err := s.db.View(func(tx *bolt.Tx) error { bkt := tx.Bucket(feeRecipientBucket) addr = bkt.Get(bytesutil.Uint64ToBytesBigEndian(uint64(id))) + // IF the fee recipient is not found in the standard fee recipient bucket, then + // check the registration bucket. The fee recipient may be there. + // This is to resolve imcompatility until we fully migrate to the registration bucket. if addr == nil { - return errors.Wrapf(ErrNotFoundFeeRecipient, "validator id %d", id) + bkt = tx.Bucket(registrationBucket) + enc := bkt.Get(bytesutil.Uint64ToBytesBigEndian(uint64(id))) + if enc == nil { + return errors.Wrapf(ErrNotFoundFeeRecipient, "validator id %d", id) + } + reg := ðpb.ValidatorRegistrationV1{} + if err := decode(ctx, enc, reg); err != nil { + return err + } + addr = reg.FeeRecipient } return nil }) @@ -498,6 +510,48 @@ func (s *Store) SaveFeeRecipientsByValidatorIDs(ctx context.Context, ids []types }) } +// RegistrationByValidatorID returns the validator registration object for a validator id. +// `ErrNotFoundFeeRecipient` is returned if the validator id is not found. +func (s *Store) RegistrationByValidatorID(ctx context.Context, id types.ValidatorIndex) (*ethpb.ValidatorRegistrationV1, error) { + ctx, span := trace.StartSpan(ctx, "BeaconDB.RegistrationByValidatorID") + defer span.End() + reg := ðpb.ValidatorRegistrationV1{} + err := s.db.View(func(tx *bolt.Tx) error { + bkt := tx.Bucket(registrationBucket) + enc := bkt.Get(bytesutil.Uint64ToBytesBigEndian(uint64(id))) + if enc == nil { + return errors.Wrapf(ErrNotFoundFeeRecipient, "validator id %d", id) + } + return decode(ctx, enc, reg) + }) + return reg, err +} + +// SaveRegistrationsByValidatorIDs saves the validator registrations for validator ids. +// Error is returned if `ids` and `registrations` are not the same length. +func (s *Store) SaveRegistrationsByValidatorIDs(ctx context.Context, ids []types.ValidatorIndex, regs []*ethpb.ValidatorRegistrationV1) error { + _, span := trace.StartSpan(ctx, "BeaconDB.SaveRegistrationsByValidatorIDs") + defer span.End() + + if len(ids) != len(regs) { + return errors.New("ids and registrations must be the same length") + } + + return s.db.Update(func(tx *bolt.Tx) error { + bkt := tx.Bucket(registrationBucket) + for i, id := range ids { + enc, err := encode(ctx, regs[i]) + if err != nil { + return err + } + if err := bkt.Put(bytesutil.Uint64ToBytesBigEndian(uint64(id)), enc); err != nil { + return err + } + } + return nil + }) +} + // blockRootsByFilter retrieves the block roots given the filter criteria. func blockRootsByFilter(ctx context.Context, tx *bolt.Tx, f *filters.QueryFilter) ([][]byte, error) { ctx, span := trace.StartSpan(ctx, "BeaconDB.blockRootsByFilter") diff --git a/beacon-chain/db/kv/blocks_test.go b/beacon-chain/db/kv/blocks_test.go index aca6dcf75..d10c205fb 100644 --- a/beacon-chain/db/kv/blocks_test.go +++ b/beacon-chain/db/kv/blocks_test.go @@ -693,4 +693,78 @@ func TestStore_FeeRecipientByValidatorID(t *testing.T) { _, err = db.FeeRecipientByValidatorID(ctx, 3) want := errors.Wrap(ErrNotFoundFeeRecipient, "validator id 3") require.Equal(t, want.Error(), err.Error()) + + regs := []*ethpb.ValidatorRegistrationV1{ + { + FeeRecipient: bytesutil.PadTo([]byte("a"), 20), + GasLimit: 1, + Timestamp: 2, + Pubkey: bytesutil.PadTo([]byte("b"), 48), + }} + require.NoError(t, db.SaveRegistrationsByValidatorIDs(ctx, []types.ValidatorIndex{3}, regs)) + f, err = db.FeeRecipientByValidatorID(ctx, 3) + require.NoError(t, err) + require.Equal(t, common.Address{'a'}, f) + + _, err = db.FeeRecipientByValidatorID(ctx, 4) + want = errors.Wrap(ErrNotFoundFeeRecipient, "validator id 4") + require.Equal(t, want.Error(), err.Error()) +} + +func TestStore_RegistrationsByValidatorID(t *testing.T) { + db := setupDB(t) + ctx := context.Background() + ids := []types.ValidatorIndex{0, 0, 0} + regs := []*ethpb.ValidatorRegistrationV1{{}, {}, {}, {}} + require.ErrorContains(t, "ids and registrations must be the same length", db.SaveRegistrationsByValidatorIDs(ctx, ids, regs)) + + ids = []types.ValidatorIndex{0, 1, 2} + regs = []*ethpb.ValidatorRegistrationV1{ + { + FeeRecipient: bytesutil.PadTo([]byte("a"), 20), + GasLimit: 1, + Timestamp: 2, + Pubkey: bytesutil.PadTo([]byte("b"), 48), + }, + { + FeeRecipient: bytesutil.PadTo([]byte("c"), 20), + GasLimit: 3, + Timestamp: 4, + Pubkey: bytesutil.PadTo([]byte("d"), 48), + }, + { + FeeRecipient: bytesutil.PadTo([]byte("e"), 20), + GasLimit: 5, + Timestamp: 6, + Pubkey: bytesutil.PadTo([]byte("f"), 48), + }, + } + require.NoError(t, db.SaveRegistrationsByValidatorIDs(ctx, ids, regs)) + f, err := db.RegistrationByValidatorID(ctx, 0) + require.NoError(t, err) + require.DeepEqual(t, ðpb.ValidatorRegistrationV1{ + FeeRecipient: bytesutil.PadTo([]byte("a"), 20), + GasLimit: 1, + Timestamp: 2, + Pubkey: bytesutil.PadTo([]byte("b"), 48), + }, f) + f, err = db.RegistrationByValidatorID(ctx, 1) + require.NoError(t, err) + require.DeepEqual(t, ðpb.ValidatorRegistrationV1{ + FeeRecipient: bytesutil.PadTo([]byte("c"), 20), + GasLimit: 3, + Timestamp: 4, + Pubkey: bytesutil.PadTo([]byte("d"), 48), + }, f) + f, err = db.RegistrationByValidatorID(ctx, 2) + require.NoError(t, err) + require.DeepEqual(t, ðpb.ValidatorRegistrationV1{ + FeeRecipient: bytesutil.PadTo([]byte("e"), 20), + GasLimit: 5, + Timestamp: 6, + Pubkey: bytesutil.PadTo([]byte("f"), 48), + }, f) + _, err = db.RegistrationByValidatorID(ctx, 3) + want := errors.Wrap(ErrNotFoundFeeRecipient, "validator id 3") + require.Equal(t, want.Error(), err.Error()) } diff --git a/beacon-chain/db/kv/encoding.go b/beacon-chain/db/kv/encoding.go index cfcc348ca..97016a843 100644 --- a/beacon-chain/db/kv/encoding.go +++ b/beacon-chain/db/kv/encoding.go @@ -70,6 +70,8 @@ func isSSZStorageFormat(obj interface{}) bool { return true case *ethpb.VoluntaryExit: return true + case *ethpb.ValidatorRegistrationV1: + return true default: return false } diff --git a/beacon-chain/db/kv/kv.go b/beacon-chain/db/kv/kv.go index a7d00ab5c..1aba818dd 100644 --- a/beacon-chain/db/kv/kv.go +++ b/beacon-chain/db/kv/kv.go @@ -192,6 +192,7 @@ func NewKVStore(ctx context.Context, dirPath string, config *Config) (*Store, er migrationsBucket, feeRecipientBucket, + registrationBucket, ) }); err != nil { log.WithField("elapsed", time.Since(start)).Error("Failed to update db and create buckets") diff --git a/beacon-chain/db/kv/schema.go b/beacon-chain/db/kv/schema.go index f1b813ed2..4bfe429e3 100644 --- a/beacon-chain/db/kv/schema.go +++ b/beacon-chain/db/kv/schema.go @@ -19,6 +19,7 @@ var ( powchainBucket = []byte("powchain") stateValidatorsBucket = []byte("state-validators") feeRecipientBucket = []byte("fee-recipient") + registrationBucket = []byte("registration") // Deprecated: This bucket was migrated in PR 6461. Do not use, except for migrations. slotsHasObjectBucket = []byte("slots-has-objects") diff --git a/proto/prysm/v1alpha1/BUILD.bazel b/proto/prysm/v1alpha1/BUILD.bazel index c05941872..2cdbb2403 100644 --- a/proto/prysm/v1alpha1/BUILD.bazel +++ b/proto/prysm/v1alpha1/BUILD.bazel @@ -108,6 +108,7 @@ ssz_gen_marshal( "BlindedBeaconBlockBellatrix", "BlindedBeaconBlockBodyBellatrix", "SignedValidatorRegistrationV1", + "ValidatorRegistrationV1", ], )