From 6e3fab854fb8bf30692049fef0456bde37c3c22d Mon Sep 17 00:00:00 2001
From: vincent <vincent.ch.cn@gmail.com>
Date: Fri, 13 Nov 2020 23:27:52 +0800
Subject: [PATCH] Fix upgrade store loader (#7817)

* fix upgrade store loader

* fix test

* resolve comment

Co-authored-by: billy rennekamp <billy.rennekamp@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 store/rootmulti/store.go            |  4 +++-
 x/upgrade/types/storeloader.go      |  4 ++--
 x/upgrade/types/storeloader_test.go | 29 +++++++++++++++++++++--------
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index 90633a922..72238cbb1 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -315,7 +315,9 @@ func (rs *Store) TracingEnabled() bool {
 // LastCommitID implements Committer/CommitStore.
 func (rs *Store) LastCommitID() types.CommitID {
 	if rs.lastCommitInfo == nil {
-		return types.CommitID{}
+		return types.CommitID{
+			Version: getLatestVersion(rs.db),
+		}
 	}
 
 	return rs.lastCommitInfo.CommitID()
diff --git a/x/upgrade/types/storeloader.go b/x/upgrade/types/storeloader.go
index 0ff168dcc..00538e071 100644
--- a/x/upgrade/types/storeloader.go
+++ b/x/upgrade/types/storeloader.go
@@ -10,9 +10,9 @@ import (
 // pattern. This is useful for custom upgrade loading logic.
 func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) baseapp.StoreLoader {
 	return func(ms sdk.CommitMultiStore) error {
-		if upgradeHeight == ms.LastCommitID().Version {
+		if upgradeHeight == ms.LastCommitID().Version+1 {
 			// Check if the current commit version and upgrade height matches
-			if len(storeUpgrades.Renamed) > 0 || len(storeUpgrades.Deleted) > 0 {
+			if len(storeUpgrades.Renamed) > 0 || len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
 				return ms.LoadLatestVersionAndUpgrade(storeUpgrades)
 			}
 		}
diff --git a/x/upgrade/types/storeloader_test.go b/x/upgrade/types/storeloader_test.go
index f1afa4176..ec2bfa824 100644
--- a/x/upgrade/types/storeloader_test.go
+++ b/x/upgrade/types/storeloader_test.go
@@ -65,11 +65,13 @@ func checkStore(t *testing.T, db dbm.DB, ver int64, storeKey string, k, v []byte
 // Test that we can make commits and then reload old versions.
 // Test that LoadLatestVersion actually does.
 func TestSetLoader(t *testing.T) {
+	upgradeHeight := int64(5)
+
 	// set a temporary home dir
 	homeDir := t.TempDir()
 	upgradeInfoFilePath := filepath.Join(homeDir, "upgrade-info.json")
 	upgradeInfo := &store.UpgradeInfo{
-		Name: "test", Height: 0,
+		Name: "test", Height: upgradeHeight,
 	}
 
 	data, err := json.Marshal(upgradeInfo)
@@ -92,7 +94,7 @@ func TestSetLoader(t *testing.T) {
 			loadStoreKey: "foo",
 		},
 		"rename with inline opts": {
-			setLoader: useUpgradeLoader(0, &store.StoreUpgrades{
+			setLoader: useUpgradeLoader(upgradeHeight, &store.StoreUpgrades{
 				Renamed: []store.StoreRename{{
 					OldKey: "foo",
 					NewKey: "bar",
@@ -116,25 +118,36 @@ func TestSetLoader(t *testing.T) {
 
 			// load the app with the existing db
 			opts := []func(*baseapp.BaseApp){baseapp.SetPruning(store.PruneNothing)}
+
+			origapp := baseapp.NewBaseApp(t.Name(), defaultLogger(), db, nil, opts...)
+			origapp.MountStores(sdk.NewKVStoreKey(tc.origStoreKey))
+			err := origapp.LoadLatestVersion()
+			require.Nil(t, err)
+
+			for i := int64(2); i <= upgradeHeight-1; i++ {
+				origapp.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: i}})
+				res := origapp.Commit()
+				require.NotNil(t, res.Data)
+			}
+
 			if tc.setLoader != nil {
 				opts = append(opts, tc.setLoader)
 			}
 
+			// load the new app with the original app db
 			app := baseapp.NewBaseApp(t.Name(), defaultLogger(), db, nil, opts...)
-			capKey := sdk.NewKVStoreKey("main")
-			app.MountStores(capKey)
 			app.MountStores(sdk.NewKVStoreKey(tc.loadStoreKey))
-			err := app.LoadLatestVersion()
+			err = app.LoadLatestVersion()
 			require.Nil(t, err)
 
 			// "execute" one block
-			app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: 2}})
+			app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: upgradeHeight}})
 			res := app.Commit()
 			require.NotNil(t, res.Data)
 
 			// check db is properly updated
-			checkStore(t, db, 2, tc.loadStoreKey, k, v)
-			checkStore(t, db, 2, tc.loadStoreKey, []byte("foo"), nil)
+			checkStore(t, db, upgradeHeight, tc.loadStoreKey, k, v)
+			checkStore(t, db, upgradeHeight, tc.loadStoreKey, []byte("foo"), nil)
 		})
 	}
 }