Merge #4256

4256: Add version beacon badger storage r=janezpodhostnik a=janezpodhostnik

Extracted from #3736, and added more cases.

Co-authored-by: Janez Podhostnik <[email protected]>

Author: bors[bot]

model/flow/version_beacon.go

@@ -30,6 +30,13 @@ type VersionBeacon struct {
 	Sequence          uint64
 }
 
+// SealedVersionBeacon is a VersionBeacon with a SealHeight field.
+// Version beacons are effective only after they are sealed.
+type SealedVersionBeacon struct {
+	*VersionBeacon
+	SealHeight uint64
+}
+
 func (v *VersionBeacon) ServiceEvent() ServiceEvent {
 	return ServiceEvent{
 		Type:  ServiceEventVersionBeacon,

storage/badger/operation/common.go

@@ -521,6 +521,43 @@ func traverse(prefix []byte, iteration iterationFunc) func(*badger.Txn) error {
 	}
 }
 
+// findHighestAtOrBelow searches for the highest key with the given prefix and a height
+// at or below the target height, and retrieves and decodes the value associated with the
+// key into the given entity.
+// If no key is found, the function returns storage.ErrNotFound.
+func findHighestAtOrBelow(
+	prefix []byte,
+	height uint64,
+	entity interface{},
+) func(*badger.Txn) error {
+	return func(tx *badger.Txn) error {
+		if len(prefix) == 0 {
+			return fmt.Errorf("prefix must not be empty")
+		}
+
+		opts := badger.DefaultIteratorOptions
+		opts.Prefix = prefix
+		opts.Reverse = true
+
+		it := tx.NewIterator(opts)
+		defer it.Close()
+
+		it.Seek(append(prefix, b(height)...))
+
+		if !it.Valid() {
+			return storage.ErrNotFound
+		}
+
+		return it.Item().Value(func(val []byte) error {
+			err := msgpack.Unmarshal(val, entity)
+			if err != nil {
+				return fmt.Errorf("could not decode entity: %w", err)
+			}
+			return nil
+		})
+	}
+}
+
 // Fail returns a DB operation function that always fails with the given error.
 func Fail(err error) func(*badger.Txn) error {
 	return func(_ *badger.Txn) error {

storage/badger/operation/common_test.go

@@ -614,3 +614,97 @@ func TestIterateBoundaries(t *testing.T) {
 		assert.ElementsMatch(t, keysInRange, found, "backward iteration should go over correct keys")
 	})
 }
+
+func TestFindHighestAtOrBelow(t *testing.T) {
+	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+		prefix := []byte("test_prefix")
+
+		type Entity struct {
+			Value uint64
+		}
+
+		entity1 := Entity{Value: 41}
+		entity2 := Entity{Value: 42}
+		entity3 := Entity{Value: 43}
+
+		err := db.Update(func(tx *badger.Txn) error {
+			key := append(prefix, b(uint64(15))...)
+			val, err := msgpack.Marshal(entity3)
+			if err != nil {
+				return err
+			}
+			err = tx.Set(key, val)
+			if err != nil {
+				return err
+			}
+
+			key = append(prefix, b(uint64(5))...)
+			val, err = msgpack.Marshal(entity1)
+			if err != nil {
+				return err
+			}
+			err = tx.Set(key, val)
+			if err != nil {
+				return err
+			}
+
+			key = append(prefix, b(uint64(10))...)
+			val, err = msgpack.Marshal(entity2)
+			if err != nil {
+				return err
+			}
+			err = tx.Set(key, val)
+			if err != nil {
+				return err
+			}
+			return nil
+		})
+		require.NoError(t, err)
+
+		var entity Entity
+
+		t.Run("target height exists", func(t *testing.T) {
+			err = findHighestAtOrBelow(
+				prefix,
+				10,
+				&entity)(db.NewTransaction(false))
+			require.NoError(t, err)
+			require.Equal(t, uint64(42), entity.Value)
+		})
+
+		t.Run("target height above", func(t *testing.T) {
+			err = findHighestAtOrBelow(
+				prefix,
+				11,
+				&entity)(db.NewTransaction(false))
+			require.NoError(t, err)
+			require.Equal(t, uint64(42), entity.Value)
+		})
+
+		t.Run("target height above highest", func(t *testing.T) {
+			err = findHighestAtOrBelow(
+				prefix,
+				20,
+				&entity)(db.NewTransaction(false))
+			require.NoError(t, err)
+			require.Equal(t, uint64(43), entity.Value)
+		})
+
+		t.Run("target height below lowest", func(t *testing.T) {
+			err = findHighestAtOrBelow(
+				prefix,
+				4,
+				&entity)(db.NewTransaction(false))
+			require.ErrorIs(t, err, storage.ErrNotFound)
+		})
+
+		t.Run("empty prefix", func(t *testing.T) {
+			err = findHighestAtOrBelow(
+				[]byte{},
+				5,
+				&entity)(db.NewTransaction(false))
+			require.Error(t, err)
+			require.Contains(t, err.Error(), "prefix must not be empty")
+		})
+	})
+}

storage/badger/operation/prefix.go

@@ -66,12 +66,13 @@ const (
 	codePayloadResults    = 58 // index mapping block ID to payload results
 	codeAllBlockReceipts  = 59 // index mapping of blockID to multiple receipts
 
-	// codes related to epoch information
+	// codes related to protocol level information
 	codeEpochSetup       = 61 // EpochSetup service event, keyed by ID
 	codeEpochCommit      = 62 // EpochCommit service event, keyed by ID
 	codeBeaconPrivateKey = 63 // BeaconPrivateKey, keyed by epoch counter
 	codeDKGStarted       = 64 // flag that the DKG for an epoch has been started
 	codeDKGEnded         = 65 // flag that the DKG for an epoch has ended (stores end state)
+	codeVersionBeacon    = 67 // flag for storing version beacons
 
 	// code for ComputationResult upload status storage
 	// NOTE: for now only GCP uploader is supported. When other uploader (AWS e.g.) needs to

storage/badger/operation/version_beacon.go

@@ -0,0 +1,31 @@
+package operation
+
+import (
+	"github.com/dgraph-io/badger/v2"
+
+	"github.com/onflow/flow-go/model/flow"
+)
+
+// IndexVersionBeaconByHeight stores a sealed version beacon indexed by
+// flow.SealedVersionBeacon.SealHeight.
+//
+// No errors are expected during normal operation.
+func IndexVersionBeaconByHeight(
+	beacon flow.SealedVersionBeacon,
+) func(*badger.Txn) error {
+	return upsert(makePrefix(codeVersionBeacon, beacon.SealHeight), beacon)
+}
+
+// LookupLastVersionBeaconByHeight finds the highest flow.VersionBeacon but no higher
+// than maxHeight. Returns storage.ErrNotFound if no version beacon exists at or below
+// the given height.
+func LookupLastVersionBeaconByHeight(
+	maxHeight uint64,
+	versionBeacon *flow.SealedVersionBeacon,
+) func(*badger.Txn) error {
+	return findHighestAtOrBelow(
+		makePrefix(codeVersionBeacon),
+		maxHeight,
+		versionBeacon,
+	)
+}

storage/badger/operation/version_beacon_test.go

@@ -0,0 +1,106 @@
+package operation
+
+import (
+	"testing"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/stretchr/testify/require"
+
+	"github.com/onflow/flow-go/model/flow"
+	"github.com/onflow/flow-go/storage"
+	"github.com/onflow/flow-go/utils/unittest"
+)
+
+func TestResults_IndexByServiceEvents(t *testing.T) {
+	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+		height1 := uint64(21)
+		height2 := uint64(37)
+		height3 := uint64(55)
+		vb1 := flow.SealedVersionBeacon{
+			VersionBeacon: unittest.VersionBeaconFixture(
+				unittest.WithBoundaries(
+					flow.VersionBoundary{
+						Version:     "1.0.0",
+						BlockHeight: height1 + 5,
+					},
+				),
+			),
+			SealHeight: height1,
+		}
+		vb2 := flow.SealedVersionBeacon{
+			VersionBeacon: unittest.VersionBeaconFixture(
+				unittest.WithBoundaries(
+					flow.VersionBoundary{
+						Version:     "1.1.0",
+						BlockHeight: height2 + 5,
+					},
+				),
+			),
+			SealHeight: height2,
+		}
+		vb3 := flow.SealedVersionBeacon{
+			VersionBeacon: unittest.VersionBeaconFixture(
+				unittest.WithBoundaries(
+					flow.VersionBoundary{
+						Version:     "2.0.0",
+						BlockHeight: height3 + 5,
+					},
+				),
+			),
+			SealHeight: height3,
+		}
+
+		// indexing 3 version beacons at different heights
+		err := db.Update(IndexVersionBeaconByHeight(vb1))
+		require.NoError(t, err)
+
+		err = db.Update(IndexVersionBeaconByHeight(vb2))
+		require.NoError(t, err)
+
+		err = db.Update(IndexVersionBeaconByHeight(vb3))
+		require.NoError(t, err)
+
+		// index version beacon 2 again to make sure we tolerate duplicates
+		// it is possible for two or more events of the same type to be from the same height
+		err = db.Update(IndexVersionBeaconByHeight(vb2))
+		require.NoError(t, err)
+
+		t.Run("retrieve exact height match", func(t *testing.T) {
+			var actualVB flow.SealedVersionBeacon
+			err := db.View(LookupLastVersionBeaconByHeight(height1, &actualVB))
+			require.NoError(t, err)
+			require.Equal(t, vb1, actualVB)
+
+			err = db.View(LookupLastVersionBeaconByHeight(height2, &actualVB))
+			require.NoError(t, err)
+			require.Equal(t, vb2, actualVB)
+
+			err = db.View(LookupLastVersionBeaconByHeight(height3, &actualVB))
+			require.NoError(t, err)
+			require.Equal(t, vb3, actualVB)
+		})
+
+		t.Run("finds highest but not higher than given", func(t *testing.T) {
+			var actualVB flow.SealedVersionBeacon
+
+			err := db.View(LookupLastVersionBeaconByHeight(height3-1, &actualVB))
+			require.NoError(t, err)
+			require.Equal(t, vb2, actualVB)
+		})
+
+		t.Run("finds highest", func(t *testing.T) {
+			var actualVB flow.SealedVersionBeacon
+
+			err := db.View(LookupLastVersionBeaconByHeight(height3+1, &actualVB))
+			require.NoError(t, err)
+			require.Equal(t, vb3, actualVB)
+		})
+
+		t.Run("height below lowest entry returns nothing", func(t *testing.T) {
+			var actualVB flow.SealedVersionBeacon
+
+			err := db.View(LookupLastVersionBeaconByHeight(height1-1, &actualVB))
+			require.ErrorIs(t, err, storage.ErrNotFound)
+		})
+	})
+}

storage/badger/version_beacon.go

@@ -0,0 +1,38 @@
+package badger
+
+import (
+	"github.com/dgraph-io/badger/v2"
+
+	"github.com/onflow/flow-go/model/flow"
+	"github.com/onflow/flow-go/storage"
+	"github.com/onflow/flow-go/storage/badger/operation"
+)
+
+type VersionBeacons struct {
+	db *badger.DB
+}
+
+var _ storage.VersionBeacons = (*VersionBeacons)(nil)
+
+func NewVersionBeacons(db *badger.DB) *VersionBeacons {
+	res := &VersionBeacons{
+		db: db,
+	}
+
+	return res
+}
+
+func (r *VersionBeacons) Highest(
+	belowOrEqualTo uint64,
+) (*flow.SealedVersionBeacon, error) {
+	tx := r.db.NewTransaction(false)
+	defer tx.Discard()
+
+	var beacon *flow.SealedVersionBeacon
+
+	err := operation.LookupLastVersionBeaconByHeight(belowOrEqualTo, beacon)(tx)
+	if err != nil {
+		return nil, err
+	}
+	return beacon, nil
+}