From a27f7e0b8488fea8fef4b5eea1da890dfec8d3dc Mon Sep 17 00:00:00 2001
From: zhangkai <zhangkai.gis@163.com>
Date: Tue, 30 May 2023 19:35:22 +0800
Subject: [PATCH] support repair snapshot

---
 app/repair_state.go                               |  3 ++-
 cmd/okbchaind/repair_data.go                      |  4 ----
 libs/cosmos-sdk/store/mpt/snapshot.go             | 15 +++++++++++++--
 .../cosmos-sdk/store/rootmulti/rootmulti_store.go | 15 ++-------------
 libs/tendermint/global/status.go                  | 11 +++++++++++
 5 files changed, 28 insertions(+), 20 deletions(-)
 create mode 100644 libs/tendermint/global/status.go

diff --git a/app/repair_state.go b/app/repair_state.go
index 8811e27e4..f3931c2ae 100644
--- a/app/repair_state.go
+++ b/app/repair_state.go
@@ -89,7 +89,8 @@ func repairStateOnStart(ctx *server.Context) {
 func RepairState(ctx *server.Context, onStart bool) {
 	sm.SetIgnoreSmbCheck(true)
 	iavl.SetIgnoreVersionCheck(true)
-	rootmulti.SetRepair()
+	global.SetRepairState(true)
+	defer global.SetRepairState(false)
 
 	// load latest block height
 	dataDir := filepath.Join(ctx.Config.RootDir, "data")
diff --git a/cmd/okbchaind/repair_data.go b/cmd/okbchaind/repair_data.go
index 15c11c43f..366da6b4d 100644
--- a/cmd/okbchaind/repair_data.go
+++ b/cmd/okbchaind/repair_data.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"fmt"
-	"github.com/okx/okbchain/libs/cosmos-sdk/store/mpt"
 	"log"
 	"net/http"
 	_ "net/http/pprof"
@@ -64,7 +63,4 @@ func setExternalPackageValue() {
 		tmiavl.SetEnableFastStorage(true)
 		tmiavl.SetIgnoreAutoUpgrade(true)
 	}
-	if !viper.GetBool(tmiavl.FlagIavlDiscardFastStorage) {
-		mpt.SetSnapshotRebuild(true)
-	}
 }
diff --git a/libs/cosmos-sdk/store/mpt/snapshot.go b/libs/cosmos-sdk/store/mpt/snapshot.go
index 45e52d141..2826fbc83 100644
--- a/libs/cosmos-sdk/store/mpt/snapshot.go
+++ b/libs/cosmos-sdk/store/mpt/snapshot.go
@@ -6,6 +6,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	mpttypes "github.com/okx/okbchain/libs/cosmos-sdk/store/mpt/types"
+	"github.com/okx/okbchain/libs/tendermint/global"
 )
 
 var (
@@ -15,8 +16,10 @@ var (
 
 const (
 	// snapshotMemoryLayerCount snapshot memory layer count
-	// as we dont rollback transactions so we only keep 1 memory layer
-	snapshotMemoryLayerCount = 1
+	// snapshotMemoryLayerCount controls the snapshot Journal height,
+	// if repair start-height is lower than snapshot Journal height,
+	// snapshot will not be repaired anymore
+	snapshotMemoryLayerCount = 10
 )
 
 func DisableSnapshot() {
@@ -42,6 +45,9 @@ func (ms *MptStore) openSnapshot() error {
 		ms.logger.Error("Enabling snapshot recovery", "chainhead", version, "diskbase", *layer)
 		recovery = true
 	}
+	if global.GetRepairState() {
+		recovery = true
+	}
 	var err error
 	ms.snaps, err = snapshot.NewCustom(ms.db.TrieDB().DiskDB(), ms.db.TrieDB(), 256, ms.originalRoot, false, gSnapshotRebuild, recovery, ms.retriever)
 	if err != nil {
@@ -74,6 +80,8 @@ func (ms *MptStore) prepareSnap(root common.Hash) {
 		ms.snapDestructs = make(map[common.Hash]struct{})
 		ms.snapAccounts = make(map[common.Hash][]byte)
 		ms.snapStorage = make(map[common.Hash]map[common.Hash][]byte)
+	} else {
+		ms.logger.Error("prepare snapshot error", "root", root)
 	}
 }
 
@@ -94,6 +102,9 @@ func (ms *MptStore) commitSnap(root common.Hash) {
 		if err := ms.snaps.Cap(root, snapshotMemoryLayerCount); err != nil {
 			ms.logger.Error("Failed to cap snapshot tree", "root", root, "layers", snapshotMemoryLayerCount, "err", err)
 		}
+		if _, err := ms.snaps.Journal(root); err != nil {
+			ms.logger.Error("Failed to journal snapshot tree", "root", root, "err", err)
+		}
 	}
 	ms.snap, ms.snapDestructs, ms.snapAccounts, ms.snapStorage = nil, nil, nil, nil
 
diff --git a/libs/cosmos-sdk/store/rootmulti/rootmulti_store.go b/libs/cosmos-sdk/store/rootmulti/rootmulti_store.go
index 0ef7d61b8..c757c945f 100644
--- a/libs/cosmos-sdk/store/rootmulti/rootmulti_store.go
+++ b/libs/cosmos-sdk/store/rootmulti/rootmulti_store.go
@@ -3,6 +3,7 @@ package rootmulti
 import (
 	"encoding/binary"
 	"fmt"
+	"github.com/okx/okbchain/libs/tendermint/global"
 	"io"
 	"log"
 	"path/filepath"
@@ -47,18 +48,6 @@ const (
 	maxPruneHeightsLength = 100
 )
 
-var (
-	repairing bool
-)
-
-func SetRepair() {
-	repairing = true
-}
-
-func getRepair() bool {
-	return repairing
-}
-
 // Store is composed of many CommitStores. Name contrasts with
 // cacheMultiStore which is for cache-wrapping other MultiStores. It implements
 // the CommitMultiStore interface.
@@ -408,7 +397,7 @@ func (rs *Store) loadVersion(ver int64, upgrades *types.StoreUpgrades) error {
 						// we can not get the upgrade version before the expect height,
 						// and we should not use the original 0 too, because 0 means the latest height,
 						// so when we repair data before the milestone. we open a empty tree by cur version.
-						if getRepair() && version == 0 {
+						if global.GetRepairState() && version == 0 {
 							param.upgradeVersion = uint64(ver)
 						}
 						rs.storesParams[key] = param
diff --git a/libs/tendermint/global/status.go b/libs/tendermint/global/status.go
new file mode 100644
index 000000000..8a1416681
--- /dev/null
+++ b/libs/tendermint/global/status.go
@@ -0,0 +1,11 @@
+package global
+
+var repairState bool
+
+func SetRepairState(state bool) {
+	repairState = state
+}
+
+func GetRepairState() bool {
+	return repairState
+}