From 174cc94076eca1bbd23a5bbdb14767720dbb9914 Mon Sep 17 00:00:00 2001
From: Francesco4203 <longngocpro@gmail.com>
Date: Mon, 8 Jul 2024 17:51:50 +0700
Subject: [PATCH] consortium-v2/snapshot: add pruneSnapshotPeriodically

pruneSnapshot: delete the nSnapshotsPrune oldest snapshots, keep the latestSnapshotsKeep snapshots
pruneSnapshotPeriodically: prune the snapshots at the start of each pruningPeriod
---
 consensus/consortium/v2/consortium.go |  7 ++++
 consensus/consortium/v2/snapshot.go   | 48 +++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/consensus/consortium/v2/consortium.go b/consensus/consortium/v2/consortium.go
index c2171e740..430b578b5 100644
--- a/consensus/consortium/v2/consortium.go
+++ b/consensus/consortium/v2/consortium.go
@@ -723,6 +723,9 @@ func (c *Consortium) snapshot(chain consensus.ChainHeaderReader, number uint64,
 			if err := snap.store(c.db); err != nil {
 				return nil, err
 			}
+			if err := snap.pruneSnapshotPeriodically(c.db, chain); err != nil {
+				return nil, err
+			}
 			log.Info("Stored checkpoint snapshot to disk", "number", number, "hash", hash)
 			figure.NewColorFigure("Welcome to DPOS", "", "green", true).Print()
 			break
@@ -782,6 +785,10 @@ func (c *Consortium) snapshot(chain consensus.ChainHeaderReader, number uint64,
 		if err = snap.store(c.db); err != nil {
 			return nil, err
 		}
+		// Prune the snapshot periodically
+		if err := snap.pruneSnapshotPeriodically(c.db, chain); err != nil {
+			return nil, err
+		}
 		log.Trace("Stored snapshot to disk", "number", snap.Number, "hash", snap.Hash)
 	}
 	log.Trace("Checking snapshot data", "number", snap.Number, "validators", snap.validators())
diff --git a/consensus/consortium/v2/snapshot.go b/consensus/consortium/v2/snapshot.go
index 2d1995531..cb0b40e6e 100644
--- a/consensus/consortium/v2/snapshot.go
+++ b/consensus/consortium/v2/snapshot.go
@@ -16,10 +16,22 @@ import (
 	blsCommon "github.com/ethereum/go-ethereum/crypto/bls/common"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/internal/ethapi"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/hashicorp/golang-lru/arc/v2"
 )
 
+const (
+	blocksPerEpoch  = 200
+	epochsPerPeriod = 144
+)
+
+var (
+	latestSnapshotsKeep = blocksPerEpoch * epochsPerPeriod * 5 // 5 days
+	snapshotsToBePruned = epochsPerPeriod * 2                  // 2 days
+	pruningPeriod       = blocksPerEpoch * epochsPerPeriod * 1 // every 1 day
+)
+
 // Snapshot is the state of the authorization validators at a given point in time.
 type Snapshot struct {
 	// private fields are not json.Marshalled
@@ -113,6 +125,42 @@ func loadSnapshot(
 	return snap, nil
 }
 
+// snapshot pruning
+// delete the nSnapshotsPrune oldest snapshots, keep the latestSnapshotsKeep snapshots
+func (s *Snapshot) pruneSnapshot(db ethdb.Database, nSnapshotPrune int, chain consensus.ChainHeaderReader) error {
+	log.Info("Pruning snapshots at block", "block", s.Number, "nSnapshotPrune", nSnapshotPrune)
+	// Get block number to start pruning
+	curBlockNumber := s.Number
+	curBlockNumber -= curBlockNumber % uint64(blocksPerEpoch) // start of the current epoch
+	curBlockNumber -= uint64(latestSnapshotsKeep)             // start of the oldest epoch to keep
+
+	// delete nSnapshotPrune snapshots starting from curBlockNumber to the older ones
+	batch := db.NewBatch()
+	for nSnapshotPrune > 0 {
+		nSnapshotPrune--
+		header := chain.GetHeaderByNumber(curBlockNumber)
+		if header == nil {
+			// no more snapshots to prune
+			break
+		}
+		curHash := header.Hash()
+		if err := batch.Delete(append(rawdb.ConsortiumSnapshotPrefix, curHash[:]...)); err != nil {
+			return err
+		}
+		curBlockNumber -= uint64(blocksPerEpoch)
+	}
+	log.Info("Pruned snapshots done")
+	return batch.Write()
+}
+
+// periodically prune the snapshots at the start of each pruningPeriod
+func (s *Snapshot) pruneSnapshotPeriodically(db ethdb.Database, chain consensus.ChainHeaderReader) error {
+	if s.Number%uint64(pruningPeriod) == 0 {
+		return s.pruneSnapshot(db, snapshotsToBePruned, chain)
+	}
+	return nil
+}
+
 // store inserts the snapshot into the database.
 func (s *Snapshot) store(db ethdb.Database) error {
 	blob, err := json.Marshal(s)