From cecf71a58af4841cf75c9408651ea87d7af03dcb Mon Sep 17 00:00:00 2001
From: Dev Ojha <ValarDragon@users.noreply.github.com>
Date: Fri, 29 Oct 2021 12:51:24 -0500
Subject: [PATCH] Lower epoch I/O time by 2-3x (#561)

---
 x/incentives/keeper/gauge.go    | 33 ++++++++++++++++++++++++++++++---
 x/incentives/keeper/iterator.go | 11 +++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/x/incentives/keeper/gauge.go b/x/incentives/keeper/gauge.go
index 66ecf90bb49..682ebdb2962 100644
--- a/x/incentives/keeper/gauge.go
+++ b/x/incentives/keeper/gauge.go
@@ -221,6 +221,22 @@ func (k Keeper) GetLocksToDistribution(ctx sdk.Context, distrTo lockuptypes.Quer
 	return []lockuptypes.PeriodLock{}
 }
 
+// getLocksToDistributionWithMaxDuration get locks that are associated to a condition
+// and if its by duration, then use the min Duration
+func (k Keeper) getLocksToDistributionWithMaxDuration(ctx sdk.Context, distrTo lockuptypes.QueryCondition, minDuration time.Duration) []lockuptypes.PeriodLock {
+	switch distrTo.LockQueryType {
+	case lockuptypes.ByDuration:
+		if distrTo.Duration > minDuration {
+			return k.lk.GetLocksLongerThanDurationDenom(ctx, distrTo.Denom, minDuration)
+		}
+		return k.lk.GetLocksLongerThanDurationDenom(ctx, distrTo.Denom, distrTo.Duration)
+	case lockuptypes.ByTime:
+		panic("Gauge by time is present!?!? Should have been blocked in ValidateBasic")
+	default:
+	}
+	return []lockuptypes.PeriodLock{}
+}
+
 // FilteredLocksDistributionEst estimate distribution amount coins from gauge for fitting conditions
 // Expectation: gauge is a valid gauge
 // filteredLocks are all locks that are valid for gauge
@@ -344,10 +360,10 @@ func (k Keeper) doDistributionSends(ctx sdk.Context, distrs *distributionInfo) e
 
 // distributeInternal runs the distribution logic for a gauge, and adds the sends to
 // the distrInfo computed. It also updates the gauge for the distribution.
+// locks is expected to be the correct set of lock recipients for this gauge.
 func (k Keeper) distributeInternal(
-	ctx sdk.Context, gauge types.Gauge, distrInfo *distributionInfo) (sdk.Coins, error) {
+	ctx sdk.Context, gauge types.Gauge, locks []lockuptypes.PeriodLock, distrInfo *distributionInfo) (sdk.Coins, error) {
 	totalDistrCoins := sdk.NewCoins()
-	locks := k.GetLocksToDistribution(ctx, gauge.DistributeTo)
 	lockSum := lockuptypes.SumLocksByDenom(locks, gauge.DistributeTo.Denom)
 
 	if lockSum.IsZero() {
@@ -398,9 +414,20 @@ func (k Keeper) distributeInternal(
 func (k Keeper) Distribute(ctx sdk.Context, gauges []types.Gauge) (sdk.Coins, error) {
 	distrInfo := newDistributionInfo()
 
+	locksByDenomCache := make(map[string][]lockuptypes.PeriodLock)
+
 	totalDistributedCoins := sdk.Coins{}
 	for _, gauge := range gauges {
-		gaugeDistributedCoins, err := k.distributeInternal(ctx, gauge, &distrInfo)
+		// All gauges have a precondition of being ByDuration
+		if _, ok := locksByDenomCache[gauge.DistributeTo.Denom]; !ok {
+			locksByDenomCache[gauge.DistributeTo.Denom] = k.getLocksToDistributionWithMaxDuration(
+				ctx, gauge.DistributeTo, time.Millisecond)
+		}
+		// get this from memory instead of hitting iterators / underlying stores.
+		// due to many details of cacheKVStore, iteration will still cause expensive IAVL reads.
+		allLocks := locksByDenomCache[gauge.DistributeTo.Denom]
+		filteredLocks := FilterLocksByMinDuration(allLocks, gauge.DistributeTo.Duration)
+		gaugeDistributedCoins, err := k.distributeInternal(ctx, gauge, filteredLocks, &distrInfo)
 		if err != nil {
 			return nil, err
 		}
diff --git a/x/incentives/keeper/iterator.go b/x/incentives/keeper/iterator.go
index f56a73680e6..c295063cbde 100644
--- a/x/incentives/keeper/iterator.go
+++ b/x/incentives/keeper/iterator.go
@@ -6,6 +6,7 @@ import (
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/osmosis-labs/osmosis/x/incentives/types"
+	lockuptypes "github.com/osmosis-labs/osmosis/x/lockup/types"
 )
 
 // Returns an iterator over all gauges in the {prefix} space of state, that begin distributing rewards after a specific time
@@ -57,3 +58,13 @@ func (k Keeper) ActiveGaugesIterator(ctx sdk.Context) sdk.Iterator {
 func (k Keeper) FinishedGaugesIterator(ctx sdk.Context) sdk.Iterator {
 	return k.iterator(ctx, types.KeyPrefixFinishedGauges)
 }
+
+func FilterLocksByMinDuration(locks []lockuptypes.PeriodLock, minDuration time.Duration) []lockuptypes.PeriodLock {
+	filteredLocks := make([]lockuptypes.PeriodLock, 0, len(locks)/2)
+	for _, lock := range locks {
+		if lock.Duration >= minDuration {
+			filteredLocks = append(filteredLocks, lock)
+		}
+	}
+	return filteredLocks
+}