continue tracking from same accounting height pre upgrade

app/upgrades/v22/upgrades.go

@@ -23,8 +23,23 @@ func CreateUpgradeHandler(
 			return nil, err
 		}
 
-		// For sake of simplicity, we restart the taker fee tracker accounting height.
-		keepers.PoolManagerKeeper.SetTakerFeeTrackerStartHeight(ctx, ctx.BlockHeight())
+		// Migrate legacy taker fee tracker to new taker fee tracker (for performance reasons)
+
+		oldTakerFeeTrackerForStakers := keepers.PoolManagerKeeper.GetLegacyTakerFeeTrackerForStakers(ctx)
+		for _, coin := range oldTakerFeeTrackerForStakers {
+			err := keepers.PoolManagerKeeper.UpdateTakerFeeTrackerForStakersByDenom(ctx, coin.Denom, coin.Amount)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		oldTakerFeeTrackerForCommunityPool := keepers.PoolManagerKeeper.GetLegacyTakerFeeTrackerForCommunityPool(ctx)
+		for _, coin := range oldTakerFeeTrackerForCommunityPool {
+			err := keepers.PoolManagerKeeper.UpdateTakerFeeTrackerForCommunityPoolByDenom(ctx, coin.Denom, coin.Amount)
+			if err != nil {
+				return nil, err
+			}
+		}
 
 		return migrations, nil
 	}

app/upgrades/v22/upgrades_test.go

@@ -0,0 +1,84 @@
+package v22_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+
+	abci "github.com/cometbft/cometbft/abci/types"
+
+	"github.com/osmosis-labs/osmosis/osmomath"
+	"github.com/osmosis-labs/osmosis/osmoutils"
+	"github.com/osmosis-labs/osmosis/v21/app/apptesting"
+	"github.com/osmosis-labs/osmosis/v21/x/protorev/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	poolmanagertypes "github.com/osmosis-labs/osmosis/v21/x/poolmanager/types"
+)
+
+const (
+	v22UpgradeHeight = int64(10)
+)
+
+type UpgradeTestSuite struct {
+	apptesting.KeeperTestHelper
+}
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+func (s *UpgradeTestSuite) TestUpgrade() {
+	s.Setup()
+
+	expectedTakerFeeForStakers := []sdk.Coin{sdk.NewCoin("uakt", osmomath.NewInt(3000)), sdk.NewCoin("uatom", osmomath.NewInt(1000)), sdk.NewCoin("uosmo", osmomath.NewInt(2000))}
+	expectedTakerFeeForCommunityPool := []sdk.Coin{sdk.NewCoin("uakt", osmomath.NewInt(2000)), sdk.NewCoin("uatom", osmomath.NewInt(3000)), sdk.NewCoin("uosmo", osmomath.NewInt(1000))}
+	expectedTrackerStartHeight := int64(3)
+
+	// Set up old protorev tracker prior to upgrade
+	s.App.PoolManagerKeeper.SetTakerFeeTrackerStartHeight(s.Ctx, expectedTrackerStartHeight)
+	newTakerFeeForStakers := poolmanagertypes.TrackedVolume{
+		Amount: expectedTakerFeeForStakers,
+	}
+	osmoutils.MustSet(s.Ctx.KVStore(s.App.GetKey(poolmanagertypes.StoreKey)), poolmanagertypes.KeyTakerFeeStakersProtoRev, &newTakerFeeForStakers)
+
+	newTakerFeeForCommunityPool := poolmanagertypes.TrackedVolume{
+		Amount: expectedTakerFeeForCommunityPool,
+	}
+	osmoutils.MustSet(s.Ctx.KVStore(s.App.GetKey(poolmanagertypes.StoreKey)), poolmanagertypes.KeyTakerFeeCommunityPoolProtoRev, &newTakerFeeForCommunityPool)
+
+	// Set up cyclic arb tracker just to double check that it is not affected by the upgrade
+	s.App.ProtoRevKeeper.SetCyclicArbProfitTrackerStartHeight(s.Ctx, expectedTrackerStartHeight)
+	cyclicArbProfits := sdk.NewCoins(sdk.NewCoin(types.OsmosisDenomination, osmomath.NewInt(9000)), sdk.NewCoin("Atom", osmomath.NewInt(3000)))
+	err := s.App.AppKeepers.ProtoRevKeeper.UpdateStatistics(s.Ctx, poolmanagertypes.SwapAmountInRoutes{}, cyclicArbProfits[0].Denom, cyclicArbProfits[0].Amount)
+	s.Require().NoError(err)
+	err = s.App.AppKeepers.ProtoRevKeeper.UpdateStatistics(s.Ctx, poolmanagertypes.SwapAmountInRoutes{}, cyclicArbProfits[1].Denom, cyclicArbProfits[1].Amount)
+	s.Require().NoError(err)
+
+	dummyUpgrade(s)
+	s.Require().NotPanics(func() {
+		s.App.BeginBlocker(s.Ctx, abci.RequestBeginBlock{})
+	})
+
+	allProtocolRevenue := s.App.ProtoRevKeeper.GetAllProtocolRevenue(s.Ctx)
+
+	// Check that the taker fee tracker for stakers has been migrated correctly
+	s.Require().Equal(expectedTakerFeeForStakers, allProtocolRevenue.TakerFeesTracker.TakerFeesToStakers)
+	s.Require().Equal(expectedTakerFeeForCommunityPool, allProtocolRevenue.TakerFeesTracker.TakerFeesToCommunityPool)
+	s.Require().Equal(expectedTrackerStartHeight, allProtocolRevenue.TakerFeesTracker.HeightAccountingStartsFrom)
+	s.Require().Equal(expectedTrackerStartHeight, allProtocolRevenue.CyclicArbTracker.HeightAccountingStartsFrom)
+}
+
+func dummyUpgrade(s *UpgradeTestSuite) {
+	s.Ctx = s.Ctx.WithBlockHeight(v22UpgradeHeight - 1)
+	plan := upgradetypes.Plan{Name: "v22", Height: v22UpgradeHeight}
+	err := s.App.UpgradeKeeper.ScheduleUpgrade(s.Ctx, plan)
+	s.Require().NoError(err)
+	_, exists := s.App.UpgradeKeeper.GetUpgradePlan(s.Ctx)
+	s.Require().True(exists)
+
+	s.Ctx = s.Ctx.WithBlockHeight(v22UpgradeHeight)
+}

x/poolmanager/protorev.go

@@ -54,3 +54,47 @@ func (k Keeper) GetTakerFeeTrackerStartHeight(ctx sdk.Context) int64 {
 func (k Keeper) SetTakerFeeTrackerStartHeight(ctx sdk.Context, startHeight int64) {
 	osmoutils.MustSet(ctx.KVStore(k.storeKey), types.KeyTakerFeeProtoRevAccountingHeight, &gogotypes.Int64Value{Value: startHeight})
 }
+
+// GetLegacyTakerFeeTrackerForStakers is the legacy getter, to be used in the v22 upgrade handler and removed in v23.
+func (k Keeper) GetLegacyTakerFeeTrackerForStakers(ctx sdk.Context) (currentTakerFeeForStakers sdk.Coins) {
+	var takerFeeForStakers types.TrackedVolume
+	takerFeeFound, err := osmoutils.Get(ctx.KVStore(k.storeKey), types.KeyTakerFeeStakersProtoRev, &takerFeeForStakers)
+	if err != nil {
+		// We can only encounter an error if a database or serialization errors occurs, so we panic here.
+		// Normally this would be handled by `osmoutils.MustGet`, but since we want to specifically use `osmoutils.Get`,
+		// we also have to manually panic here.
+		panic(err)
+	}
+
+	// If no volume was found, we treat the existing volume as 0.
+	// While we can technically require volume to exist, we would need to store empty coins in state for each pool (past and present),
+	// which is a high storage cost to pay for a weak guardrail.
+	currentTakerFeeForStakers = sdk.Coins(nil)
+	if takerFeeFound {
+		currentTakerFeeForStakers = takerFeeForStakers.Amount
+	}
+
+	return currentTakerFeeForStakers
+}
+
+// GetLegacyTakerFeeTrackerForCommunityPool is the legacy getter, to be used in the v22 upgrade handler and removed in v23.
+func (k Keeper) GetLegacyTakerFeeTrackerForCommunityPool(ctx sdk.Context) (currentTakerFeeForCommunityPool sdk.Coins) {
+	var takerFeeForCommunityPool types.TrackedVolume
+	takerFeeFound, err := osmoutils.Get(ctx.KVStore(k.storeKey), types.KeyTakerFeeCommunityPoolProtoRev, &takerFeeForCommunityPool)
+	if err != nil {
+		// We can only encounter an error if a database or serialization errors occurs, so we panic here.
+		// Normally this would be handled by `osmoutils.MustGet`, but since we want to specifically use `osmoutils.Get`,
+		// we also have to manually panic here.
+		panic(err)
+	}
+
+	// If no volume was found, we treat the existing volume as 0.
+	// While we can technically require volume to exist, we would need to store empty coins in state for each pool (past and present),
+	// which is a high storage cost to pay for a weak guardrail.
+	currentTakerFeeForCommunityPool = sdk.Coins(nil)
+	if takerFeeFound {
+		currentTakerFeeForCommunityPool = takerFeeForCommunityPool.Amount
+	}
+
+	return currentTakerFeeForCommunityPool
+}