Add registration DB methods (#10812)

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 := &ethpb.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 := &ethpb.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")

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, &ethpb.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, &ethpb.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, &ethpb.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())
 }

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
 	}

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")

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")

proto/prysm/v1alpha1/BUILD.bazel

@@ -108,6 +108,7 @@ ssz_gen_marshal(
         "BlindedBeaconBlockBellatrix",
         "BlindedBeaconBlockBodyBellatrix",
         "SignedValidatorRegistrationV1",
+        "ValidatorRegistrationV1",
     ],
 )