Added faster pruneBlocks implementation

p2pool/sidechain/sidechain.go

@@ -71,8 +71,10 @@ type SideChain struct {
 	watchBlock            *ChainMain
 	watchBlockSidechainId types.Hash
 
-	blocksByTemplateId *swiss.Map[types.Hash, *PoolBlock]
-	blocksByHeight     *swiss.Map[uint64, []*PoolBlock]
+	blocksByTemplateId       *swiss.Map[types.Hash, *PoolBlock]
+	blocksByHeight           *swiss.Map[uint64, []*PoolBlock]
+	blocksByHeightKeysSorted bool
+	blocksByHeightKeys       []uint64
 
 	preAllocatedBuffer []byte
 
@@ -425,6 +427,10 @@ func (c *SideChain) AddPoolBlock(block *PoolBlock) (err error) {
 		c.blocksByHeight.Put(block.Side.Height, append(l, block))
 	} else {
 		c.blocksByHeight.Put(block.Side.Height, []*PoolBlock{block})
+		if !(c.blocksByHeightKeysSorted && len(c.blocksByHeightKeys) > 0 && c.blocksByHeightKeys[len(c.blocksByHeightKeys)-1]+1 == block.Side.Height) {
+			c.blocksByHeightKeysSorted = false
+		}
+		c.blocksByHeightKeys = append(c.blocksByHeightKeys, block.Side.Height)
 	}
 
 	c.updateDepths(block)
@@ -882,12 +888,21 @@ func (c *SideChain) pruneOldBlocks() {
 
 	h := tip.Side.Height - pruneDistance
 
+	if !c.blocksByHeightKeysSorted {
+		slices.Sort(c.blocksByHeightKeys)
+		c.blocksByHeightKeysSorted = true
+	}
+
 	numBlocksPruned := 0
-	c.blocksByHeight.Iter(func(height uint64, v []*PoolBlock) (stop bool) {
+
+	for keyIndex, height := range c.blocksByHeightKeys {
+		// Early exit
 		if height > h {
-			return false
+			break
 		}
 
+		v, _ := c.blocksByHeight.Get(height)
+
 		// loop backwards for proper deletions
 		for i := len(v) - 1; i >= 0; i-- {
 			block := v[i]
@@ -908,11 +923,11 @@ func (c *SideChain) pruneOldBlocks() {
 
 		if len(v) == 0 {
 			c.blocksByHeight.Delete(height)
+			c.blocksByHeightKeys = slices.Delete(c.blocksByHeightKeys, keyIndex, keyIndex+1)
 		} else {
 			c.blocksByHeight.Put(height, v)
 		}
-		return false
-	})
+	}
 
 	if numBlocksPruned > 0 {
 		log.Printf("[SideChain] pruned %d old blocks at heights <= %d", numBlocksPruned, h)

p2pool/sidechain/sidechain_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path"
 	"runtime"
+	"slices"
 	"testing"
 )
 
@@ -133,9 +134,11 @@ func TestSideChainMiniPreFork(t *testing.T) {
 	testSideChain(s, t, f, 2424349, 2696040, block2420028, block2420027)
 }
 
-func benchmarkResetState(tip, parent *PoolBlock, templateId types.Hash, fullId FullId, difficulty types.Difficulty, s *SideChain) {
+func benchmarkResetState(tip, parent *PoolBlock, templateId types.Hash, fullId FullId, difficulty types.Difficulty, blocksByHeightKeys []uint64, s *SideChain) {
 	//Remove states in maps
 	s.blocksByHeight.Delete(tip.Side.Height)
+	s.blocksByHeightKeys = blocksByHeightKeys
+	s.blocksByHeightKeysSorted = true
 	s.blocksByTemplateId.Delete(templateId)
 	s.seenBlocks.Delete(fullId)
 
@@ -160,6 +163,9 @@ func benchSideChain(b *testing.B, s *SideChain, tipHash types.Hash) {
 
 	for tip.SideTemplateId(s.Consensus()) != tipHash {
 		s.blocksByHeight.Delete(tip.Side.Height)
+		s.blocksByHeightKeys = slices.DeleteFunc(s.blocksByHeightKeys, func(u uint64) bool {
+			return u == tip.Side.Height
+		})
 		s.blocksByTemplateId.Delete(tip.SideTemplateId(s.Consensus()))
 		s.seenBlocks.Delete(tip.FullId())
 
@@ -176,13 +182,20 @@ func benchSideChain(b *testing.B, s *SideChain, tipHash types.Hash) {
 
 	difficulty, _, _ := s.GetDifficulty(parent)
 
-	benchmarkResetState(tip, parent, templateId, fullId, difficulty, s)
+	slices.Sort(s.blocksByHeightKeys)
+	blocksByHeightKeys := slices.DeleteFunc(s.blocksByHeightKeys, func(u uint64) bool {
+		return u == tip.Side.Height
+	})
+
+	benchmarkResetState(tip, parent, templateId, fullId, difficulty, slices.Clone(blocksByHeightKeys), s)
 
 	var err error
 
 	b.StartTimer()
 	for i := 0; i < b.N; i++ {
-		benchmarkResetState(tip, parent, templateId, fullId, difficulty, s)
+		b.StopTimer()
+		benchmarkResetState(tip, parent, templateId, fullId, difficulty, slices.Clone(blocksByHeightKeys), s)
+		b.StartTimer()
 		_, err, _ = s.AddPoolBlockExternal(tip)
 		if err != nil {
 			b.Error(err)