stategen: Pre-populate bls pubkey cache as part of stategen's Resume function (#11482)

* Pre-populate bls pubkey cache as part of state gen's Resume function. This adds some helpers and a benchmark to blst

* Do it async

* fix missing import

* lint

---------

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
Co-authored-by: rauljordan <raul@prysmaticlabs.com>

beacon-chain/state/stategen/BUILD.bazel

@@ -36,6 +36,7 @@ go_library(
         "//consensus-types/blocks:go_default_library",
         "//consensus-types/interfaces:go_default_library",
         "//consensus-types/primitives:go_default_library",
+        "//crypto/bls:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//monitoring/tracing:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",

beacon-chain/state/stategen/service.go

@@ -6,6 +6,7 @@ package stategen
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/v4/beacon-chain/db"
@@ -14,12 +15,15 @@ import (
 	"github.com/prysmaticlabs/prysm/v4/beacon-chain/sync/backfill"
 	"github.com/prysmaticlabs/prysm/v4/config/params"
 	"github.com/prysmaticlabs/prysm/v4/consensus-types/primitives"
+	"github.com/prysmaticlabs/prysm/v4/crypto/bls"
 	"github.com/prysmaticlabs/prysm/v4/encoding/bytesutil"
 	"go.opencensus.io/trace"
 )
 
 var defaultHotStateDBInterval primitives.Slot = 128
 
+var populatePubkeyCacheOnce sync.Once
+
 // StateManager represents a management object that handles the internal
 // logic of maintaining both hot and cold states in DB.
 type StateManager interface {
@@ -138,6 +142,24 @@ func (s *State) Resume(ctx context.Context, fState state.BeaconState) (state.Bea
 
 	s.finalizedInfo = &finalizedInfo{slot: fState.Slot(), root: fRoot, state: fState.Copy()}
 
+	// Pre-populate the pubkey cache with the validator public keys from the finalized state.
+	// This process takes about 30 seconds on mainnet with 450,000 validators.
+	go populatePubkeyCacheOnce.Do(func() {
+		log.Debug("Populating pubkey cache")
+		start := time.Now()
+		if err := fState.ReadFromEveryValidator(func(_ int, val state.ReadOnlyValidator) error {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			pub := val.PublicKey()
+			_, err := bls.PublicKeyFromBytes(pub[:])
+			return err
+		}); err != nil {
+			log.WithError(err).Error("Failed to populate pubkey cache")
+		}
+		log.WithField("duration", time.Since(start)).Debug("Done populating pubkey cache")
+	})
+
 	return fState, nil
 }
 

crypto/bls/blst/BUILD.bazel

@@ -34,6 +34,7 @@ go_test(
         "public_key_test.go",
         "secret_key_test.go",
         "signature_test.go",
+        "test_helper_test.go",
     ],
     embed = [":go_default_library"],
     deps = [

crypto/bls/blst/public_key.go

@@ -12,7 +12,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/common"
 )
 
-var maxKeys = 1000000
+var maxKeys = 1_000_000
 var pubkeyCache = lruwrpr.New(maxKeys)
 
 // PublicKey used in the BLS signature scheme.

crypto/bls/blst/public_key_test.go

@@ -97,3 +97,27 @@ func TestPublicKeysEmpty(t *testing.T) {
 	_, err := blst.AggregatePublicKeys(pubs)
 	require.ErrorContains(t, "nil or empty public keys", err)
 }
+
+func BenchmarkPublicKeyFromBytes(b *testing.B) {
+	priv, err := blst.RandKey()
+	require.NoError(b, err)
+	pubkey := priv.PublicKey()
+	pubkeyBytes := pubkey.Marshal()
+
+	b.Run("cache on", func(b *testing.B) {
+		blst.EnableCaches()
+		for i := 0; i < b.N; i++ {
+			_, err := blst.PublicKeyFromBytes(pubkeyBytes)
+			require.NoError(b, err)
+		}
+	})
+
+	b.Run("cache off", func(b *testing.B) {
+		blst.DisableCaches()
+		for i := 0; i < b.N; i++ {
+			_, err := blst.PublicKeyFromBytes(pubkeyBytes)
+			require.NoError(b, err)
+		}
+	})
+
+}

crypto/bls/blst/test_helper_test.go

@@ -0,0 +1,13 @@
+package blst
+
+// Note: These functions are for tests to access private globals, such as pubkeyCache.
+
+// DisableCaches sets the cache sizes to 0.
+func DisableCaches() {
+	pubkeyCache.Resize(0)
+}
+
+// EnableCaches sets the cache sizes to the default values.
+func EnableCaches() {
+	pubkeyCache.Resize(maxKeys)
+}