feat: twap key refactor

CHANGELOG.md

@@ -45,6 +45,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### State Breaking
 
 * [#7250](https://github.com/osmosis-labs/osmosis/pull/7250) Further filter spam gauges from epoch distribution.
+* [#7472](https://github.com/osmosis-labs/osmosis/pull/7472) Refactor TWAP keys to only require a single key format.
 
 ## v23.0.0
 

Makefile

@@ -168,7 +168,6 @@ go-mock-update:
 	mockgen -source=x/poolmanager/types/pool.go -destination=tests/mocks/pool.go -package=mocks
 	mockgen -source=x/gamm/types/pool.go -destination=tests/mocks/cfmm_pool.go -package=mocks
 	mockgen -source=x/concentrated-liquidity/types/cl_pool_extensionI.go -destination=tests/mocks/cl_pool.go -package=mocks
-	mockgen -source=ingest/sqs/domain/pools.go -destination=tests/mocks/sqs_pool.go -package=mocks -mock_names=PoolI=MockSQSPoolI
 
 ###############################################################################
 ###                                Release                                  ###

app/upgrades/v17/upgrades_test.go

@@ -218,9 +218,6 @@ func (s *UpgradeTestSuite) TestUpgrade() {
 				clPool2TwapRecordHistoricalPoolIndexPreUpgrade, err := keepers.TwapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(s.Ctx, lastPoolIdMinusOne)
 				s.Require().NoError(err)
 
-				clPoolsTwapRecordHistoricalTimeIndexPreUpgrade, err := keepers.TwapKeeper.GetAllHistoricalTimeIndexedTWAPs(s.Ctx)
-				s.Require().NoError(err)
-
 				// Run upgrade handler.
 				dummyUpgrade(s)
 				s.Require().NotPanics(func() {
@@ -239,15 +236,11 @@ func (s *UpgradeTestSuite) TestUpgrade() {
 				clPool2TwapRecordHistoricalPoolIndexPostUpgrade, err := keepers.TwapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(s.Ctx, lastPoolIdMinusOne)
 				s.Require().NoError(err)
 
-				clPoolsTwapRecordHistoricalTimeIndexPostUpgrade, err := keepers.TwapKeeper.GetAllHistoricalTimeIndexedTWAPs(s.Ctx)
-				s.Require().NoError(err)
-
 				// check that all TWAP records aren't empty
 				s.Require().NotEmpty(clPool1TwapRecordPostUpgrade)
 				s.Require().NotEmpty(clPool1TwapRecordHistoricalPoolIndexPostUpgrade)
 				s.Require().NotEmpty(clPool2TwapRecordPostUpgrade)
 				s.Require().NotEmpty(clPool2TwapRecordHistoricalPoolIndexPostUpgrade)
-				s.Require().NotEmpty(clPoolsTwapRecordHistoricalTimeIndexPostUpgrade)
 
 				for _, data := range []struct {
 					pre, post []types.TwapRecord
@@ -262,15 +255,6 @@ func (s *UpgradeTestSuite) TestUpgrade() {
 					}
 				}
 
-				for i := range clPoolsTwapRecordHistoricalTimeIndexPostUpgrade {
-					record := clPoolsTwapRecordHistoricalTimeIndexPostUpgrade[i]
-					if record.PoolId == lastPoolIdMinusOne || record.PoolId == lastPoolIdMinusTwo {
-						assertTwapFlipped(s, clPoolsTwapRecordHistoricalTimeIndexPreUpgrade[i], record)
-					} else if record.PoolId == lastPoolID {
-						assertEqual(s, clPoolsTwapRecordHistoricalTimeIndexPreUpgrade[i], record)
-					}
-				}
-
 				// Retrieve the community pool balance (and the feePool balance) after the upgrade
 				communityPoolBalancePost := s.App.BankKeeper.GetAllBalances(s.Ctx, communityPoolAddress)
 				feePoolCommunityPoolPost := s.App.DistrKeeper.GetFeePool(s.Ctx).CommunityPool

app/upgrades/v24/upgrades.go

@@ -23,6 +23,10 @@ func CreateUpgradeHandler(
 			return nil, err
 		}
 
+		// Now that the TWAP keys are refactored, we can delete all time indexed TWAPs
+		// since we only need the pool indexed TWAPs.
+		keepers.TwapKeeper.DeleteAllHistoricalTimeIndexedTWAPs(ctx)
+
 		return migrations, nil
 	}
 }

app/upgrades/v24/upgrades_test.go

@@ -0,0 +1,104 @@
+package v24_test
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+	"time"
+
+	"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/v23/app/apptesting"
+
+	"github.com/osmosis-labs/osmosis/v23/x/twap/types"
+	twaptypes "github.com/osmosis-labs/osmosis/v23/x/twap/types"
+)
+
+const (
+	v24UpgradeHeight              = int64(10)
+	HistoricalTWAPTimeIndexPrefix = "historical_time_index"
+	KeySeparator                  = "|"
+)
+
+type UpgradeTestSuite struct {
+	apptesting.KeeperTestHelper
+}
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+func (s *UpgradeTestSuite) TestUpgrade() {
+	s.Setup()
+
+	// Manually set up TWAP records indexed by both pool ID and time.
+	twapStoreKey := s.App.GetKey(twaptypes.ModuleName)
+	store := s.Ctx.KVStore(twapStoreKey)
+	twap := twaptypes.TwapRecord{
+		PoolId:                      1,
+		Asset0Denom:                 "foo",
+		Asset1Denom:                 "bar",
+		Height:                      1,
+		Time:                        time.Date(2023, 0o2, 1, 0, 0, 0, 0, time.UTC),
+		P0LastSpotPrice:             osmomath.OneDec(),
+		P1LastSpotPrice:             osmomath.OneDec(),
+		P0ArithmeticTwapAccumulator: osmomath.ZeroDec(),
+		P1ArithmeticTwapAccumulator: osmomath.ZeroDec(),
+		GeometricTwapAccumulator:    osmomath.ZeroDec(),
+		LastErrorTime:               time.Time{}, // no previous error
+	}
+	poolIndexKey := types.FormatHistoricalPoolIndexTWAPKey(twap.PoolId, twap.Asset0Denom, twap.Asset1Denom, twap.Time)
+	osmoutils.MustSet(store, poolIndexKey, &twap)
+
+	// The time index key is a bit manual since we removed the old code that did this programmatically.
+	var buffer bytes.Buffer
+	timeS := osmoutils.FormatTimeString(twap.Time)
+	fmt.Fprintf(&buffer, "%s%d%s%s%s%s%s%s", HistoricalTWAPTimeIndexPrefix, twap.PoolId, KeySeparator, twap.Asset0Denom, KeySeparator, twap.Asset1Denom, KeySeparator, timeS)
+	timeIndexKey := buffer.Bytes()
+	osmoutils.MustSet(store, timeIndexKey, &twap)
+
+	// TWAP records indexed by time should exist
+	twapRecords, err := osmoutils.GatherValuesFromStorePrefix(store, []byte(HistoricalTWAPTimeIndexPrefix), types.ParseTwapFromBz)
+	s.Require().NoError(err)
+	s.Require().Len(twapRecords, 1)
+	s.Require().Equal(twap, twapRecords[0])
+
+	// TWAP records indexed by pool ID should exist.
+	twapRecords, err = s.App.TwapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(s.Ctx, twap.PoolId)
+	s.Require().NoError(err)
+	s.Require().Len(twapRecords, 1)
+	s.Require().Equal(twap, twapRecords[0])
+
+	dummyUpgrade(s)
+	s.Require().NotPanics(func() {
+		s.App.BeginBlocker(s.Ctx, abci.RequestBeginBlock{})
+	})
+
+	// TWAP records indexed by time should be completely removed.
+	twapRecords, err = osmoutils.GatherValuesFromStorePrefix(store, []byte(HistoricalTWAPTimeIndexPrefix), types.ParseTwapFromBz)
+	s.Require().NoError(err)
+	s.Require().Len(twapRecords, 0)
+
+	// TWAP records indexed by pool ID should be untouched.
+	twapRecords, err = s.App.TwapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(s.Ctx, twap.PoolId)
+	s.Require().NoError(err)
+	s.Require().Len(twapRecords, 1)
+	s.Require().Equal(twap, twapRecords[0])
+}
+
+func dummyUpgrade(s *UpgradeTestSuite) {
+	s.Ctx = s.Ctx.WithBlockHeight(v24UpgradeHeight - 1)
+	plan := upgradetypes.Plan{Name: "v24", Height: v24UpgradeHeight}
+	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(v24UpgradeHeight)
+}

proto/osmosis/twap/v1beta1/twap_record.proto

@@ -2,9 +2,6 @@ syntax = "proto3";
 package osmosis.twap.v1beta1;
 
 import "gogoproto/gogo.proto";
-import "google/protobuf/any.proto";
-import "cosmos_proto/cosmos.proto";
-import "cosmos/base/v1beta1/coin.proto";
 import "google/protobuf/timestamp.proto";
 
 option go_package = "github.com/osmosis-labs/osmosis/v23/x/twap/types";
@@ -90,7 +87,11 @@ message PruningState {
     (gogoproto.stdtime) = true,
     (gogoproto.moretags) = "yaml:\"last_kept_time\""
   ];
-  // last_key_seen is the last key of the TWAP records that were pruned
-  // before reaching the block's prune limit
-  bytes last_key_seen = 3;
+  // Deprecated: This field is deprecated.
+  bytes last_key_seen = 3 [ deprecated = true ];
+  // last_seen_pool_id is the pool_id that we will begin pruning in the next
+  // block. This value starts at the highest pool_id at time of epoch, and
+  // decreases until it reaches 1. When it reaches 1, the pruning
+  // process is complete.
+  uint64 last_seen_pool_id = 4;
 }

x/twap/keeper.go

@@ -79,9 +79,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 
 // ExportGenesis returns the twap module's exported genesis.
 func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
-	// These are ordered in increasing order, guaranteed by the iterator
-	// that is prefixed by time.
-	twapRecords, err := k.GetAllHistoricalTimeIndexedTWAPs(ctx)
+	twapRecords, err := k.getAllHistoricalPoolIndexedTWAPs(ctx)
 	if err != nil {
 		panic(err)
 	}

x/twap/keeper_test.go

@@ -341,12 +341,6 @@ func (s *TestSuite) getAllHistoricalRecordsForPool(poolId uint64) []types.TwapRe
 func (s *TestSuite) validateExpectedRecords(expectedRecords []types.TwapRecord) {
 	twapKeeper := s.twapkeeper
 
-	// validate that the time indexed TWAPs are cleared.
-	timeIndexedTwaps, err := twapKeeper.GetAllHistoricalTimeIndexedTWAPs(s.Ctx)
-	s.Require().NoError(err)
-	s.Require().Len(timeIndexedTwaps, len(expectedRecords))
-	s.Require().Equal(timeIndexedTwaps, expectedRecords)
-
 	// validate that the pool indexed TWAPs are cleared.
 	poolIndexedTwaps, err := twapKeeper.GetAllHistoricalPoolIndexedTWAPs(s.Ctx)
 	s.Require().NoError(err)

x/twap/listeners.go

@@ -26,11 +26,14 @@ func (k Keeper) EpochHooks() epochtypes.EpochHooks {
 func (hook *epochhook) AfterEpochEnd(ctx sdk.Context, epochIdentifier string, epochNumber int64) error {
 	if epochIdentifier == hook.k.PruneEpochIdentifier(ctx) {
 		lastKeptTime := ctx.BlockTime().Add(-hook.k.RecordHistoryKeepPeriod(ctx))
-		hook.k.SetPruningState(ctx, types.PruningState{
-			IsPruning:    true,
-			LastKeptTime: lastKeptTime,
-			LastKeySeen:  types.FormatHistoricalTimeIndexTWAPKey(lastKeptTime, 0, "", ""),
-		})
+		poolIdToStartFrom := hook.k.poolmanagerKeeper.GetNextPoolId(ctx) - 1
+		if poolIdToStartFrom > 0 {
+			hook.k.SetPruningState(ctx, types.PruningState{
+				IsPruning:      true,
+				LastKeptTime:   lastKeptTime,
+				LastSeenPoolId: poolIdToStartFrom,
+			})
+		}
 	}
 	return nil
 }

x/twap/listeners_test.go

@@ -257,7 +257,7 @@ func (s *TestSuite) TestAfterEpochEnd() {
 
 	s.twapkeeper.StoreNewRecord(s.Ctx, newestRecord)
 
-	twapsBeforeEpoch, err := s.twapkeeper.GetAllHistoricalTimeIndexedTWAPs(s.Ctx)
+	twapsBeforeEpoch, err := s.twapkeeper.GetAllHistoricalPoolIndexedTWAPs(s.Ctx)
 	s.Require().NoError(err)
 	s.Require().Equal(2, len(twapsBeforeEpoch))