fix: pass a SnapshotOption func when init a new pruner (bnb-chain#11)

galaio · Oct 18, 2023 · d69d749 · d69d749
1 parent 38bc840
commit d69d749
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/cmd/geth/snapshot.go b/cmd/geth/snapshot.go
@@ -23,6 +23,8 @@ import (
 	"os"
 	"time"
 
+	cli "github.com/urfave/cli/v2"
+
 	"github.com/ethereum/go-ethereum/cmd/utils"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
@@ -35,7 +37,6 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
-	cli "github.com/urfave/cli/v2"
 )
 
 var (
@@ -173,8 +174,14 @@ func pruneState(ctx *cli.Context) error {
 		Cachedir:  stack.ResolvePath(config.Eth.TrieCleanCacheJournal),
 		BloomSize: ctx.Uint64(utils.BloomFilterSizeFlag.Name),
 	}
-	pruner, err := pruner.NewPruner(chaindb, prunerconfig,
-		pruner.WithTriesInMemory(ctx.Uint64(utils.TriesInMemoryFlag.Name)))
+	pruner, err := pruner.NewPruner(chaindb, prunerconfig, pruner.CombinedOptions{
+		PrunerOptions: []pruner.PrunerOption{
+			pruner.WithTriesInMemory(ctx.Uint64(utils.TriesInMemoryFlag.Name)),
+		},
+		SnapshotOptions: []snapshot.SnapshotOption{
+			snapshot.SetCapLimit(int(ctx.Uint64(utils.TriesInMemoryFlag.Name))),
+		},
+	})
 	if err != nil {
 		log.Error("Failed to open snapshot tree", "err", err)
 		return err

diff --git a/core/state/pruner/pruner.go b/core/state/pruner/pruner.go
@@ -92,8 +92,13 @@ type Pruner struct {
 	triesInMemory uint64
 }
 
+type CombinedOptions struct {
+	PrunerOptions   []PrunerOption
+	SnapshotOptions []snapshot.SnapshotOption
+}
+
 // NewPruner creates the pruner instance.
-func NewPruner(db ethdb.Database, config Config, opts ...PrunerOption) (*Pruner, error) {
+func NewPruner(db ethdb.Database, config Config, opts CombinedOptions) (*Pruner, error) {
 	headBlock := rawdb.ReadHeadBlock(db)
 	if headBlock == nil {
 		return nil, errors.New("failed to load head block")
@@ -105,6 +110,11 @@ func NewPruner(db ethdb.Database, config Config, opts ...PrunerOption) (*Pruner,
 		AsyncBuild: false,
 	}
 	snaptree, err := snapshot.New(snapconfig, db, trie.NewDatabase(db), headBlock.Root())
+	for _, snapshotOption := range opts.SnapshotOptions {
+		if snapshotOption != nil {
+			snapshotOption(snaptree)
+		}
+	}
 	if err != nil {
 		return nil, err // The relevant snapshot(s) might not exist
 	}
@@ -125,9 +135,9 @@ func NewPruner(db ethdb.Database, config Config, opts ...PrunerOption) (*Pruner,
 		snaptree:      snaptree,
 		triesInMemory: core.TriesInMemory,
 	}
-	for _, opt := range opts {
-		if opt != nil {
-			opt(pruner)
+	for _, prunerOption := range opts.PrunerOptions {
+		if prunerOption != nil {
+			prunerOption(pruner)
 		}
 	}
 	return pruner, nil