Switch accumulation store to ResetAllLocks, only remove after unbond complete (#480)

* This commit introduces a ResetAllLocks function

This prevents the code from hitting the accumulation store for every lock,
and to instead only hit it once per duration.

* Switch to version that fixes bug in underlying CacheKVStore under unfortunate access pattern

* Fix go.sum

* Make accumulation store only get decremented on lock finishing unbond

* Cleanup GRPC test

* Update upgrade_test

* Fix comments and function names per code review

* Update changelog

CHANGELOG.md

@@ -38,16 +38,18 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
+## [v4.0.0]
 
 * Significantly speedup epoch times
 * Fix bug in the lockup module code that caused it to take a linear amount of gas.
 * Make unbonding tokens from the lockup module get automatically claimed when unbonding is done.
 * Add events for all tx types in the gamm module.
+* Add events for adding LP rewards.
 * Make queries to bank total chain balance account for developer vesting correctly.
-* Add ability for nodes to query 
+* Add ability for nodes to query the total amount locked for each denomination.
 * Embedded seeds in init.go
 * Added changelog and info about changelog format.
+* Fix accumulation store only counting bonded tokens, not unbonding tokens, that prevented the front-end from using more correct APY estimates. (Previously, the front-end could only underestimate rewards)
 
 ## [v3.2.0](https://github.com/osmosis/osmosis-labs/releases/tag/v2.0.0) - 2021-06-28
 

app/app.go

 package app
 
 import (
-	"fmt"
 	"io"
 	"net/http"
 	"os"
@@ -312,15 +311,7 @@ func NewOsmosisApp(
 			app.LockupKeeper.ClearAllAccumulationStores(ctx)
 
 			// reset all lock and references
-			for i, lock := range locks {
-				if i%10000 == 0 {
-					ctx.Logger().Info(fmt.Sprintf("Reset %d locks", i))
-				}
-				err = app.LockupKeeper.ResetLock(ctx, lock)
-				if err != nil {
-					panic(err)
-				}
-			}
+			app.LockupKeeper.ResetAllLocks(ctx, locks)
 
 			// configure upgrade for gamm module's pool creation fee param add
 			app.GAMMKeeper.SetParams(ctx, gammtypes.NewParams(sdk.Coins{sdk.NewInt64Coin("uosmo", 1)})) // 1 uOSMO

go.mod

@@ -31,6 +31,6 @@ replace github.com/gogo/protobuf => github.com/regen-network/protobuf v1.3.2-alp
 
 replace github.com/tendermint/tendermint => github.com/tendermint/tendermint v0.34.12
 
-replace github.com/cosmos/cosmos-sdk => github.com/osmosis-labs/cosmos-sdk v0.42.10-0.20210911044605-92805585b7ea
+replace github.com/cosmos/cosmos-sdk => github.com/osmosis-labs/cosmos-sdk v0.42.10-0.20210915013958-01114e89a579
 
 replace github.com/tendermint/tm-db => github.com/osmosis-labs/tm-db v0.6.5-0.20210911033928-ba9154613417

go.sum

@@ -400,8 +400,8 @@ github.com/openzipkin-contrib/zipkin-go-opentracing v0.4.5/go.mod h1:/wsWhb9smxS
 github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJc5AZX7/PBEpw=
 github.com/openzipkin/zipkin-go v0.2.1/go.mod h1:NaW6tEwdmWMaCDZzg8sh+IBNOxHMPnhQw8ySjnjRyN4=
 github.com/openzipkin/zipkin-go v0.2.2/go.mod h1:NaW6tEwdmWMaCDZzg8sh+IBNOxHMPnhQw8ySjnjRyN4=
-github.com/osmosis-labs/cosmos-sdk v0.42.10-0.20210911044605-92805585b7ea h1:M9+MPGB1TMTIsPCnqBwDr77vZ31bdWQXV/WdHcmnlNM=
-github.com/osmosis-labs/cosmos-sdk v0.42.10-0.20210911044605-92805585b7ea/go.mod h1:QByRyM6mubdL7o30ZuJTFWT2cwjmT+GT19XSYwyS7pc=
+github.com/osmosis-labs/cosmos-sdk v0.42.10-0.20210915013958-01114e89a579 h1:wF1dMC97AEdUE4kKCEK4ep1mTnpaPEN6G8IUcL2vu2E=
+github.com/osmosis-labs/cosmos-sdk v0.42.10-0.20210915013958-01114e89a579/go.mod h1:QByRyM6mubdL7o30ZuJTFWT2cwjmT+GT19XSYwyS7pc=
 github.com/osmosis-labs/tm-db v0.6.5-0.20210911033928-ba9154613417 h1:otchJDd2SjFWfs7Tse3ULblGcVWqMJ50BE02XCaqXOo=
 github.com/osmosis-labs/tm-db v0.6.5-0.20210911033928-ba9154613417/go.mod h1:dptYhIpJ2M5kUuenLr+Yyf3zQOv1SgBZcl8/BmWlMBw=
 github.com/otiai10/copy v1.6.0 h1:IinKAryFFuPONZ7cm6T6E2QX/vcJwSnlaA5lfoaXIiQ=

x/lockup/genesis.go

 package lockup
 
 import (
-	"fmt"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/osmosis-labs/osmosis/x/lockup/keeper"
 	"github.com/osmosis-labs/osmosis/x/lockup/types"
@@ -12,15 +10,7 @@ import (
 // state.
 func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState) {
 	k.SetLastLockID(ctx, genState.LastLockId)
-	for i, lock := range genState.Locks {
-		if i%10000 == 0 {
-			ctx.Logger().Debug(fmt.Sprintf("lock number %d, entry %d\n", lock.ID, i))
-		}
-		// reset lock's main operation is to store reference queues for iteration
-		if err := k.ResetLock(ctx, lock); err != nil {
-			panic(err)
-		}
-	}
+	k.ResetAllLocks(ctx, genState.Locks)
 }
 
 // ExportGenesis returns the capability module's exported genesis.

x/lockup/keeper/bench_test.go

@@ -72,9 +72,7 @@ func benchmarkResetLogic(numLockups int, b *testing.B) {
 	b.StartTimer()
 	b.ReportAllocs()
 	// distribute coins from gauges to lockup owners
-	for i := 0; i < numLockups; i++ {
-		app.LockupKeeper.ResetLock(ctx, locks[i])
-	}
+	app.LockupKeeper.ResetAllLocks(ctx, locks)
 }
 
 func BenchmarkResetLogicMedium(b *testing.B) {

x/lockup/keeper/grpc_query_test.go

@@ -18,6 +18,10 @@ func (suite *KeeperTestSuite) BeginUnlocking(addr sdk.AccAddress) {
 	suite.Require().NoError(err)
 }
 
+func (suite *KeeperTestSuite) WithdrawAllMaturedLocks() {
+	suite.app.LockupKeeper.WithdrawAllMaturedLocks(suite.ctx)
+}
+
 func (suite *KeeperTestSuite) TestModuleBalance() {
 	suite.SetupTest()
 
@@ -470,29 +474,24 @@ func (suite *KeeperTestSuite) TestLockedDenom() {
 	suite.SetupTest()
 	addr1 := sdk.AccAddress([]byte("addr1---------------"))
 
+	testTotalLockedDuration := func(durationStr string, expectedAmount int64) {
+		duration, _ := time.ParseDuration(durationStr)
+		res, err := suite.app.LockupKeeper.LockedDenom(
+			sdk.WrapSDKContext(suite.ctx),
+			&types.LockedDenomRequest{Denom: "stake", Duration: duration})
+		suite.Require().NoError(err)
+		suite.Require().Equal(res.Amount, sdk.NewInt(expectedAmount))
+	}
+
 	// lock coins
 	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
 	suite.LockTokens(addr1, coins, time.Hour)
 
 	// test with single lockup
-	res, err := suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: 0})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(10))
-
-	duration, _ := time.ParseDuration("30m")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(10))
-
-	duration, _ = time.ParseDuration("1h")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(10))
-
-	duration, _ = time.ParseDuration("2h")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(0))
+	testTotalLockedDuration("0s", 10)
+	testTotalLockedDuration("30m", 10)
+	testTotalLockedDuration("1h", 10)
+	testTotalLockedDuration("2h", 0)
 
 	// adds different account and lockup for testing
 	addr2 := sdk.AccAddress([]byte("addr2---------------"))
@@ -500,25 +499,18 @@ func (suite *KeeperTestSuite) TestLockedDenom() {
 	coins = sdk.Coins{sdk.NewInt64Coin("stake", 20)}
 	suite.LockTokens(addr2, coins, time.Hour*2)
 
-	duration, _ = time.ParseDuration("30m")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(30))
-
-	duration, _ = time.ParseDuration("1h")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(30))
-
-	duration, _ = time.ParseDuration("2h")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(20))
+	testTotalLockedDuration("30m", 30)
+	testTotalLockedDuration("1h", 30)
+	testTotalLockedDuration("2h", 20)
 
 	// test unlocking
 	suite.BeginUnlocking(addr2)
-	duration, _ = time.ParseDuration("2h")
-	res, err = suite.app.LockupKeeper.LockedDenom(sdk.WrapSDKContext(suite.ctx), &types.LockedDenomRequest{Denom: "stake", Duration: duration})
-	suite.Require().NoError(err)
-	suite.Require().Equal(res.Amount, sdk.NewInt(0))
+	testTotalLockedDuration("2h", 20)
+
+	// finish unlocking
+	duration, _ := time.ParseDuration("2h")
+	suite.ctx = suite.ctx.WithBlockTime(suite.ctx.BlockTime().Add(duration))
+	suite.WithdrawAllMaturedLocks()
+	testTotalLockedDuration("2h", 0)
+	testTotalLockedDuration("1h", 10)
 }

x/lockup/keeper/lock.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"fmt"
+	"sort"
 	"time"
 
 	"github.com/cosmos/cosmos-sdk/store/prefix"
@@ -366,8 +367,72 @@ func (k Keeper) ClearAccumulationStores(ctx sdk.Context) {
 	k.clearKeysByPrefix(ctx, types.KeyPrefixLockAccumulation)
 }
 
-// ResetLock reset lock to lock's previous state on InitGenesis
-func (k Keeper) ResetLock(ctx sdk.Context, lock types.PeriodLock) error {
+// ResetAllLocks takes a set of locks, and initializes state to be storing
+// them all correctly. This utilizes batch optimizations to improve efficiency,
+// as this becomes a bottleneck at chain initialization & upgrades.
+func (k Keeper) ResetAllLocks(ctx sdk.Context, locks []types.PeriodLock) error {
+	// index by coin.Denom, them duration -> amt
+	// We accumulate the accumulation store entries separately,
+	// to avoid hitting the myriad of slowdowns in the SDK iterator creation process.
+	// We then save these once to the accumulation store at the end.
+	accumulationStoreEntries := make(map[string]map[time.Duration]sdk.Int)
+	denoms := []string{}
+	for i, lock := range locks {
+		if i%25000 == 0 {
+			msg := fmt.Sprintf("Reset %d lock refs, cur lock ID %d", i, lock.ID)
+			ctx.Logger().Info(msg)
+		}
+		err := k.setLockAndResetLockRefs(ctx, lock)
+		if err != nil {
+			return err
+		}
+
+		// Add to the accumlation store cache
+		for _, coin := range lock.Coins {
+			// update or create the new map from duration -> Int for this denom.
+			var curDurationMap map[time.Duration]sdk.Int
+			if durationMap, ok := accumulationStoreEntries[coin.Denom]; ok {
+				curDurationMap = durationMap
+				// update or create new amount in the duration map
+				newAmt := coin.Amount
+				if curAmt, ok := durationMap[lock.Duration]; ok {
+					newAmt = newAmt.Add(curAmt)
+				}
+				curDurationMap[lock.Duration] = newAmt
+			} else {
+				denoms = append(denoms, coin.Denom)
+				curDurationMap = map[time.Duration]sdk.Int{lock.Duration: coin.Amount}
+			}
+			accumulationStoreEntries[coin.Denom] = curDurationMap
+		}
+	}
+
+	// deterministically iterate over durationMap cache.
+	sort.Strings(denoms)
+	for _, denom := range denoms {
+		curDurationMap := accumulationStoreEntries[denom]
+		durations := make([]time.Duration, 0, len(curDurationMap))
+		for duration, _ := range curDurationMap {
+			durations = append(durations, duration)
+		}
+		sort.Slice(durations, func(i, j int) bool { return durations[i] < durations[j] })
+		// now that we have a sorted list of durations for this denom,
+		// add them all to accumulation store
+		msg := fmt.Sprintf("Setting accumulation entries for locks for %s, there are %d distinct durations",
+			denom, len(durations))
+		ctx.Logger().Info(msg)
+		for _, d := range durations {
+			amt := curDurationMap[d]
+			k.accumulationStore(ctx, denom).Increase(accumulationKey(d), amt)
+		}
+	}
+
+	return nil
+}
+
+// setLockAndResetLockRefs sets the lock, and resets all of its lock references
+// This puts the lock into a 'clean' state, aside from the AccumulationStore.
+func (k Keeper) setLockAndResetLockRefs(ctx sdk.Context, lock types.PeriodLock) error {
 	err := k.setLock(ctx, lock)
 	if err != nil {
 		return err
@@ -378,11 +443,6 @@ func (k Keeper) ResetLock(ctx sdk.Context, lock types.PeriodLock) error {
 		return k.addLockRefs(ctx, types.KeyPrefixUnlocking, lock)
 	}
 
-	// add to accumulation store when unlocking is not started
-	for _, coin := range lock.Coins {
-		k.accumulationStore(ctx, coin.Denom).Increase(accumulationKey(lock.Duration), coin.Amount)
-	}
-
 	return k.addLockRefs(ctx, types.KeyPrefixNotUnlocking, lock)
 }
 
@@ -453,11 +513,6 @@ func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock) error {
 		return err
 	}
 
-	// remove from accumulation store
-	for _, coin := range lock.Coins {
-		k.accumulationStore(ctx, coin.Denom).Decrease(accumulationKey(lock.Duration), coin.Amount)
-	}
-
 	return nil
 }
 
@@ -492,6 +547,11 @@ func (k Keeper) Unlock(ctx sdk.Context, lock types.PeriodLock) error {
 		return err
 	}
 
+	// remove from accumulation store
+	for _, coin := range lock.Coins {
+		k.accumulationStore(ctx, coin.Denom).Decrease(accumulationKey(lock.Duration), coin.Amount)
+	}
+
 	k.hooks.OnTokenUnlocked(ctx, owner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
 	return nil
 }

x/lockup/keeper/store.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"time"
 
-  "github.com/cosmos/cosmos-sdk/store/prefix"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/osmosis-labs/osmosis/x/lockup/types"
 )
@@ -68,7 +68,8 @@ func (k Keeper) deleteLockRefByKey(ctx sdk.Context, key []byte, lockID uint64) e
 }
 
 func accumulationStorePrefix(denom string) (res []byte) {
-	res = make([]byte, len(types.KeyPrefixLockAccumulation))
+	capacity := len(types.KeyPrefixLockAccumulation) + len(denom) + 1
+	res = make([]byte, len(types.KeyPrefixLockAccumulation), capacity)
 	copy(res, types.KeyPrefixLockAccumulation)
 	res = append(res, []byte(denom+"/")...)
 	return

x/lockup/keeper/upgrade_test.go

@@ -133,6 +133,20 @@ func (suite *KeeperTestSuite) TestUpgradeStoreManagement() {
 					Denom:         "stake",
 					Duration:      time.Second,
 				})
+				suite.Require().Equal(accum.String(), "20")
+
+				accum = suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, lockuptypes.QueryCondition{
+					LockQueryType: lockuptypes.ByDuration,
+					Denom:         "stake",
+					Duration:      50 * time.Second,
+				})
+				suite.Require().Equal(accum.String(), "20")
+
+				accum = suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, lockuptypes.QueryCondition{
+					LockQueryType: lockuptypes.ByDuration,
+					Denom:         "stake",
+					Duration:      200 * time.Second,
+				})
 				suite.Require().Equal(accum.String(), "10")
 			},
 			true,