Use Memory Pool for Randao Mixes (#4896)

* add mem pool

* use mem pool

* Update shared/memorypool/memorypool.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>

* Update shared/memorypool/memorypool.go

Co-Authored-By: terence tsao <terence@prysmaticlabs.com>

Co-authored-by: terence tsao <terence@prysmaticlabs.com>

beacon-chain/state/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/hashutil:go_default_library",
+        "//shared/memorypool:go_default_library",
         "//shared/params:go_default_library",
         "//shared/stateutil:go_default_library",
         "@com_github_gogo_protobuf//proto:go_default_library",

beacon-chain/state/getters.go

@@ -2,11 +2,12 @@ package state
 
 import (
 	"fmt"
-	"github.com/prysmaticlabs/go-bitfield"
 
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-bitfield"
 	pbp2p "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/memorypool"
 )
 
 // EffectiveBalance returns the effective balance of the
@@ -452,7 +453,7 @@ func (b *BeaconState) RandaoMixes() [][]byte {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
-	mixes := make([][]byte, len(b.state.RandaoMixes))
+	mixes := memorypool.GetDoubleByteSlice(len(b.state.RandaoMixes))
 	for i, r := range b.state.RandaoMixes {
 		tmpRt := make([]byte, len(r))
 		copy(tmpRt, r)

beacon-chain/state/types.go

@@ -12,6 +12,7 @@ import (
 	pbp2p "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/hashutil"
+	"github.com/prysmaticlabs/prysm/shared/memorypool"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/stateutil"
 )
@@ -141,8 +142,11 @@ func (b *BeaconState) Copy() *BeaconState {
 
 	// Finalizer runs when dst is being destroyed in garbage collection.
 	runtime.SetFinalizer(dst, func(b *BeaconState) {
-		for _, v := range b.sharedFieldReferences {
+		for field, v := range b.sharedFieldReferences {
 			v.refs--
+			if field == randaoMixes && v.refs == 0 {
+				memorypool.PutDoubleByteSlice(b.state.RandaoMixes)
+			}
 		}
 	})
 

shared/memorypool/BUILD.bazel

@@ -0,0 +1,14 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "go_default_library",
+    srcs = ["memorypool.go"],
+    importpath = "github.com/prysmaticlabs/prysm/shared/memorypool",
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["memorypool_test.go"],
+    embed = [":go_default_library"],
+)

shared/memorypool/memorypool.go

@@ -0,0 +1,27 @@
+package memorypool
+
+import "sync"
+
+// DoubleByteSlicePool represents the memory pool
+// for 2d byte slices.
+var DoubleByteSlicePool = new(sync.Pool)
+
+// GetDoubleByteSlice retrieves the 2d byte slice of
+// the desired size from the memory pool.
+func GetDoubleByteSlice(size int) [][]byte {
+	rawObj := DoubleByteSlicePool.Get()
+	if rawObj == nil {
+		return make([][]byte, size)
+	}
+	byteSlice := rawObj.([][]byte)
+	if len(byteSlice) >= size {
+		return byteSlice[:size]
+	}
+	return append(byteSlice, make([][]byte, size-len(byteSlice))...)
+}
+
+// PutDoubleByteSlice places the provided 2d byte slice
+// in the memory pool.
+func PutDoubleByteSlice(data [][]byte) {
+	DoubleByteSlicePool.Put(data)
+}

shared/memorypool/memorypool_test.go

@@ -0,0 +1,16 @@
+package memorypool
+
+import (
+	"testing"
+)
+
+func TestRoundTripMemoryRetrieval(t *testing.T) {
+	byteSlice := make([][]byte, 1000)
+	PutDoubleByteSlice(byteSlice)
+	newSlice := GetDoubleByteSlice(1000)
+
+	if len(newSlice) != 1000 {
+		t.Errorf("Wanted same slice object, but got different object. "+
+			"Wanted  slice with length %d but got length %d", 1000, len(newSlice))
+	}
+}