Add mempool feature flag (prysmaticlabs#4824)

* Add mempool feature flag

* gate put too

* fix

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

shared/featureconfig/config.go

@@ -42,6 +42,7 @@ type Flags struct {
 	ProtectAttester                            bool   // ProtectAttester prevents the validator client from signing any attestations that would be considered a slashable offense.
 	DisableStrictAttestationPubsubVerification bool   // DisableStrictAttestationPubsubVerification will disabling strict signature verification in pubsub.
 	DisableUpdateHeadPerAttestation            bool   // DisableUpdateHeadPerAttestation will disabling update head on per attestation basis.
+	EnableByteMempool                          bool   // EnaableByteMempool memory management.
 
 	// DisableForkChoice disables using LMD-GHOST fork choice to update
 	// the head of the chain based on attestations and instead accepts any valid received block
@@ -148,6 +149,10 @@ func ConfigureBeaconChain(ctx *cli.Context) {
 		log.Warn("Disabled update head on per attestation basis")
 		cfg.DisableUpdateHeadPerAttestation = true
 	}
+	if ctx.GlobalBool(enableByteMempool.Name) {
+		log.Warn("Enabling experimental memory management for beacon state")
+		cfg.EnableByteMempool = true
+	}
 
 	Init(cfg)
 }

shared/featureconfig/flags.go

@@ -98,6 +98,10 @@ var (
 		Name:  "disable-update-head-attestation",
 		Usage: "Disable update fork choice head on per attestation. See PR 4802 for details.",
 	}
+	enableByteMempool = cli.BoolFlag{
+		Name: "enable-byte-mempool",
+		Usage: "Enable use of sync.Pool for certain byte arrays in the beacon state",
+	}
 )
 
 // Deprecated flags list.
@@ -256,6 +260,7 @@ var BeaconChainFlags = append(deprecatedFlags, []cli.Flag{
 	cacheFilteredBlockTreeFlag,
 	disableStrictAttestationPubsubVerificationFlag,
 	disableUpdateHeadPerAttestation,
+	enableByteMempool,
 }...)
 
 // E2EBeaconChainFlags contains a list of the beacon chain feature flags to be tested in E2E.
@@ -268,4 +273,5 @@ var E2EBeaconChainFlags = []string{
 	"--enable-eth1-data-vote-cache",
 	"--initial-sync-cache-state",
 	"--proto-array-forkchoice",
+	"--enable-byte-mempool",
 }

shared/memorypool/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = ["memorypool.go"],
     importpath = "github.com/prysmaticlabs/prysm/shared/memorypool",
     visibility = ["//visibility:public"],
+    deps = ["//shared/featureconfig:go_default_library"],
 )
 
 go_test(

shared/memorypool/memorypool.go

@@ -1,6 +1,10 @@
 package memorypool
 
-import "sync"
+import (
+	"sync"
+
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+)
 
 // DoubleByteSlicePool represents the memory pool
 // for 2d byte slices
@@ -9,6 +13,10 @@ var DoubleByteSlicePool = new(sync.Pool)
 // GetDoubleByteSlice retrieves the 2d byte slice of
 // the desired size from the memory pool.
 func GetDoubleByteSlice(size int) [][]byte {
+	if !featureconfig.Get().EnableByteMempool {
+		return make([][]byte, size)
+	}
+
 	rawObj := DoubleByteSlicePool.Get()
 	if rawObj == nil {
 		return make([][]byte, size)
@@ -23,5 +31,7 @@ func GetDoubleByteSlice(size int) [][]byte {
 // PutDoubleByteSlice places the provided 2d byte slice
 // in the memory pool
 func PutDoubleByteSlice(data [][]byte) {
-	DoubleByteSlicePool.Put(data)
+	if featureconfig.Get().EnableByteMempool {
+		DoubleByteSlicePool.Put(data)
+	}
 }