feat: add inactive session handler for node and subscription modules

sentinel-official · Nov 20, 2024 · b41842a · b41842a
1 parent f6042b2
commit b41842a
Show file tree

Hide file tree

Showing 11 changed files with 122 additions and 31 deletions.
diff --git a/x/node/expected/keeper.go b/x/node/expected/keeper.go
@@ -34,6 +34,9 @@ type LeaseKeeper interface {
 }
 
 type SessionKeeper interface {
+	DeleteSession(ctx sdk.Context, id uint64)
+	DeleteSessionForAccount(ctx sdk.Context, addr sdk.AccAddress, id uint64)
+	DeleteSessionForNode(ctx sdk.Context, addr base.NodeAddress, id uint64)
 	GetCount(ctx sdk.Context) uint64
 	GetSession(ctx sdk.Context, id uint64) (sessiontypes.Session, bool)
 	NodeInactivePreHook(ctx sdk.Context, addr base.NodeAddress) error

diff --git a/x/node/keeper/hooks.go b/x/node/keeper/hooks.go
@@ -11,7 +11,7 @@ import (
 	"github.com/sentinel-official/hub/v12/x/node/types/v3"
 )
 
-func (k *Keeper) SessionInactivePreHook(ctx sdk.Context, id uint64) error {
+func (k *Keeper) HandleInactiveSession(ctx sdk.Context, id uint64) error {
 	item, found := k.session.GetSession(ctx, id)
 	if !found {
 		return fmt.Errorf("session %d does not exist", id)
@@ -71,5 +71,9 @@ func (k *Keeper) SessionInactivePreHook(ctx sdk.Context, id uint64) error {
 		},
 	)
 
+	k.session.DeleteSession(ctx, item.GetID())
+	k.session.DeleteSessionForAccount(ctx, accAddr, item.GetID())
+	k.session.DeleteSessionForNode(ctx, nodeAddr, item.GetID())
+
 	return nil
 }
diff --git a/x/session/expected/keeper.go b/x/session/expected/keeper.go
@@ -19,9 +19,9 @@ type DepositKeeper interface {
 }
 
 type NodeKeeper interface {
-	SessionInactivePreHook(ctx sdk.Context, id uint64) error
+	HandleInactiveSession(ctx sdk.Context, id uint64) error
 }
 
 type SubscriptionKeeper interface {
-	SessionInactivePreHook(ctx sdk.Context, id uint64) error
+	HandleInactiveSession(ctx sdk.Context, id uint64) error
 }
diff --git a/x/session/keeper/abci.go b/x/session/keeper/abci.go
@@ -3,7 +3,6 @@ package keeper
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
-	base "github.com/sentinel-official/hub/v12/types"
 	v1base "github.com/sentinel-official/hub/v12/types/v1"
 	"github.com/sentinel-official/hub/v12/x/session/types/v3"
 )
@@ -36,28 +35,12 @@ func (k *Keeper) handleInactiveSessions(ctx sdk.Context) {
 			return false
 		}
 
-		accAddr, err := sdk.AccAddressFromBech32(item.GetAccAddress())
-		if err != nil {
-			panic(err)
-		}
-
-		nodeAddr, err := base.NodeAddressFromBech32(item.GetNodeAddress())
-		if err != nil {
-			panic(err)
-		}
-
 		k.DeleteSessionForInactiveAt(ctx, item.GetInactiveAt(), item.GetID())
 
-		if err := k.SessionInactivePreHook(ctx, item.GetID()); err != nil {
+		if err := k.HandleInactiveSession(ctx, item.GetID()); err != nil {
 			panic(err)
 		}
 
-		k.DeleteSession(ctx, item.GetID())
-		k.DeleteSessionForAccount(ctx, accAddr, item.GetID())
-		k.DeleteSessionForAllocation(ctx, 0, accAddr, item.GetID())
-		k.DeleteSessionForNode(ctx, nodeAddr, item.GetID())
-		k.DeleteSessionForSubscription(ctx, 0, item.GetID())
-
 		ctx.EventManager().EmitTypedEvent(
 			&v3.EventUpdateStatus{
 				ID:          item.GetID(),

diff --git a/x/session/keeper/alias.go b/x/session/keeper/alias.go
@@ -20,11 +20,11 @@ func (k *Keeper) SendCoinFromDepositToModule(ctx sdk.Context, from sdk.AccAddres
 	return k.deposit.SendCoinsFromDepositToModule(ctx, from, to, sdk.NewCoins(coin))
 }
 
-func (k *Keeper) SessionInactivePreHook(ctx sdk.Context, id uint64) error {
-	if err := k.node.SessionInactivePreHook(ctx, id); err != nil {
+func (k *Keeper) HandleInactiveSession(ctx sdk.Context, id uint64) error {
+	if err := k.node.HandleInactiveSession(ctx, id); err != nil {
 		return err
 	}
-	if err := k.subscription.SessionInactivePreHook(ctx, id); err != nil {
+	if err := k.subscription.HandleInactiveSession(ctx, id); err != nil {
 		return err
 	}
 

diff --git a/x/session/keeper/hooks.go b/x/session/keeper/hooks.go
@@ -8,10 +8,6 @@ import (
 	"github.com/sentinel-official/hub/v12/x/session/types/v3"
 )
 
-func (k *Keeper) LeaseInactivePreHook(_ sdk.Context, _ uint64) error {
-	return nil
-}
-
 func (k *Keeper) NodeInactivePreHook(ctx sdk.Context, addr base.NodeAddress) error {
 	k.IterateSessionsForNode(ctx, addr, func(_ int, item v3.Session) bool {
 		if !item.GetStatus().Equal(v1base.StatusActive) {
@@ -55,3 +51,25 @@ func (k *Keeper) SubscriptionInactivePendingPreHook(ctx sdk.Context, id uint64)
 
 	return nil
 }
+
+func (k *Keeper) PlanUnlinkNodePreHook(ctx sdk.Context, id uint64, addr base.NodeAddress) error {
+	k.IterateSessionsForPlanByNode(ctx, id, addr, func(_ int, item v3.Session) bool {
+		if !item.GetStatus().Equal(v1base.StatusActive) {
+			return false
+		}
+
+		msg := &v3.MsgCancelSessionRequest{
+			From: item.GetAccAddress(),
+			ID:   item.GetID(),
+		}
+
+		handler := k.router.Handler(msg)
+		if _, err := handler(ctx, msg); err != nil {
+			panic(err)
+		}
+
+		return false
+	})
+
+	return nil
+}
diff --git a/x/session/keeper/session.go b/x/session/keeper/session.go
@@ -183,6 +183,44 @@ func (k *Keeper) IterateSessionsForNode(ctx sdk.Context, addr base.NodeAddress,
 	}
 }
 
+// SetSessionForPlanByNode links a session ID to a plan ID and node address in the module's KVStore.
+func (k *Keeper) SetSessionForPlanByNode(ctx sdk.Context, planID uint64, addr base.NodeAddress, sessionID uint64) {
+	store := k.Store(ctx)
+	key := types.SessionForPlanByNodeKey(planID, addr, sessionID)
+	value := k.cdc.MustMarshal(&protobuf.BoolValue{Value: true})
+
+	store.Set(key, value)
+}
+
+// DeleteSessionForPlanByNode removes the association between a session ID, a plan ID, and a node address from the module's KVStore.
+func (k *Keeper) DeleteSessionForPlanByNode(ctx sdk.Context, planID uint64, addr base.NodeAddress, sessionID uint64) {
+	store := k.Store(ctx)
+	key := types.SessionForPlanByNodeKey(planID, addr, sessionID)
+
+	store.Delete(key)
+}
+
+// IterateSessionsForPlanByNode iterates over all sessions for a specific plan ID and node address, calling the provided function for each session.
+// The iteration stops when the provided function returns 'true'.
+func (k *Keeper) IterateSessionsForPlanByNode(ctx sdk.Context, id uint64, addr base.NodeAddress, fn func(index int, item v3.Session) (stop bool)) {
+	store := k.Store(ctx)
+	iterator := sdk.KVStorePrefixIterator(store, types.GetSessionForPlanByNodeKeyPrefix(id, addr))
+
+	defer iterator.Close()
+
+	for i := 0; iterator.Valid(); iterator.Next() {
+		item, found := k.GetSession(ctx, types.IDFromSessionForPlanByNodeKey(iterator.Key()))
+		if !found {
+			panic(fmt.Errorf("session for plan by node key %X does not exist", iterator.Key()))
+		}
+
+		if stop := fn(i, item); stop {
+			break
+		}
+		i++
+	}
+}
+
 // SetSessionForSubscription links a session ID to a subscription ID in the module's KVStore.
 func (k *Keeper) SetSessionForSubscription(ctx sdk.Context, subscriptionID, sessionID uint64) {
 	store := k.Store(ctx)

diff --git a/x/session/types/keys.go b/x/session/types/keys.go
@@ -22,8 +22,9 @@ var (
 	SessionForAccountKeyPrefix      = []byte{0x11}
 	SessionForAllocationKeyPrefix   = []byte{0x12}
 	SessionForNodeKeyPrefix         = []byte{0x13}
-	SessionForSubscriptionKeyPrefix = []byte{0x14}
-	SessionForInactiveAtKeyPrefix   = []byte{0x15}
+	SessionForPlanKeyPrefix         = []byte{0x14}
+	SessionForSubscriptionKeyPrefix = []byte{0x15}
+	SessionForInactiveAtKeyPrefix   = []byte{0x16}
 )
 
 func SessionKey(id uint64) []byte {
@@ -46,6 +47,18 @@ func SessionForNodeKey(addr base.NodeAddress, id uint64) []byte {
 	return append(GetSessionForNodeKeyPrefix(addr), sdk.Uint64ToBigEndian(id)...)
 }
 
+func GetSessionForPlanKeyPrefix(id uint64) []byte {
+	return append(SessionForPlanKeyPrefix, sdk.Uint64ToBigEndian(id)...)
+}
+
+func GetSessionForPlanByNodeKeyPrefix(id uint64, addr base.NodeAddress) []byte {
+	return append(GetSessionForPlanKeyPrefix(id), sdkaddress.MustLengthPrefix(addr.Bytes())...)
+}
+
+func SessionForPlanByNodeKey(planID uint64, addr base.NodeAddress, sessionID uint64) []byte {
+	return append(GetSessionForPlanByNodeKeyPrefix(planID, addr), sdk.Uint64ToBigEndian(sessionID)...)
+}
+
 func GetSessionForSubscriptionKeyPrefix(id uint64) []byte {
 	return append(SessionForSubscriptionKeyPrefix, sdk.Uint64ToBigEndian(id)...)
 }
@@ -92,6 +105,17 @@ func IDFromSessionForNodeKey(key []byte) uint64 {
 	return sdk.BigEndianToUint64(key[2+addrLen:])
 }
 
+func IDFromSessionForPlanByNodeKey(key []byte) uint64 {
+	// prefix (1 byte) | planID (8 bytes) | addrLen (1 byte) | addr (addrLen bytes) | sessionID (8 bytes)
+
+	addrLen := int(key[9])
+	if len(key) != 18+addrLen {
+		panic(fmt.Errorf("invalid key length %d; expected %d", len(key), 18+addrLen))
+	}
+
+	return sdk.BigEndianToUint64(key[10+addrLen:])
+}
+
 func IDFromSessionForSubscriptionKey(key []byte) uint64 {
 	// prefix (1 byte) | subscriptionID (8 bytes) | sessionID (8 bytes)
 

diff --git a/x/subscription/expected/keeper.go b/x/subscription/expected/keeper.go
@@ -44,6 +44,12 @@ type PlanKeeper interface {
 }
 
 type SessionKeeper interface {
+	DeleteSession(ctx sdk.Context, id uint64)
+	DeleteSessionForAccount(ctx sdk.Context, addr sdk.AccAddress, id uint64)
+	DeleteSessionForAllocation(ctx sdk.Context, subscriptionID uint64, addr sdk.AccAddress, sessionID uint64)
+	DeleteSessionForNode(ctx sdk.Context, addr base.NodeAddress, id uint64)
+	DeleteSessionForPlanByNode(ctx sdk.Context, planID uint64, addr base.NodeAddress, sessionID uint64)
+	DeleteSessionForSubscription(ctx sdk.Context, subscriptionID, sessionID uint64)
 	GetCount(ctx sdk.Context) uint64
 	GetSession(ctx sdk.Context, id uint64) (sessiontypes.Session, bool)
 	SetCount(ctx sdk.Context, count uint64)
@@ -52,6 +58,7 @@ type SessionKeeper interface {
 	SetSessionForAllocation(ctx sdk.Context, subscriptionID uint64, addr sdk.AccAddress, sessionID uint64)
 	SetSessionForInactiveAt(ctx sdk.Context, at time.Time, id uint64)
 	SetSessionForNode(ctx sdk.Context, addr base.NodeAddress, id uint64)
+	SetSessionForPlanByNode(ctx sdk.Context, planID uint64, addr base.NodeAddress, sessionID uint64)
 	SetSessionForSubscription(ctx sdk.Context, subscriptionID, sessionID uint64)
 	StatusChangeDelay(ctx sdk.Context) time.Duration
 	SubscriptionInactivePendingPreHook(ctx sdk.Context, id uint64) error

diff --git a/x/subscription/keeper/hooks.go b/x/subscription/keeper/hooks.go
@@ -5,11 +5,12 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
+	base "github.com/sentinel-official/hub/v12/types"
 	v1base "github.com/sentinel-official/hub/v12/types/v1"
 	"github.com/sentinel-official/hub/v12/x/subscription/types/v3"
 )
 
-func (k *Keeper) SessionInactivePreHook(ctx sdk.Context, id uint64) error {
+func (k *Keeper) HandleInactiveSession(ctx sdk.Context, id uint64) error {
 	item, found := k.session.GetSession(ctx, id)
 	if !found {
 		return fmt.Errorf("session %d does not exist", id)
@@ -33,6 +34,11 @@ func (k *Keeper) SessionInactivePreHook(ctx sdk.Context, id uint64) error {
 		panic(err)
 	}
 
+	nodeAddr, err := base.NodeAddressFromBech32(item.GetNodeAddress())
+	if err != nil {
+		return err
+	}
+
 	alloc, found := k.GetAllocation(ctx, subscription.ID, accAddr)
 	if !found {
 		return fmt.Errorf("subscription allocation %d/%s does not exist", subscription.ID, accAddr)
@@ -55,5 +61,12 @@ func (k *Keeper) SessionInactivePreHook(ctx sdk.Context, id uint64) error {
 		},
 	)
 
+	k.session.DeleteSession(ctx, item.GetID())
+	k.session.DeleteSessionForAccount(ctx, accAddr, item.GetID())
+	k.session.DeleteSessionForAllocation(ctx, subscription.ID, accAddr, item.GetID())
+	k.session.DeleteSessionForNode(ctx, nodeAddr, item.GetID())
+	k.session.DeleteSessionForPlanByNode(ctx, subscription.PlanID, nodeAddr, item.GetID())
+	k.session.DeleteSessionForSubscription(ctx, subscription.ID, item.GetID())
+
 	return nil
 }
diff --git a/x/subscription/keeper/msg_handler.go b/x/subscription/keeper/msg_handler.go
@@ -416,6 +416,7 @@ func (k *Keeper) HandleMsgStartSession(ctx sdk.Context, msg *v3.MsgStartSessionR
 	k.session.SetSession(ctx, session)
 	k.session.SetSessionForAccount(ctx, accAddr, session.ID)
 	k.session.SetSessionForNode(ctx, nodeAddr, session.ID)
+	k.session.SetSessionForPlanByNode(ctx, subscription.PlanID, nodeAddr, session.ID)
 	k.session.SetSessionForSubscription(ctx, subscription.ID, session.ID)
 	k.session.SetSessionForAllocation(ctx, subscription.ID, accAddr, session.ID)
 	k.session.SetSessionForInactiveAt(ctx, session.InactiveAt, session.ID)