Skip to content

Commit

Permalink
Partial unlocking implementation (#893)
Browse files Browse the repository at this point in the history
* add partial unlock logic

* add partial unlock test

* fix lint

* Apply suggestions from code review

Co-authored-by: Dev Ojha <[email protected]>

* add test for case after unlocking period

Co-authored-by: Dev Ojha <[email protected]>
  • Loading branch information
2 people authored and UnityChaos committed Feb 20, 2022
1 parent f7db3e1 commit a5bee64
Show file tree
Hide file tree
Showing 12 changed files with 288 additions and 53 deletions.
5 changes: 5 additions & 0 deletions proto/osmosis/lockup/tx.proto
Original file line number Diff line number Diff line change
Expand Up @@ -42,5 +42,10 @@ message MsgBeginUnlockingAllResponse { repeated PeriodLock unlocks = 1; }
message MsgBeginUnlocking {
string owner = 1 [ (gogoproto.moretags) = "yaml:\"owner\"" ];
uint64 ID = 2;
// Amount of unlocking coins. Unlock all if not set.
repeated cosmos.base.v1beta1.Coin coins = 3 [
(gogoproto.nullable) = false,
(gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coins"
];
}
message MsgBeginUnlockingResponse { bool success = 1; }
1 change: 1 addition & 0 deletions x/lockup/client/cli/tx.go
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ func NewBeginUnlockByIDCmd() *cobra.Command {
msg := types.NewMsgBeginUnlocking(
clientCtx.GetFromAddress(),
uint64(id),
nil,
)

return tx.GenerateOrBroadcastTxWithFactory(clientCtx, txf, msg)
Expand Down
2 changes: 1 addition & 1 deletion x/lockup/keeper/iterator.go
Original file line number Diff line number Diff line change
Expand Up @@ -201,7 +201,7 @@ func (k Keeper) beginUnlockFromIterator(ctx sdk.Context, iterator db.Iterator) (

locks := k.getLocksFromIterator(ctx, iterator)
for _, lock := range locks {
err := k.BeginUnlock(ctx, lock)
err := k.BeginUnlock(ctx, lock, nil)
if err != nil {
return locks, err
}
Expand Down
56 changes: 45 additions & 11 deletions x/lockup/keeper/lock.go
Original file line number Diff line number Diff line change
Expand Up @@ -429,22 +429,60 @@ func (k Keeper) Lock(ctx sdk.Context, lock types.PeriodLock) error {
return nil
}

// splitLock splits a lock with the given amount, and stores split new lock to the state
func (k Keeper) splitLock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Coins) (types.PeriodLock, error) {
if lock.IsUnlocking() {
return types.PeriodLock{}, fmt.Errorf("cannot split unlocking lock")
}
lock.Coins = lock.Coins.Sub(coins)
err := k.setLock(ctx, lock)
if err != nil {
return types.PeriodLock{}, err
}

splitLockID := k.GetLastLockID(ctx) + 1
k.SetLastLockID(ctx, splitLockID)

splitLock := types.NewPeriodLock(splitLockID, lock.OwnerAddress(), lock.Duration, lock.EndTime, coins)
err = k.setLock(ctx, splitLock)
return splitLock, err
}

// BeginUnlock is a utility to start unlocking coins from NotUnlocking queue
func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock) error {
// remove lock refs from not unlocking queue
func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Coins) error {
// sanity check
if !coins.IsAllLTE(lock.Coins) {
return fmt.Errorf("requested amount to unlock exceedes locked tokens")
}

// If the amount were unlocking is empty, or the entire coins amount, unlock the entire lock.
// Otherwise, split the lock into two locks, and fully unlock the newly created lock.
// (By virtue, the newly created lock we split into should have the unlock amount)
if len(coins) != 0 && !coins.IsEqual(lock.Coins) {
// prohibit partial unlock if other locks are referring
if k.HasAnySyntheticLockups(ctx, lock.ID) {
return fmt.Errorf("cannot partial unlock a lock with synthetic lockup")
}

splitLock, err := k.splitLock(ctx, lock, coins)
if err != nil {
return err
}
lock = splitLock
}

// remove lock refs from not unlocking queue if exists
err := k.deleteLockRefs(ctx, types.KeyPrefixNotUnlocking, lock)
if err != nil {
return err
}

// store lock with end time set
lock.EndTime = ctx.BlockTime().Add(lock.Duration)
store := ctx.KVStore(k.storeKey)
bz, err := proto.Marshal(&lock)
err = k.setLock(ctx, lock)
if err != nil {
return err
}
store.Set(lockStoreKey(lock.ID), bz)

// add lock refs into unlocking queue
err = k.addLockRefs(ctx, types.KeyPrefixUnlocking, lock)
Expand All @@ -456,11 +494,7 @@ func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock) error {
return nil
}

lockOwner, err := sdk.AccAddressFromBech32(lock.Owner)
if err != nil {
panic(err)
}
k.hooks.OnStartUnlock(ctx, lockOwner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
k.hooks.OnStartUnlock(ctx, lock.OwnerAddress(), lock.ID, lock.Coins, lock.Duration, lock.EndTime)

return nil
}
Expand All @@ -484,7 +518,7 @@ func (k Keeper) Unlock(ctx sdk.Context, lock types.PeriodLock) error {
// TODO: Revisit for Superfluid Staking
func (k Keeper) ForceUnlock(ctx sdk.Context, lock types.PeriodLock) error {
if !lock.IsUnlocking() {
err := k.BeginUnlock(ctx, lock)
err := k.BeginUnlock(ctx, lock, nil)
if err != nil {
return err
}
Expand Down
74 changes: 71 additions & 3 deletions x/lockup/keeper/lock_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ func (suite *KeeperTestSuite) TestBeginUnlockPeriodLock() {
suite.Require().Equal(locks[0].IsUnlocking(), false)

// begin unlock
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, locks[0])
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, locks[0], nil)
suite.Require().NoError(err)

// check locks
Expand Down Expand Up @@ -127,7 +127,7 @@ func (suite *KeeperTestSuite) TestUnlockPeriodLockByID() {
// begin unlock
lock, err = lockKeeper.GetLockByID(suite.ctx, 1)
suite.Require().NoError(err)
err = lockKeeper.BeginUnlock(suite.ctx, *lock)
err = lockKeeper.BeginUnlock(suite.ctx, *lock, nil)
suite.Require().NoError(err)

// unlock 1s after begin unlock
Expand Down Expand Up @@ -190,14 +190,82 @@ func (suite *KeeperTestSuite) TestUnlock() {
suite.Require().NoError(err)

// begin unlock with lock object
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock)
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, nil)
suite.Require().NoError(err)

// unlock with lock object
err = suite.app.LockupKeeper.Unlock(suite.ctx.WithBlockTime(now.Add(time.Second)), lock)
suite.Require().NoError(err)
}

func (suite *KeeperTestSuite) TestPartialUnlock() {
suite.SetupTest()
now := suite.ctx.BlockTime()

addr1 := sdk.AccAddress([]byte("addr1---------------"))
coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}

// lock with balance
err := simapp.FundAccount(suite.app.BankKeeper, suite.ctx, addr1, coins)
suite.Require().NoError(err)
lock, err := suite.app.LockupKeeper.LockTokens(suite.ctx, addr1, coins, time.Second)
suite.Require().NoError(err)

// check unlocking coins
unlockings := suite.app.LockupKeeper.GetAccountUnlockingCoins(suite.ctx, addr1)
suite.Require().Equal(len(unlockings), 0)

// check locked coins
locked := suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx, addr1)
suite.Require().Equal(len(locked), 1)
suite.Require().Equal(locked[0].Amount.Int64(), int64(10))

// test exceeding coins
exceedingCoins := sdk.Coins{sdk.NewInt64Coin("stake", 15)}
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, exceedingCoins)
suite.Require().Error(err)

// test invalid coins
invalidCoins := sdk.Coins{sdk.NewInt64Coin("unknown", 1)}
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, invalidCoins)
suite.Require().Error(err)

// begin unlock partial amount
partialCoins := sdk.Coins{sdk.NewInt64Coin("stake", 1)}
err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, partialCoins)
suite.Require().NoError(err)

// check unlocking coins
unlockings = suite.app.LockupKeeper.GetAccountUnlockingCoins(suite.ctx, addr1)
suite.Require().Equal(len(unlockings), 1)
suite.Require().Equal(unlockings[0].Amount.Int64(), int64(1))

// check locked coins
locked = suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx, addr1)
suite.Require().Equal(len(locked), 1)
suite.Require().Equal(locked[0].Amount.Int64(), int64(10))

// check locked coins after the unlocking period
locked = suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx.WithBlockTime(now.Add(time.Second)), addr1)
suite.Require().Equal(len(locked), 1)
suite.Require().Equal(locked[0].Amount.Int64(), int64(9))

// Finish unlocking partial unlock
partialUnlock := suite.app.LockupKeeper.GetAccountPeriodLocks(suite.ctx, addr1)[1]
err = suite.app.LockupKeeper.Unlock(suite.ctx.WithBlockTime(now.Add(time.Second)), partialUnlock)
suite.Require().NoError(err)

// check unlocking coins
unlockings = suite.app.LockupKeeper.GetAccountUnlockingCoins(suite.ctx, addr1)
suite.Require().Equal(len(unlockings), 0)

// check locked coins
locked = suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx, addr1)
suite.Require().Equal(len(locked), 1)
suite.Require().Equal(locked[0].Amount.Int64(), int64(9))

}

func (suite *KeeperTestSuite) TestModuleLockedCoins() {
suite.SetupTest()

Expand Down
3 changes: 2 additions & 1 deletion x/lockup/keeper/msg_server.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,8 @@ func (server msgServer) BeginUnlocking(goCtx context.Context, msg *types.MsgBegi
if err != nil {
return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
}
err = server.keeper.BeginUnlock(ctx, *lock)

err = server.keeper.BeginUnlock(ctx, *lock, msg.Coins)
if err != nil {
return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
}
Expand Down
9 changes: 9 additions & 0 deletions x/lockup/keeper/synthetic_lock.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ func (k Keeper) GetSyntheticLockup(ctx sdk.Context, lockID uint64, suffix string
func (k Keeper) GetAllSyntheticLockupsByLockup(ctx sdk.Context, lockID uint64) []types.SyntheticLock {
store := ctx.KVStore(k.storeKey)
iterator := sdk.KVStorePrefixIterator(store, combineKeys(types.KeyPrefixSyntheticLockup, sdk.Uint64ToBigEndian(lockID)))
defer iterator.Close()

synthLocks := []types.SyntheticLock{}
for ; iterator.Valid(); iterator.Next() {
Expand All @@ -59,9 +60,17 @@ func (k Keeper) GetAllSyntheticLockupsByLockup(ctx sdk.Context, lockID uint64) [
return synthLocks
}

func (k Keeper) HasAnySyntheticLockups(ctx sdk.Context, lockID uint64) bool {
store := ctx.KVStore(k.storeKey)
iterator := sdk.KVStorePrefixIterator(store, combineKeys(types.KeyPrefixSyntheticLockup, sdk.Uint64ToBigEndian(lockID)))
defer iterator.Close()
return iterator.Valid()
}

func (k Keeper) GetAllSyntheticLockups(ctx sdk.Context) []types.SyntheticLock {
store := ctx.KVStore(k.storeKey)
iterator := sdk.KVStorePrefixIterator(store, types.KeyPrefixSyntheticLockup)
defer iterator.Close()

synthLocks := []types.SyntheticLock{}
for ; iterator.Valid(); iterator.Next() {
Expand Down
9 changes: 9 additions & 0 deletions x/lockup/types/lock.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,15 @@ func (p SyntheticLock) IsUnlocking() bool {
return !p.EndTime.Equal(time.Time{})
}

// OwnerAddress returns locks owner address
func (p PeriodLock) OwnerAddress() sdk.AccAddress {
addr, err := sdk.AccAddressFromBech32(p.Owner)
if err != nil {
panic(err)
}
return addr
}

func SumLocksByDenom(locks []PeriodLock, denom string) sdk.Int {
sum := sdk.NewInt(0)
// validate the denom once, so we can avoid the expensive validate check in the hot loop.
Expand Down
3 changes: 2 additions & 1 deletion x/lockup/types/msgs.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,10 +66,11 @@ func (m MsgBeginUnlockingAll) GetSigners() []sdk.AccAddress {
var _ sdk.Msg = &MsgBeginUnlocking{}

// NewMsgBeginUnlocking creates a message to begin unlocking the tokens of a specific lock
func NewMsgBeginUnlocking(owner sdk.AccAddress, id uint64) *MsgBeginUnlocking {
func NewMsgBeginUnlocking(owner sdk.AccAddress, id uint64, coins sdk.Coins) *MsgBeginUnlocking {
return &MsgBeginUnlocking{
Owner: owner.String(),
ID: id,
Coins: coins,
}
}

Expand Down
Loading

0 comments on commit a5bee64

Please sign in to comment.