From 439f2f9d5b5884bc9df4b58d702555330549a898 Mon Sep 17 00:00:00 2001
From: Marko <marko@baricevic.me>
Date: Thu, 6 Jun 2024 15:01:41 +0200
Subject: [PATCH] feat: upstream changes from server modular  (#20584)

---
 store/v2/commitment/store.go | 41 ++++++++++++++++++++++++++----------
 store/v2/database.go         |  3 +++
 store/v2/root/store.go       |  7 +++++-
 x/upgrade/keeper/abci.go     |  6 +-----
 4 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/store/v2/commitment/store.go b/store/v2/commitment/store.go
index e9cf2c0c63..b67a1149c3 100644
--- a/store/v2/commitment/store.go
+++ b/store/v2/commitment/store.go
@@ -113,6 +113,19 @@ func (c *CommitStore) GetLatestVersion() (uint64, error) {
 	return version, nil
 }
 
+// IsEmpty returns true if the CommitStore is empty.
+func (c *CommitStore) IsEmpty() (bool, error) {
+	value, err := c.db.Get([]byte(latestVersionKey))
+	if err != nil {
+		return false, err
+	}
+	if value == nil {
+		return true, nil
+	} else {
+		return false, nil
+	}
+}
+
 func (c *CommitStore) LoadVersion(targetVersion uint64) error {
 	// Rollback the metadata to the target version.
 	latestVersion, err := c.GetLatestVersion()
@@ -167,16 +180,19 @@ func (c *CommitStore) GetCommitInfo(version uint64) (*proof.CommitInfo, error) {
 }
 
 func (c *CommitStore) flushCommitInfo(version uint64, cInfo *proof.CommitInfo) error {
+	// do nothing if commit info is nil, as will be the case for an empty, initializing store
+	if cInfo == nil {
+		return nil
+	}
+
 	batch := c.db.NewBatch()
-	if cInfo != nil {
-		cInfoKey := []byte(fmt.Sprintf(commitInfoKeyFmt, version))
-		value, err := cInfo.Marshal()
-		if err != nil {
-			return err
-		}
-		if err := batch.Set(cInfoKey, value); err != nil {
-			return err
-		}
+	cInfoKey := []byte(fmt.Sprintf(commitInfoKeyFmt, version))
+	value, err := cInfo.Marshal()
+	if err != nil {
+		return err
+	}
+	if err := batch.Set(cInfoKey, value); err != nil {
+		return err
 	}
 
 	var buf bytes.Buffer
@@ -201,8 +217,11 @@ func (c *CommitStore) Commit(version uint64) (*proof.CommitInfo, error) {
 		// If a commit event execution is interrupted, a new iavl store's version
 		// will be larger than the RMS's metadata, when the block is replayed, we
 		// should avoid committing that iavl store again.
-		var commitID proof.CommitID
-		if tree.GetLatestVersion() >= version {
+		var (
+			commitID      proof.CommitID
+			latestVersion = tree.GetLatestVersion()
+		)
+		if latestVersion != 0 && latestVersion >= version {
 			commitID.Version = version
 			commitID.Hash = tree.Hash()
 		} else {
diff --git a/store/v2/database.go b/store/v2/database.go
index a0466de18d..ae235cdab4 100644
--- a/store/v2/database.go
+++ b/store/v2/database.go
@@ -57,6 +57,9 @@ type Committer interface {
 	// GetCommitInfo returns the CommitInfo for the given version.
 	GetCommitInfo(version uint64) (*proof.CommitInfo, error)
 
+	// IsEmpty returns true if the database is empty.
+	IsEmpty() (bool, error)
+
 	// Close releases associated resources. It should NOT be idempotent. It must
 	// only be called once and any call after may panic.
 	io.Closer
diff --git a/store/v2/root/store.go b/store/v2/root/store.go
index 8c77cd0d75..e5543a91a2 100644
--- a/store/v2/root/store.go
+++ b/store/v2/root/store.go
@@ -388,8 +388,13 @@ func (s *Store) writeSC(cs *corestore.Changeset) error {
 		return fmt.Errorf("failed to write batch to SC store: %w", err)
 	}
 
+	isEmpty, err := s.stateCommitment.IsEmpty()
+	if err != nil {
+		return fmt.Errorf("failed to check if SC store is empty: %w", err)
+	}
+
 	var previousHeight, version uint64
-	if s.lastCommitInfo.GetVersion() == 0 && s.initialVersion > 1 {
+	if isEmpty {
 		// This case means that no commit has been made in the store, we
 		// start from initialVersion.
 		version = s.initialVersion
diff --git a/x/upgrade/keeper/abci.go b/x/upgrade/keeper/abci.go
index ce1a418ea2..e6bed2f5e6 100644
--- a/x/upgrade/keeper/abci.go
+++ b/x/upgrade/keeper/abci.go
@@ -5,12 +5,10 @@ import (
 	"errors"
 	"fmt"
 
-	storetypes "cosmossdk.io/store/types"
 	consensusv1 "cosmossdk.io/x/consensus/types"
 	"cosmossdk.io/x/upgrade/types"
 
 	"github.com/cosmos/cosmos-sdk/telemetry"
-	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
 // PreBlocker will check if there is a scheduled plan and if it is ready to be executed.
@@ -31,7 +29,6 @@ func (k Keeper) PreBlocker(ctx context.Context) error {
 	}
 	found := err == nil
 
-	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO remove with consensus messages
 	if !k.DowngradeVerified() {
 		k.SetDowngradeVerified(true)
 		// This check will make sure that we are using a valid binary.
@@ -97,8 +94,7 @@ func (k Keeper) PreBlocker(ctx context.Context) error {
 
 		// We have an upgrade handler for this upgrade name, so apply the upgrade
 		k.Logger.Info(fmt.Sprintf("applying upgrade \"%s\" at %s", plan.Name, plan.DueAt()))
-		sdkCtx = sdkCtx.WithBlockGasMeter(storetypes.NewInfiniteGasMeter())
-		if err := k.ApplyUpgrade(sdkCtx, plan); err != nil {
+		if err := k.ApplyUpgrade(ctx, plan); err != nil {
 			return err
 		}
 		return nil