plt-224 memoize GetDoneHeight in upgrade keeper to avoid repeated KV reads

sei-cosmos/x/upgrade/keeper/keeper.go

@@ -8,6 +8,7 @@ import (
 	"path"
 	"path/filepath"
 	"sort"
+	"sync"
 
 	"github.com/armon/go-metrics"
 	tmos "github.com/sei-protocol/sei-chain/sei-tendermint/libs/os"
@@ -40,6 +41,9 @@ type Keeper struct {
 	upgradeHandlers    map[string]types.UpgradeHandler // map of plan name to upgrade handler
 	versionSetter      xp.ProtocolVersionSetter        // implements setting the protocol version field on BaseApp
 	downgradeVerified  bool                            // tells if we've already sanity checked that this binary version isn't being used against an old state.
+	// doneHeightCache is a pointer so that value-receiver copies of Keeper all share the same underlying map.
+	// Done-heights are immutable once set, so cached values are always valid for the lifetime of the process.
+	doneHeightCache *sync.Map // upgrade name -> int64 done height
 }
 
 // NewKeeper constructs an upgrade Keeper which requires the following arguments:
@@ -56,6 +60,7 @@ func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey sdk.StoreKey, cdc cod
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
 		versionSetter:      vs,
+		doneHeightCache:    &sync.Map{},
 	}
 }
 
@@ -278,13 +283,17 @@ func parseDoneKey(key []byte) string {
 
 // GetDoneHeight returns the height at which the given upgrade was executed
 func (k Keeper) GetDoneHeight(ctx sdk.Context, name string) int64 {
+	if v, ok := k.doneHeightCache.Load(name); ok {
+		return v.(int64)
+	}
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte{types.DoneByte})
 	bz := store.Get([]byte(name))
 	if len(bz) == 0 {
 		return 0
 	}
-
-	return int64(binary.BigEndian.Uint64(bz)) //nolint:gosec // stored by SetDone from block heights which are always non-negative
+	height := int64(binary.BigEndian.Uint64(bz)) //nolint:gosec // stored by SetDone from block heights which are always non-negative
+	k.doneHeightCache.Store(name, height)
+	return height
 }
 
 func (k Keeper) GetClosestUpgrade(ctx sdk.Context, height int64) (string, int64) {
@@ -346,6 +355,7 @@ func (k Keeper) SetDone(ctx sdk.Context, name string) {
 	bz := make([]byte, 8)
 	binary.BigEndian.PutUint64(bz, uint64(ctx.BlockHeight())) //nolint:gosec // block heights are always non-negative
 	store.Set([]byte(name), bz)
+	k.doneHeightCache.Store(name, ctx.BlockHeight())
 }
 
 // HasHandler returns true iff there is a handler registered for this name

sei-cosmos/x/upgrade/keeper/keeper_bench_test.go

@@ -0,0 +1,59 @@
+package keeper_test
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	tmproto "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/types"
+
+	seiapp "github.com/sei-protocol/sei-chain/app"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/upgrade/keeper"
+	"github.com/sei-protocol/sei-chain/sei-cosmos/x/upgrade/types"
+)
+
+// BenchmarkGetDoneHeightCached measures GetDoneHeight with a warm in-memory cache
+// (current behavior after the fix). Every iteration is a sync.Map lookup.
+func BenchmarkGetDoneHeightCached(b *testing.B) {
+	app := seiapp.Setup(b, false, false, false)
+	ctx := app.BaseApp.NewContext(false, tmproto.Header{Height: 100, Time: time.Now()})
+	app.UpgradeKeeper.SetDone(ctx, "test-upgrade")
+	// Warm the cache with one read before measuring.
+	app.UpgradeKeeper.GetDoneHeight(ctx, "test-upgrade")
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		app.UpgradeKeeper.GetDoneHeight(ctx, "test-upgrade")
+	}
+}
+
+// BenchmarkGetDoneHeightUncached measures GetDoneHeight with a cold cache,
+// simulating the pre-fix behavior where every call hits the KV store.
+// A unique upgrade name is used per iteration so the cache is always cold,
+// but the value exists in the KV store — reproducing the pre-fix hot path.
+func BenchmarkGetDoneHeightUncached(b *testing.B) {
+	app := seiapp.Setup(b, false, false, false)
+	ctx := app.BaseApp.NewContext(false, tmproto.Header{Height: 100, Time: time.Now()})
+
+	// Pre-populate N unique upgrade names so each iteration is a guaranteed
+	// cache miss on a fresh keeper (empty cache, value present in KV store).
+	names := make([]string, b.N)
+	for i := range names {
+		names[i] = fmt.Sprintf("test-upgrade-%d", i)
+		app.UpgradeKeeper.SetDone(ctx, names[i])
+	}
+
+	// Single fresh keeper: cache is empty, all names are in the KV store.
+	k := keeper.NewKeeper(
+		make(map[int64]bool),
+		app.GetKey(types.StoreKey),
+		app.AppCodec(),
+		b.TempDir(),
+		app.BaseApp,
+	)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		k.GetDoneHeight(ctx, names[i])
+	}
+}

sei-cosmos/x/upgrade/keeper/keeper_test.go

@@ -312,6 +312,39 @@ func (s *KeeperTestSuite) TestGetClosestUpgrade() {
 	}
 }
 
+// TestGetDoneHeight covers both the cache-hit and cache-miss paths of GetDoneHeight.
+func (s *KeeperTestSuite) TestGetDoneHeight() {
+	// Unknown upgrade returns 0 and must NOT be cached — SetDone may be called later.
+	s.Require().Equal(int64(0), s.app.UpgradeKeeper.GetDoneHeight(s.ctx, "no-such-upgrade"))
+	s.app.UpgradeKeeper.SetUpgradeHandler("no-such-upgrade", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
+	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{Name: "no-such-upgrade", Height: s.ctx.BlockHeight()})
+	s.Require().NotEqual(int64(0), s.app.UpgradeKeeper.GetDoneHeight(s.ctx, "no-such-upgrade"))
+
+	// SetDone writes to KV and populates the cache.
+	s.app.UpgradeKeeper.SetUpgradeHandler("test-upgrade", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
+	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{Name: "test-upgrade", Height: s.ctx.BlockHeight()})
+
+	// Cache hit: same keeper instance — the value was stored by SetDone.
+	s.Require().Equal(s.ctx.BlockHeight(), s.app.UpgradeKeeper.GetDoneHeight(s.ctx, "test-upgrade"))
+
+	// Cache miss: fresh keeper with empty cache but value present in the KV store.
+	freshKeeper := keeper.NewKeeper(
+		make(map[int64]bool),
+		s.app.GetKey(types.StoreKey),
+		s.app.AppCodec(),
+		s.homeDir,
+		s.app.BaseApp,
+	)
+	s.Require().Equal(s.ctx.BlockHeight(), freshKeeper.GetDoneHeight(s.ctx, "test-upgrade"))
+
+	// Second call on fresh keeper should now hit the cache.
+	s.Require().Equal(s.ctx.BlockHeight(), freshKeeper.GetDoneHeight(s.ctx, "test-upgrade"))
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }