state diff kv functions

Author: Inspector-Butters

beacon-chain/db/kv/BUILD.bazel

@@ -27,7 +27,9 @@ go_library(
         "p2p.go",
         "schema.go",
         "state.go",
+        "state_diff.go",
         "state_diff_cache.go",
+        "state_diff_helpers.go",
         "state_summary.go",
         "state_summary_cache.go",
         "utils.go",
@@ -47,6 +49,7 @@ go_library(
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",
+        "//consensus-types/hdiff:go_default_library",
         "//consensus-types/interfaces:go_default_library",
         "//consensus-types/light-client:go_default_library",
         "//consensus-types/primitives:go_default_library",
@@ -55,6 +58,7 @@ go_library(
         "//encoding/ssz/detect:go_default_library",
         "//genesis:go_default_library",
         "//io/file:go_default_library",
+        "//math:go_default_library",
         "//monitoring/progress:go_default_library",
         "//monitoring/tracing:go_default_library",
         "//monitoring/tracing/trace:go_default_library",
@@ -100,6 +104,7 @@ go_test(
         "migration_block_slot_index_test.go",
         "migration_state_validators_test.go",
         "p2p_test.go",
+        "state_diff_test.go",
         "state_summary_test.go",
         "state_test.go",
         "utils_test.go",
@@ -113,6 +118,7 @@ go_test(
         "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/state-native:go_default_library",
+        "//cmd/beacon-chain/flags:go_default_library",
         "//config/features:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
@@ -122,6 +128,7 @@ go_test(
         "//consensus-types/primitives:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//genesis:go_default_library",
+        "//math:go_default_library",
         "//proto/dbval:go_default_library",
         "//proto/engine/v1:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",

beacon-chain/db/kv/state_diff.go

@@ -0,0 +1,232 @@
+package kv
+
+import (
+	"context"
+
+	"github.com/OffchainLabs/prysm/v6/beacon-chain/state"
+	"github.com/OffchainLabs/prysm/v6/cmd/beacon-chain/flags"
+	"github.com/OffchainLabs/prysm/v6/consensus-types/hdiff"
+	"github.com/OffchainLabs/prysm/v6/consensus-types/primitives"
+	"github.com/OffchainLabs/prysm/v6/monitoring/tracing/trace"
+	"github.com/pkg/errors"
+	bolt "go.etcd.io/bbolt"
+)
+
+const (
+	stateSuffix     = "_s"
+	validatorSuffix = "_v"
+	balancesSuffix  = "_b"
+)
+
+/*
+	We use a level-based approach to save state diffs. The levels are 0-6, where each level corresponds to an exponent of 2 (exponents[lvl]).
+	The data at level 0 is saved every 2**exponent[0] slots and always contains a full state snapshot that is used as a base for the delta saved at other levels.
+*/
+
+// saveStateByDiff takes a state and decides between saving a full state snapshot or a diff.
+func (s *Store) saveStateByDiff(ctx context.Context, st state.ReadOnlyBeaconState) error {
+	_, span := trace.StartSpan(ctx, "BeaconDB.saveStateByDiff")
+	defer span.End()
+
+	if st == nil {
+		return errors.New("state is nil")
+	}
+
+	slot := st.Slot()
+	offset := s.getOffset()
+	if uint64(slot) < offset {
+		return ErrSlotBeforeOffset
+	}
+
+	// Find the level to save the state.
+	lvl := computeLevel(offset, slot)
+	if lvl == -1 {
+		return nil
+	}
+
+	// Save full state if level is 0.
+	if lvl == 0 {
+		return s.saveFullSnapshot(st)
+	}
+
+	// Get anchor state to compute the diff from.
+	anchorState, err := s.getAnchorState(offset, lvl, slot)
+	if err != nil {
+		return err
+	}
+
+	err = s.saveHdiff(lvl, anchorState, st)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// stateByDiff retrieves the full state for a given slot.
+func (s *Store) stateByDiff(ctx context.Context, slot primitives.Slot) (state.BeaconState, error) {
+	offset := s.getOffset()
+	if uint64(slot) < offset {
+		return nil, ErrSlotBeforeOffset
+	}
+
+	snapshot, diffChain, err := s.getBaseAndDiffChain(offset, slot)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, diff := range diffChain {
+		snapshot, err = hdiff.ApplyDiff(ctx, snapshot, diff)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return snapshot, nil
+}
+
+// SaveHdiff computes the diff between the anchor state and the current state and saves it to the database.
+func (s *Store) saveHdiff(lvl int, anchor, st state.ReadOnlyBeaconState) error {
+	slot := uint64(st.Slot())
+	key := makeKey(lvl, slot)
+
+	diff, err := hdiff.Diff(anchor, st)
+	if err != nil {
+		return err
+	}
+
+	err = s.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+		buf := append(key, stateSuffix...)
+		if err := bucket.Put(buf, diff.StateDiff); err != nil {
+			return err
+		}
+		buf = append(key, validatorSuffix...)
+		if err := bucket.Put(buf, diff.ValidatorDiffs); err != nil {
+			return err
+		}
+		buf = append(key, balancesSuffix...)
+		if err := bucket.Put(buf, diff.BalancesDiff); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	// Save the full state to the cache (if not the last level).
+	if lvl != len(flags.Get().StateDiffExponents)-1 {
+		err = s.stateDiffCache.setAnchor(lvl, st)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// SaveFullSnapshot saves the full level 0 state snapshot to the database.
+func (s *Store) saveFullSnapshot(st state.ReadOnlyBeaconState) error {
+	slot := uint64(st.Slot())
+	key := makeKey(0, slot)
+	stateBytes, err := st.MarshalSSZ()
+	if err != nil {
+		return err
+	}
+	// add version key to value
+	enc, err := addKey(st.Version(), stateBytes)
+	if err != nil {
+		return err
+	}
+
+	err = s.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+
+		if err := bucket.Put(key, enc); err != nil {
+			return err
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	// Save the full state to the cache, and invalidate other levels.
+	s.stateDiffCache.clearAnchors()
+	err = s.stateDiffCache.setAnchor(0, st)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Store) getDiff(lvl int, slot uint64) (hdiff.HdiffBytes, error) {
+	key := makeKey(lvl, slot)
+	var stateDiff []byte
+	var validatorDiff []byte
+	var balancesDiff []byte
+
+	err := s.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+		buf := append(key, stateSuffix...)
+		stateDiff = bucket.Get(buf)
+		if stateDiff == nil {
+			return errors.New("state diff not found")
+		}
+		buf = append(key, validatorSuffix...)
+		validatorDiff = bucket.Get(buf)
+		if validatorDiff == nil {
+			return errors.New("validator diff not found")
+		}
+		buf = append(key, balancesSuffix...)
+		balancesDiff = bucket.Get(buf)
+		if balancesDiff == nil {
+			return errors.New("balances diff not found")
+		}
+		return nil
+	})
+
+	if err != nil {
+		return hdiff.HdiffBytes{}, err
+	}
+
+	return hdiff.HdiffBytes{
+		StateDiff:      stateDiff,
+		ValidatorDiffs: validatorDiff,
+		BalancesDiff:   balancesDiff,
+	}, nil
+}
+
+func (s *Store) getFullSnapshot(slot uint64) (state.BeaconState, error) {
+	key := makeKey(0, slot)
+	var enc []byte
+
+	err := s.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateDiffBucket)
+		if bucket == nil {
+			return bolt.ErrBucketNotFound
+		}
+		enc = bucket.Get(key)
+		if enc == nil {
+			return errors.New("state not found")
+		}
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return s.decodeStateSnapshot(enc)
+}