Add lockup duration edit method

proto/osmosis/lockup/tx.proto

@@ -17,6 +17,8 @@ service Msg {
       returns (MsgBeginUnlockingAllResponse);
   // MsgBeginUnlocking begins unlocking tokens by lock ID
   rpc BeginUnlocking(MsgBeginUnlocking) returns (MsgBeginUnlockingResponse);
+  // MsgEditLockup edits the existing lockups by lock ID
+  rpc ExtendLockup(MsgExtendLockup) returns (MsgExtendLockupResponse);
 }
 
 message MsgLockTokens {
@@ -49,3 +51,20 @@ message MsgBeginUnlocking {
   ];
 }
 message MsgBeginUnlockingResponse { bool success = 1; }
+
+message MsgExtendLockup {
+  string owner = 1 [ (gogoproto.moretags) = "yaml:\"owner\"" ];
+  uint64 ID = 2;
+
+  // duration to be set. fails if lower than the current duration, or is unlocking
+  google.protobuf.Duration duration = 3 [
+    (gogoproto.nullable) = false,
+    (gogoproto.stdduration) = true,
+    (gogoproto.jsontag) = "duration,omitempty",
+    (gogoproto.moretags) = "yaml:\"duration\""
+  ];
+
+  // extend for other edit, e.g. cancel unlocking
+}
+
+message MsgExtendLockupResponse { bool success = 1; }

x/lockup/keeper/lock.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/gogo/protobuf/proto"
+
 	"github.com/osmosis-labs/osmosis/v7/store"
 	"github.com/osmosis-labs/osmosis/v7/x/lockup/types"
 
@@ -536,3 +537,42 @@ func (k Keeper) unlockInternalLogic(ctx sdk.Context, lock types.PeriodLock) erro
 	k.hooks.OnTokenUnlocked(ctx, owner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
 	return nil
 }
+
+func (k Keeper) ExtendLockup(ctx sdk.Context, lock types.PeriodLock, newDuration time.Duration) error {
+	if lock.IsUnlocking() {
+		return fmt.Errorf("cannot edit unlocking lockup")
+	}
+
+	// check synthetic lockup exists
+	if k.HasAnySyntheticLockups(ctx, lock.ID) {
+		return fmt.Errorf("cannot edit lockup with synthetic lock")
+	}
+
+	oldLock := lock
+
+	if newDuration != 0 {
+		fmt.Println(newDuration, lock.Duration)
+		if !(newDuration > lock.Duration) {
+			return fmt.Errorf("new duration should be greater than the original")
+		}
+
+		// update accumulation store
+		for _, coin := range lock.Coins {
+			k.accumulationStore(ctx, coin.Denom).Decrease(accumulationKey(lock.Duration), coin.Amount)
+			k.accumulationStore(ctx, coin.Denom).Increase(accumulationKey(newDuration), coin.Amount)
+		}
+
+		lock.Duration = newDuration
+	}
+
+	// update lockup
+	err := k.deleteLockRefs(ctx, unlockingPrefix(oldLock.IsUnlocking()), oldLock)
+	if err != nil {
+		return err
+	}
+	err = k.addLockRefs(ctx, lock)
+	if err != nil {
+		return err
+	}
+	return k.setLock(ctx, lock)
+}

x/lockup/keeper/lock_test.go

@@ -648,3 +648,63 @@ func (suite *KeeperTestSuite) TestSlashTokensFromLockByID() {
 	_, err = suite.app.LockupKeeper.SlashTokensFromLockByID(suite.ctx, 1, sdk.Coins{sdk.NewInt64Coin("stake1", 1)})
 	suite.Require().Error(err)
 }
+
+func (suite *KeeperTestSuite) TestEditLockup() {
+	suite.SetupTest()
+
+	// initial check
+	locks, err := suite.app.LockupKeeper.GetPeriodLocks(suite.ctx)
+	suite.Require().NoError(err)
+	suite.Require().Len(locks, 0)
+
+	// lock coins
+	addr := sdk.AccAddress([]byte("addr1---------------"))
+
+	// 1 * time.Second: 10
+	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
+	suite.LockTokens(addr, coins, time.Second)
+
+	// check accumulations
+	acc := suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, types.QueryCondition{
+		Denom:    "stake",
+		Duration: time.Second,
+	})
+	suite.Require().Equal(int64(10), acc.Int64())
+
+	lock, _ := suite.app.LockupKeeper.GetLockByID(suite.ctx, 1)
+
+	// duration decrease should fail
+	err = suite.app.LockupKeeper.ExtendLockup(suite.ctx, *lock, time.Second/2)
+	suite.Require().Error(err)
+
+	// duration increase should success
+	err = suite.app.LockupKeeper.ExtendLockup(suite.ctx, *lock, time.Second*2)
+	suite.Require().NoError(err)
+
+	// check queries
+	lock, _ = suite.app.LockupKeeper.GetLockByID(suite.ctx, lock.ID)
+	suite.Require().Equal(lock.Duration, time.Second*2)
+
+	locks = suite.app.LockupKeeper.GetLocksLongerThanDurationDenom(suite.ctx, "stake", time.Second)
+	suite.Require().Equal(len(locks), 1)
+
+	locks = suite.app.LockupKeeper.GetLocksLongerThanDurationDenom(suite.ctx, "stake", time.Second*2)
+	suite.Require().Equal(len(locks), 1)
+
+	// check accumulations
+	acc = suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, types.QueryCondition{
+		Denom:    "stake",
+		Duration: time.Second,
+	})
+	suite.Require().Equal(int64(10), acc.Int64())
+	acc = suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, types.QueryCondition{
+		Denom:    "stake",
+		Duration: time.Second * 2,
+	})
+	suite.Require().Equal(int64(10), acc.Int64())
+	acc = suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, types.QueryCondition{
+		Denom:    "stake",
+		Duration: time.Second * 3,
+	})
+	suite.Require().Equal(int64(0), acc.Int64())
+}

x/lockup/keeper/msg_server.go

@@ -155,3 +155,32 @@ func createBeginUnlockEvent(lock *types.PeriodLock) sdk.Event {
 		sdk.NewAttribute(types.AttributePeriodLockUnlockTime, lock.EndTime.String()),
 	)
 }
+
+func (server msgServer) ExtendLockup(goCtx context.Context, msg *types.MsgExtendLockup) (*types.MsgExtendLockupResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	lock, err := server.keeper.GetLockByID(ctx, msg.ID)
+	if err != nil {
+		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, err.Error())
+	}
+
+	if msg.Owner != lock.Owner {
+		return nil, sdkerrors.Wrapf(types.ErrNotLockOwner, fmt.Sprintf("msg sender (%s) and lock owner (%s) does not match", msg.Owner, lock.Owner))
+	}
+
+	err = server.keeper.ExtendLockup(ctx, *lock, msg.Duration)
+	if err != nil {
+		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidRequest, err.Error())
+	}
+
+	ctx.EventManager().EmitEvents(sdk.Events{
+		sdk.NewEvent(
+			types.TypeEvtLockTokens,
+			sdk.NewAttribute(types.AttributePeriodLockID, utils.Uint64ToString(lock.ID)),
+			sdk.NewAttribute(types.AttributePeriodLockOwner, lock.Owner),
+			sdk.NewAttribute(types.AttributePeriodLockDuration, lock.Duration.String()),
+		),
+	})
+
+	return &types.MsgExtendLockupResponse{}, nil
+}

x/lockup/keeper/msg_server_test.go

@@ -266,3 +266,78 @@ func (suite *KeeperTestSuite) TestMsgBeginUnlockingAll() {
 		}
 	}
 }
+
+func (suite *KeeperTestSuite) TestMsgEditLockup() {
+	type param struct {
+		coinsToLock       sdk.Coins
+		isSyntheticLockup bool
+		lockOwner         sdk.AccAddress
+		duration          time.Duration
+		newDuration       time.Duration
+	}
+
+	tests := []struct {
+		name       string
+		param      param
+		expectPass bool
+	}{
+		{
+			name: "edit lockups by duration",
+			param: param{
+				coinsToLock:       sdk.Coins{sdk.NewInt64Coin("stake", 10)}, // setup wallet
+				isSyntheticLockup: false,
+				lockOwner:         sdk.AccAddress([]byte("addr1---------------")), // setup wallet
+				duration:          time.Second,
+				newDuration:       time.Second * 2,
+			},
+			expectPass: true,
+		},
+		{
+			name: "edit lockups by lesser duration",
+			param: param{
+				coinsToLock:       sdk.Coins{sdk.NewInt64Coin("stake", 10)}, // setup wallet
+				isSyntheticLockup: false,
+				lockOwner:         sdk.AccAddress([]byte("addr1---------------")), // setup wallet
+				duration:          time.Second,
+				newDuration:       time.Second / 2,
+			},
+			expectPass: false,
+		},
+		{
+			name: "disallow edit when synthetic lockup exists",
+			param: param{
+				coinsToLock:       sdk.Coins{sdk.NewInt64Coin("stake", 10)}, // setup wallet
+				isSyntheticLockup: true,
+				lockOwner:         sdk.AccAddress([]byte("addr1---------------")), // setup wallet
+				duration:          time.Second,
+				newDuration:       time.Second * 2,
+			},
+			expectPass: false,
+		},
+	}
+
+	for _, test := range tests {
+		suite.SetupTest()
+
+		err := simapp.FundAccount(suite.app.BankKeeper, suite.ctx, test.param.lockOwner, test.param.coinsToLock)
+		suite.Require().NoError(err)
+
+		msgServer := keeper.NewMsgServerImpl(suite.app.LockupKeeper)
+		c := sdk.WrapSDKContext(suite.ctx)
+		resp, err := msgServer.LockTokens(c, types.NewMsgLockTokens(test.param.lockOwner, test.param.duration, test.param.coinsToLock))
+		suite.Require().NoError(err)
+
+		if test.param.isSyntheticLockup {
+			err = suite.app.LockupKeeper.CreateSyntheticLockup(suite.ctx, resp.ID, "synthetic", time.Second, false)
+			suite.Require().NoError(err)
+		}
+
+		_, err = msgServer.ExtendLockup(c, types.NewMsgExtendLockup(test.param.lockOwner, resp.ID, test.param.newDuration))
+
+		if test.expectPass {
+			suite.Require().NoError(err, test.name)
+		} else {
+			suite.Require().Error(err, test.name)
+		}
+	}
+}

x/lockup/types/msgs.go

@@ -12,6 +12,7 @@ const (
 	TypeMsgLockTokens        = "lock_tokens"
 	TypeMsgBeginUnlockingAll = "begin_unlocking_all"
 	TypeMsgBeginUnlocking    = "begin_unlocking"
+	TypeMsgExtendLockup      = "edit_lockup"
 )
 
 var _ sdk.Msg = &MsgLockTokens{}
@@ -92,3 +93,30 @@ func (m MsgBeginUnlocking) GetSigners() []sdk.AccAddress {
 	owner, _ := sdk.AccAddressFromBech32(m.Owner)
 	return []sdk.AccAddress{owner}
 }
+
+// NewMsgExtendLockup creates a message to edit the properties of existing locks
+func NewMsgExtendLockup(owner sdk.AccAddress, id uint64, duration time.Duration) *MsgExtendLockup {
+	return &MsgExtendLockup{
+		Owner:    owner.String(),
+		ID:       id,
+		Duration: duration,
+	}
+}
+
+func (m MsgExtendLockup) Route() string { return RouterKey }
+func (m MsgExtendLockup) Type() string  { return TypeMsgExtendLockup }
+func (m MsgExtendLockup) ValidateBasic() error {
+	if m.Duration <= 0 {
+		return fmt.Errorf("duration should be positive: %d < 0", m.Duration)
+	}
+	return nil
+}
+
+func (m MsgExtendLockup) GetSignBytes() []byte {
+	return sdk.MustSortJSON(ModuleCdc.MustMarshalJSON((&m)))
+}
+
+func (m MsgExtendLockup) GetSigners() []sdk.AccAddress {
+	owner, _ := sdk.AccAddressFromBech32(m.Owner)
+	return []sdk.AccAddress{owner}
+}