fix build

crypto/bls/blst/public_key.go

@@ -4,16 +4,16 @@ package blst
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/pkg/errors"
-	lruwrpr "github.com/prysmaticlabs/prysm/v4/cache/lru"
 	fieldparams "github.com/prysmaticlabs/prysm/v4/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v4/config/params"
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/common"
 )
 
-var maxKeys = 1_000_000
-var pubkeyCache = lruwrpr.New(maxKeys)
+var pubkeyLock sync.RWMutex
+var pubkeyMap = make(map[[48]byte]*PublicKey)
 
 // PublicKey used in the BLS signature scheme.
 type PublicKey struct {
@@ -26,9 +26,13 @@ func PublicKeyFromBytes(pubKey []byte) (common.PublicKey, error) {
 		return nil, fmt.Errorf("public key must be %d bytes", params.BeaconConfig().BLSPubkeyLength)
 	}
 	newKey := (*[fieldparams.BLSPubkeyLength]byte)(pubKey)
-	if cv, ok := pubkeyCache.Get(*newKey); ok {
-		return cv.(*PublicKey).Copy(), nil
+	pubkeyLock.RLock()
+	if cv, ok := pubkeyMap[*newKey]; ok {
+		pubkeyLock.RUnlock()
+		return cv.Copy(), nil
 	}
+	pubkeyLock.RUnlock()
+
 	// Subgroup check NOT done when decompressing pubkey.
 	p := new(blstPublicKey).Uncompress(pubKey)
 	if p == nil {
@@ -41,8 +45,16 @@ func PublicKeyFromBytes(pubKey []byte) (common.PublicKey, error) {
 	}
 	pubKeyObj := &PublicKey{p: p}
 	copiedKey := pubKeyObj.Copy()
+	assertedKey, ok := (copiedKey).(*PublicKey)
+	// Should be impossible to happen, this is checked
+	// to satisfy lint tools.
+	if !ok {
+		return pubKeyObj, nil
+	}
 	cacheKey := *newKey
-	pubkeyCache.Add(cacheKey, copiedKey)
+	pubkeyLock.Lock()
+	pubkeyMap[cacheKey] = assertedKey
+	pubkeyLock.Unlock()
 	return pubKeyObj, nil
 }
 

crypto/bls/blst/public_key_test.go

@@ -5,6 +5,7 @@ package blst_test
 import (
 	"bytes"
 	"errors"
+	"sync"
 	"testing"
 
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
@@ -98,26 +99,51 @@ func TestPublicKeysEmpty(t *testing.T) {
 	require.ErrorContains(t, "nil or empty public keys", err)
 }
 
+func TestPublicKeyMap(t *testing.T) {
+	priv, err := blst.RandKey()
+	require.NoError(t, err)
+	pubkeyA := priv.PublicKey().Marshal()
+
+	priv2, err := blst.RandKey()
+	require.NoError(t, err)
+	pubkeyB := priv2.PublicKey().Marshal()
+	km := blst.KeyMap()
+	_, ok := km[[48]byte(pubkeyA)]
+	require.Equal(t, false, ok, "pubkey a exists")
+	_, ok = km[[48]byte(pubkeyB)]
+	require.Equal(t, false, ok, "pubkey b exists")
+	wg := new(sync.WaitGroup)
+	wg.Add(1)
+	go func() {
+		_, err := blst.PublicKeyFromBytes(pubkeyA)
+		require.NoError(t, err)
+		wg.Done()
+	}()
+
+	wg.Add(1)
+	go func() {
+		_, err := blst.PublicKeyFromBytes(pubkeyB)
+		require.NoError(t, err)
+		wg.Done()
+	}()
+	wg.Wait()
+
+	km = blst.KeyMap()
+	_, ok = km[[48]byte(pubkeyA)]
+	require.Equal(t, true, ok, "pubkey a does not exist")
+	_, ok = km[[48]byte(pubkeyB)]
+	require.Equal(t, true, ok, "pubkey b does not exist")
+
+}
+
 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)
-		}
-	})
-
+	for i := 0; i < b.N; i++ {
+		_, err := blst.PublicKeyFromBytes(pubkeyBytes)
+		require.NoError(b, err)
+	}
 }

crypto/bls/blst/test_helper_test.go

@@ -1,13 +1,8 @@
 package blst
 
-// Note: These functions are for tests to access private globals, such as pubkeyCache.
+// Note: These functions are for tests to access private globals, such as pubkeyMap.
 
-// 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)
+// KeyMap returns the pubkey cache.
+func KeyMap() map[[48]byte]*PublicKey {
+	return pubkeyMap
 }