From d69d749fc69a40a7a15f554adfc4192a86e24684 Mon Sep 17 00:00:00 2001
From: krish-z <122767080+krish-nr@users.noreply.github.com>
Date: Wed, 18 Oct 2023 17:50:42 +0800
Subject: [PATCH] fix: pass a SnapshotOption func when init a new pruner (#11)

---
 cmd/geth/snapshot.go        | 13 ++++++++++---
 core/state/pruner/pruner.go | 18 ++++++++++++++----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/cmd/geth/snapshot.go b/cmd/geth/snapshot.go
index 11636d4732..665c89e8e2 100644
--- 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
index d50033fcc8..6d4b5c373f 100644
--- 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