--beacon-db-pruning: Prune state diff buckets when the --enable-state-diff flag is used.

beacon-chain/db/iface/interface.go

@@ -127,6 +127,8 @@ type NoHeadAccessDatabase interface {
 
 	CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint primitives.Slot) error
 	DeleteHistoricalDataBeforeSlot(ctx context.Context, slot primitives.Slot, batchSize int) (int, error)
+	DeleteStateDiffBeforeSlot(ctx context.Context, slot primitives.Slot, maxEntries int) (int, error)
+	LastStateDiffBoundary(slot primitives.Slot) (primitives.Slot, error)
 
 	// Genesis operations.
 	LoadGenesis(ctx context.Context, stateBytes []byte) error

beacon-chain/db/kv/BUILD.bazel

@@ -31,6 +31,7 @@ go_library(
         "state_diff.go",
         "state_diff_cache.go",
         "state_diff_helpers.go",
+        "state_diff_prune.go",
         "state_hot_snapshots.go",
         "state_summary.go",
         "state_summary_cache.go",
@@ -109,6 +110,7 @@ go_test(
         "migration_state_validators_test.go",
         "p2p_test.go",
         "state_diff_helpers_test.go",
+        "state_diff_prune_test.go",
         "state_diff_test.go",
         "state_hot_snapshots_test.go",
         "state_summary_test.go",

beacon-chain/db/kv/state_diff_cache.go

@@ -230,6 +230,17 @@ func (c *stateDiffCache) setAnchor(level int, anchor state.ReadOnlyBeaconState)
 	return nil
 }
 
+// reanchor points the cache at a new offset and drops the cached anchors
+func (c *stateDiffCache) reanchor(offset uint64, levelsWithData []bool) {
+	c.Lock()
+	defer c.Unlock()
+
+	c.offset = offset
+	c.levelsWithData = levelsWithData
+
+	c.clearAnchorsLocked()
+}
+
 func (c *stateDiffCache) levelHasData(level int) bool {
 	c.RLock()
 	defer c.RUnlock()
@@ -264,6 +275,11 @@ func (c *stateDiffCache) setOffset(offset uint64) {
 func (c *stateDiffCache) clearAnchors() {
 	c.Lock()
 	defer c.Unlock()
+	c.clearAnchorsLocked()
+}
+
+// clearAnchorsLocked is clearAnchors, for the callers that already hold the lock.
+func (c *stateDiffCache) clearAnchorsLocked() {
 	c.anchors = make([][]byte, len(flags.Get().StateDiffExponents)-1) // -1 because last level doesn't need to be cached
 	for level := range len(c.anchors) {
 		stateDiffAnchorCacheBytes.WithLabelValues(strconv.Itoa(level)).Set(0)

beacon-chain/db/kv/state_diff_helpers.go

@@ -21,6 +21,16 @@ import (
 	"go.etcd.io/bbolt"
 )
 
+const (
+	// stateDiffTreeKeyLength is the length of a state-diff tree key, before any suffix.
+	stateDiffTreeKeyLength = 16
+
+	// stateDiffTreeKeySlotEnd is the end of the meaningful part of a state-diff tree key: a level
+	// byte followed by a little-endian slot. The bytes up to stateDiffTreeKeyLength are padding,
+	// and are always zero.
+	stateDiffTreeKeySlotEnd = 9
+)
+
 var (
 	offsetKey                   = []byte("offset")
 	exponentsKey                = []byte("exponents")
@@ -112,12 +122,38 @@ func (s *Store) loadStateDiffExponents() ([]int, error) {
 }
 
 func makeKeyForStateDiffTree(level int, slot uint64) []byte {
-	buf := make([]byte, 16)
+	buf := make([]byte, stateDiffTreeKeyLength)
 	buf[0] = byte(level)
-	binary.LittleEndian.PutUint64(buf[1:], slot)
+	binary.LittleEndian.PutUint64(buf[1:stateDiffTreeKeySlotEnd], slot)
 	return buf
 }
 
+// isStateDiffTreeKey reports whether the given key holds a tree entry, as opposed to one of the
+// metadata keys stored in the same bucket.
+func isStateDiffTreeKey(key []byte) bool {
+	if len(key) < stateDiffTreeKeyLength {
+		return false
+	}
+
+	if int(key[0]) >= len(flags.Get().StateDiffExponents) {
+		return false
+	}
+
+	for _, padding := range key[stateDiffTreeKeySlotEnd:stateDiffTreeKeyLength] {
+		if padding != 0 {
+			return false
+		}
+	}
+
+	return true
+}
+
+// stateDiffTreeKeySlot returns the slot a state-diff tree key is stored at.
+// It must only be called on a key that isStateDiffTreeKey accepts.
+func stateDiffTreeKeySlot(key []byte) uint64 {
+	return binary.LittleEndian.Uint64(key[1:stateDiffTreeKeySlotEnd])
+}
+
 func (s *Store) getAnchorState(ctx context.Context, offset uint64, lvl int, slot primitives.Slot) (anchor state.ReadOnlyBeaconState, err error) {
 	if lvl <= 0 || lvl > len(flags.Get().StateDiffExponents) {
 		return nil, errors.New("invalid value for level")
@@ -127,10 +163,8 @@ func (s *Store) getAnchorState(ctx context.Context, offset uint64, lvl int, slot
 		return nil, ErrSlotBeforeOffset
 	}
 	relSlot := uint64(slot) - offset
+	// The exponents are validated at node startup, so they always fit in a uint64 shift.
 	prevExp := flags.Get().StateDiffExponents[lvl-1]
-	if prevExp < flags.MinStateDiffExponent || prevExp >= 64 {
-		return nil, fmt.Errorf("state diff exponent %d out of range for uint64", prevExp)
-	}
 	span := math.PowerOf2(uint64(prevExp))
 	anchorSlot := primitives.Slot(uint64(slot) - relSlot%span)
 
@@ -179,10 +213,8 @@ func computeLevel(offset uint64, slot primitives.Slot) int {
 		return -1
 	}
 	rel := uint64(slot) - offset
+	// The exponents are validated at node startup, so they always fit in a uint64 shift.
 	for i, exp := range flags.Get().StateDiffExponents {
-		if exp < flags.MinStateDiffExponent || exp >= 64 {
-			return -1
-		}
 		span := math.PowerOf2(uint64(exp))
 		if rel%span == 0 {
 			return i

beacon-chain/db/kv/state_diff_helpers_test.go

@@ -1,6 +1,7 @@
 package kv
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/OffchainLabs/prysm/v7/runtime/version"
@@ -17,6 +18,38 @@ func TestMakeKeyForStateDiffTree_KeyLength(t *testing.T) {
 	require.Equal(t, 16, len(key))
 }
 
+func TestIsStateDiffTreeKey(t *testing.T) {
+	setStateDiffExponents([]int{7, 5})
+
+	// A tree key with a level byte, a slot, and zero padding, then the same with an entry suffix.
+	treeKey := makeKeyForStateDiffTree(1, 320)
+	suffixedKey := append(bytes.Clone(treeKey), stateSuffix...)
+
+	// A key that is shaped like a tree key up to its padding, which a tree key never sets.
+	paddedKey := bytes.Clone(treeKey)
+	paddedKey[stateDiffTreeKeySlotEnd] = 'x'
+
+	tests := []struct {
+		name string
+		key  []byte
+		want bool
+	}{
+		{name: "tree key", key: treeKey, want: true},
+		{name: "suffixed tree key", key: suffixedKey, want: true},
+		{name: "offset metadata key", key: offsetKey, want: false},
+		{name: "exponents metadata key", key: exponentsKey, want: false},
+		{name: "metadata key longer than a tree key", key: []byte("a-long-metadata-key-here"), want: false},
+		{name: "level byte out of range", key: append([]byte("m"), make([]byte, stateDiffTreeKeyLength)...), want: false},
+		{name: "non-zero padding", key: paddedKey, want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.want, isStateDiffTreeKey(tt.key))
+		})
+	}
+}
+
 func TestKeyForSnapshot_AllVersions(t *testing.T) {
 	for _, v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {