some changes to memory mgmt when using mempool

prysmaticlabs · Mar 29, 2020 · f5b3e7f · f5b3e7f
1 parent b297cb1
commit f5b3e7f
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 13 deletions.
diff --git a/beacon-chain/state/field_trie.go b/beacon-chain/state/field_trie.go
@@ -117,7 +117,7 @@ func (f *FieldTrie) CopyTrie() *FieldTrie {
 	case stateRoots:
 		dstFieldTrie = memorypool.GetStateRootsTrie(len(f.fieldLayers))
 	default:
-		dstFieldTrie = make([][]*[32]byte, len(f.fieldLayers))
+		dstFieldTrie = memorypool.GetFieldTrie(len(f.fieldLayers))
 	}
 
 	for i, layer := range f.fieldLayers {

diff --git a/beacon-chain/state/getters.go b/beacon-chain/state/getters.go
@@ -222,7 +222,7 @@ func (b *BeaconState) BlockRoots() [][]byte {
 	if b.state.BlockRoots == nil {
 		return nil
 	}
-	roots := make([][]byte, len(b.state.BlockRoots))
+	roots := memorypool.GetDoubleByteSlice(len(b.state.BlockRoots))
 	for i, r := range b.state.BlockRoots {
 		tmpRt := make([]byte, len(r))
 		copy(tmpRt, r)
@@ -264,7 +264,7 @@ func (b *BeaconState) StateRoots() [][]byte {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
-	roots := make([][]byte, len(b.state.StateRoots))
+	roots := memorypool.GetDoubleByteSlice(len(b.state.StateRoots))
 	for i, r := range b.state.StateRoots {
 		tmpRt := make([]byte, len(r))
 		copy(tmpRt, r)
@@ -285,7 +285,7 @@ func (b *BeaconState) HistoricalRoots() [][]byte {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
-	roots := make([][]byte, len(b.state.HistoricalRoots))
+	roots := memorypool.GetDoubleByteSlice(len(b.state.HistoricalRoots))
 	for i, r := range b.state.HistoricalRoots {
 		tmpRt := make([]byte, len(r))
 		copy(tmpRt, r)
@@ -563,7 +563,7 @@ func (b *BeaconState) RandaoMixes() [][]byte {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
-	mixes := memorypool.GetDoubleByteSlice(len(b.state.RandaoMixes))
+	mixes := memorypool.GetRandoMixes(len(b.state.RandaoMixes))
 	for i, r := range b.state.RandaoMixes {
 		tmpRt := make([]byte, len(r))
 		copy(tmpRt, r)

diff --git a/beacon-chain/state/state_trie.go b/beacon-chain/state/state_trie.go
@@ -158,18 +158,30 @@ func (b *BeaconState) Copy() *BeaconState {
 			v.refs--
 			if b.stateFieldLeaves[field].reference != nil {
 				b.stateFieldLeaves[field].MinusRef()
+				if b.stateFieldLeaves[field].refs == 0 && field != randaoMixes && field != blockRoots && field != stateRoots && field != historicalRoots {
+					memorypool.PutFieldTrie(b.stateFieldLeaves[field].fieldLayers)
+				}
 			}
 			if field == randaoMixes && v.refs == 0 {
-				memorypool.PutDoubleByteSlice(b.state.RandaoMixes)
+				memorypool.PutRandaoMixes(b.state.RandaoMixes)
 				if b.stateFieldLeaves[field].refs == 0 {
 					memorypool.PutRandaoMixesTrie(b.stateFieldLeaves[randaoMixes].fieldLayers)
 				}
 			}
-			if field == blockRoots && v.refs == 0 && b.stateFieldLeaves[field].refs == 0 {
-				memorypool.PutBlockRootsTrie(b.stateFieldLeaves[blockRoots].fieldLayers)
+			if field == blockRoots && v.refs == 0 {
+				memorypool.PutDoubleByteSlice(b.state.BlockRoots)
+				if b.stateFieldLeaves[field].refs == 0 {
+					memorypool.PutBlockRootsTrie(b.stateFieldLeaves[blockRoots].fieldLayers)
+				}
+			}
+			if field == stateRoots && v.refs == 0 {
+				memorypool.PutDoubleByteSlice(b.state.StateRoots)
+				if b.stateFieldLeaves[field].refs == 0 {
+					memorypool.PutStateRootsTrie(b.stateFieldLeaves[stateRoots].fieldLayers)
+				}
 			}
-			if field == stateRoots && v.refs == 0 && b.stateFieldLeaves[field].refs == 0 {
-				memorypool.PutStateRootsTrie(b.stateFieldLeaves[stateRoots].fieldLayers)
+			if field == historicalRoots && v.refs == 0 {
+				memorypool.PutDoubleByteSlice(b.state.HistoricalRoots)
 			}
 		}
 	})

diff --git a/shared/memorypool/memorypool.go b/shared/memorypool/memorypool.go
@@ -22,6 +22,8 @@ var StateRootsMemoryPool = new(sync.Pool)
 // for randao mixes trie.
 var RandaoMixesMemoryPool = new(sync.Pool)
 
+var GenericFieldFieldTriePool = new(sync.Pool)
+
 // GetDoubleByteSlice retrieves the 2d byte slice of
 // the desired size from the memory pool.
 func GetDoubleByteSlice(size int) [][]byte {
@@ -48,6 +50,29 @@ func PutDoubleByteSlice(data [][]byte) {
 	}
 }
 
+func GetRandoMixes(size int) [][]byte {
+	if !featureconfig.Get().EnableByteMempool {
+		return make([][]byte, size)
+	}
+	rawObj := RandaoMixesMemoryPool.Get()
+	if rawObj == nil {
+		return GetDoubleByteSlice(size)
+	}
+
+	byteSlice := rawObj.([][]byte)
+	if len(byteSlice) >= size {
+		return byteSlice[:size]
+	}
+	return append(byteSlice, make([][]byte, size-len(byteSlice))...)
+}
+
+func PutRandaoMixes(data [][]byte) {
+	if !featureconfig.Get().EnableByteMempool {
+		return
+	}
+	RandaoMixesMemoryPool.Put(data)
+}
+
 // GetBlockRootsTrie retrieves the 3d byte trie of
 // the desired size from the memory pool.
 func GetBlockRootsTrie(size int) [][]*[32]byte {
@@ -57,7 +82,7 @@ func GetBlockRootsTrie(size int) [][]*[32]byte {
 
 	rawObj := BlockRootsMemoryPool.Get()
 	if rawObj == nil {
-		return make([][]*[32]byte, size)
+		return GetFieldTrie(size)
 	}
 	byteSlice := rawObj.([][]*[32]byte)
 	if len(byteSlice) >= size {
@@ -83,7 +108,7 @@ func GetStateRootsTrie(size int) [][]*[32]byte {
 
 	rawObj := BlockRootsMemoryPool.Get()
 	if rawObj == nil {
-		return make([][]*[32]byte, size)
+		return GetFieldTrie(size)
 	}
 	byteSlice := rawObj.([][]*[32]byte)
 	if len(byteSlice) >= size {
@@ -109,7 +134,7 @@ func GetRandaoMixesTrie(size int) [][]*[32]byte {
 
 	rawObj := StateRootsMemoryPool.Get()
 	if rawObj == nil {
-		return make([][]*[32]byte, size)
+		return GetFieldTrie(size)
 	}
 	byteSlice := rawObj.([][]*[32]byte)
 	if len(byteSlice) >= size {
@@ -125,3 +150,25 @@ func PutRandaoMixesTrie(data [][]*[32]byte) {
 		StateRootsMemoryPool.Put(data)
 	}
 }
+
+func GetFieldTrie(size int) [][]*[32]byte {
+	if !featureconfig.Get().EnableByteMempool {
+		return make([][]*[32]byte, size)
+	}
+
+	rawObj := GenericFieldFieldTriePool.Get()
+	if rawObj == nil {
+		return make([][]*[32]byte, size)
+	}
+	byteSlice := rawObj.([][]*[32]byte)
+	if len(byteSlice) >= size {
+		return byteSlice[:size]
+	}
+	return append(byteSlice, make([][]*[32]byte, size-len(byteSlice))...)
+}
+
+func PutFieldTrie(data [][]*[32]byte) {
+	if featureconfig.Get().EnableByteMempool {
+		GenericFieldFieldTriePool.Put(data)
+	}
+}