From b99c73867c89dfe041d71b70186c23d71ae58711 Mon Sep 17 00:00:00 2001
From: Dev Ojha <dojha@berkeley.edu>
Date: Mon, 5 Feb 2024 15:34:08 -0500
Subject: [PATCH] Speedup TWAP pruning logic

Previously we were unmarshalling the value of all records we prune,
and paying lots of time in re-formatting keys.

Instead we can gather all the data we need from the key we iterate over
(and save us expensive unmarshals - .3s to .5s depending on GC).
We also skip one key formatting, since thats the representation we
are already iterating over (.12s to .21s depending on GC).
We skip time formatting in the second key format (30ms AKA .03s)
We could save more in the second key formatting w/ more code complexity.

I'd estimate we're adding .05s of overhead, which we should measure.

The net change of this should be between .42s to .72s off of epoch,
and be state compatible.
---
 x/twap/store.go      | 16 +++++++++++-----
 x/twap/types/keys.go | 26 ++++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/x/twap/store.go b/x/twap/store.go
index d7e28d772a7..52940feeb00 100644
--- a/x/twap/store.go
+++ b/x/twap/store.go
@@ -94,15 +94,16 @@ func (k Keeper) pruneRecordsBeforeTimeButNewest(ctx sdk.Context, lastKeptTime ti
 	seenPoolAssetTriplets := map[uniqueTriplet]struct{}{}
 
 	for ; iter.Valid(); iter.Next() {
-		twapToRemove, err := types.ParseTwapFromBz(iter.Value())
+		timeIndexKey := iter.Key()
+		timeS, poolId, asset0, asset1, err := types.ParseFieldsFromHistoricalTimeKey(timeIndexKey)
 		if err != nil {
 			return err
 		}
 
 		poolKey := uniqueTriplet{
-			poolId: twapToRemove.PoolId,
-			asset0: twapToRemove.Asset0Denom,
-			asset1: twapToRemove.Asset1Denom,
+			poolId,
+			asset0,
+			asset1,
 		}
 		_, hasSeenPoolRecord := seenPoolAssetTriplets[poolKey]
 		if !hasSeenPoolRecord {
@@ -110,7 +111,12 @@ func (k Keeper) pruneRecordsBeforeTimeButNewest(ctx sdk.Context, lastKeptTime ti
 			continue
 		}
 
-		k.DeleteHistoricalRecord(ctx, twapToRemove)
+		// Now we need to delete the historical record, formatted by both historical time and pool index.
+		// We already are iterating over the historical time index, so we delete that key. Then we
+		// reformat the key to delete the historical pool index key.
+		store.Delete(timeIndexKey)
+		poolIndexKey := types.FormatHistoricalPoolIndexTWAPKeyFromStrTime(poolId, asset0, asset1, timeS)
+		store.Delete(poolIndexKey)
 	}
 	return nil
 }
diff --git a/x/twap/types/keys.go b/x/twap/types/keys.go
index cf1c57087bb..fe69bcba2b3 100644
--- a/x/twap/types/keys.go
+++ b/x/twap/types/keys.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"strconv"
 	time "time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -66,12 +67,33 @@ func FormatHistoricalTimeIndexTWAPKey(accumulatorWriteTime time.Time, poolId uin
 }
 
 func FormatHistoricalPoolIndexTWAPKey(poolId uint64, denom1, denom2 string, accumulatorWriteTime time.Time) []byte {
-	var buffer bytes.Buffer
 	timeS := osmoutils.FormatTimeString(accumulatorWriteTime)
-	fmt.Fprintf(&buffer, "%s%d%s%s%s%s%s%s", HistoricalTWAPPoolIndexPrefix, poolId, KeySeparator, denom1, KeySeparator, denom2, KeySeparator, timeS)
+	return FormatHistoricalPoolIndexTWAPKeyFromStrTime(poolId, denom1, denom2, timeS)
+}
+
+func FormatHistoricalPoolIndexTWAPKeyFromStrTime(poolId uint64, denom1, denom2 string, accumulatorWriteTimeString string) []byte {
+	var buffer bytes.Buffer
+	fmt.Fprintf(&buffer, "%s%d%s%s%s%s%s%s", HistoricalTWAPPoolIndexPrefix, poolId, KeySeparator, denom1, KeySeparator, denom2, KeySeparator, accumulatorWriteTimeString)
 	return buffer.Bytes()
 }
 
+// returns timeString, poolIdString, denom1, denom2, error
+// nolint: revive
+func ParseFieldsFromHistoricalTimeKey(bz []byte) (string, uint64, string, string, error) {
+	split := bytes.Split(bz, []byte(KeySeparator))
+	if len(split) != 5 {
+		return "", 0, "", "", errors.New("invalid key")
+	}
+	timeS := string(split[1])
+	poolId, err := strconv.Atoi(string(split[2]))
+	if err != nil {
+		return "", 0, "", "", err
+	}
+	denom1 := string(split[3])
+	denom2 := string(split[4])
+	return timeS, uint64(poolId), denom1, denom2, err
+}
+
 func FormatHistoricalPoolIndexTimePrefix(poolId uint64, denom1, denom2 string) []byte {
 	return []byte(fmt.Sprintf("%s%d%s%s%s%s%s", HistoricalTWAPPoolIndexPrefix, poolId, KeySeparator, denom1, KeySeparator, denom2, KeySeparator))
 }