Add LockupHooks to the lockup module

x/lockup/keeper/keeper.go

@@ -15,17 +15,23 @@ import (
 type Keeper struct {
 	cdc      codec.Marshaler
 	storeKey sdk.StoreKey
-	ak       authkeeper.AccountKeeper
-	bk       types.BankKeeper
+
+	hooks types.LockupHooks
+
+	ak authkeeper.AccountKeeper
+	bk types.BankKeeper
 }
 
 // NewKeeper returns an instance of Keeper
-func NewKeeper(cdc codec.Marshaler, storeKey sdk.StoreKey, ak authkeeper.AccountKeeper, bk types.BankKeeper) *Keeper {
+func NewKeeper(cdc codec.Marshaler, storeKey sdk.StoreKey, hooks types.LockupHooks, ak authkeeper.AccountKeeper, bk types.BankKeeper) *Keeper {
 	return &Keeper{
 		cdc:      cdc,
 		storeKey: storeKey,
-		ak:       ak,
-		bk:       bk,
+
+		hooks: hooks,
+
+		ak: ak,
+		bk: bk,
 	}
 }
 

x/lockup/keeper/lock.go

@@ -253,6 +253,8 @@ func (k Keeper) Lock(ctx sdk.Context, lock types.PeriodLock) error {
 			return err
 		}
 	}
+
+	k.hooks.OnTokenLocked(ctx, lock.Owner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
 	return nil
 }
 
@@ -306,5 +308,7 @@ func (k Keeper) Unlock(ctx sdk.Context, lock types.PeriodLock) error {
 			return err
 		}
 	}
+
+	k.hooks.OnTokenUnlocked(ctx, lock.Owner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
 	return nil
 }

x/lockup/spec/06_hooks.md

@@ -11,6 +11,6 @@ In this section we describe the "hooks" that `lockup` module provide for other m
 Upon successful coin lock/unlock, other modules might need to do few actions automatically instead of endblocker basis synchronization.
 
 ```go
-  onTokenLocked(address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time)
-  onTokenUnlocked(address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time)
+  OnTokenLocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time)
+  OnTokenUnlocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time)
 ```

x/lockup/types/hooks.go

@@ -0,0 +1,31 @@
+package types
+
+import (
+	"time"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+type LockupHooks interface {
+	OnTokenLocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time)
+	OnTokenUnlocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time)
+}
+
+// combine multiple gamm hooks, all hook functions are run in array sequence
+type MultiLockupHooks []LockupHooks
+
+func NewMultiLockupHooks(hooks ...LockupHooks) MultiLockupHooks {
+	return hooks
+}
+
+func (h MultiLockupHooks) onTokenLocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time) {
+	for i := range h {
+		h[i].OnTokenLocked(ctx, address, lockID, amount, lockDuration, unlockTime)
+	}
+}
+
+func (h MultiLockupHooks) onTokenUnlocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time) {
+	for i := range h {
+		h[i].OnTokenUnlocked(ctx, address, lockID, amount, lockDuration, unlockTime)
+	}
+}

x/pool-yield/keeper/hooks.go

@@ -1,6 +1,10 @@
 package keeper
 
 import (
+	"time"
+
+	lockuptypes "github.com/c-osmosis/osmosis/x/lockup/types"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	gammtypes "github.com/c-osmosis/osmosis/x/gamm/types"
@@ -11,6 +15,7 @@ type Hooks struct {
 }
 
 var _ gammtypes.GammHooks = Hooks{}
+var _ lockuptypes.LockupHooks = Hooks{}
 
 // Create new pool yield hooks
 func (k Keeper) Hooks() Hooks { return Hooks{k} }
@@ -19,3 +24,43 @@ func (k Keeper) Hooks() Hooks { return Hooks{k} }
 func (h Hooks) AfterPoolCreated(ctx sdk.Context, poolId uint64) {
 	h.k.CreatePoolFarms(ctx, poolId)
 }
+
+func (h Hooks) OnTokenLocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time) {
+	// If the locked token in the lockup module is a pool’s share, attempt to add/remove the share to the farm’s pool
+	for _, coin := range amount {
+		poolId, err := gammtypes.GetPoolIdFromShareDenom(coin.Denom)
+		if err == nil {
+			farmId, err := h.k.GetPoolFarmId(ctx, poolId, lockDuration)
+
+			if err == nil {
+				farm, err := h.k.farmKeeper.GetFarm(ctx, farmId)
+
+				if err == nil {
+					// Note that the Farm module doesn’t custody shares within the module, and leaves other modules to manage the balance.
+					// In this case, the shares are not managed in the pool-yield module, but the lockup module.
+					h.k.farmKeeper.DepositShareToFarm(ctx, farm.FarmId, address, coin.Amount)
+				}
+			}
+		}
+	}
+}
+
+func (h Hooks) OnTokenUnlocked(ctx sdk.Context, address sdk.AccAddress, lockID uint64, amount sdk.Coins, lockDuration time.Duration, unlockTime time.Time) {
+	// If the locked token in the lockup module is a pool’s share, attempt to add/remove the share to the farm’s pool
+	for _, coin := range amount {
+		poolId, err := gammtypes.GetPoolIdFromShareDenom(coin.Denom)
+		if err == nil {
+			farmId, err := h.k.GetPoolFarmId(ctx, poolId, lockDuration)
+
+			if err == nil {
+				farm, err := h.k.farmKeeper.GetFarm(ctx, farmId)
+
+				if err == nil {
+					// Note that the Farm module doesn’t custody shares within the module, and leaves other modules to manage the balance.
+					// In this case, the shares are not managed in the pool-yield module, but the lockup module.
+					h.k.farmKeeper.WithdrawShareFromFarm(ctx, farm.FarmId, address, coin.Amount)
+				}
+			}
+		}
+	}
+}