From a842136a36fb7e852b7e5dce973881c11e3610ef Mon Sep 17 00:00:00 2001
From: Srinivas Baride <srini@exidio.co>
Date: Wed, 16 Jun 2021 00:20:50 +0530
Subject: [PATCH] Remove active subscriptions from all quotas while cancel

---
 x/subscription/abci.go              | 37 +++++++++++++++--------------
 x/subscription/keeper/msg_server.go | 10 +++++++-
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/x/subscription/abci.go b/x/subscription/abci.go
index af1dfccd..2ec88e4c 100644
--- a/x/subscription/abci.go
+++ b/x/subscription/abci.go
@@ -18,6 +18,25 @@ func EndBlock(ctx sdk.Context, k keeper.Keeper) []abcitypes.ValidatorUpdate {
 	k.IterateInactiveSubscriptions(ctx, ctx.BlockTime(), func(_ int, key []byte, item types.Subscription) bool {
 		log.Info("inactive subscription", "key", key, "value", item)
 
+		if item.Status.Equal(hubtypes.StatusActive) {
+			k.DeleteInactiveSubscriptionAt(ctx, item.Expiry, item.Id)
+			k.IterateQuotas(ctx, item.Id, func(_ int, quota types.Quota) bool {
+				address := quota.GetAddress()
+				k.DeleteActiveSubscriptionForAddress(ctx, address, item.Id)
+				k.SetInactiveSubscriptionForAddress(ctx, address, item.Id)
+
+				return false
+			})
+
+			item.Status = hubtypes.StatusInactivePending
+			item.StatusAt = ctx.BlockTime()
+
+			k.SetSubscription(ctx, item)
+			k.SetInactiveSubscriptionAt(ctx, item.StatusAt.Add(inactiveDuration), item.Id)
+
+			return false
+		}
+
 		if item.Plan == 0 {
 			consumed := sdk.ZeroInt()
 			k.IterateQuotas(ctx, item.Id, func(_ int, quota types.Quota) bool {
@@ -33,27 +52,9 @@ func EndBlock(ctx sdk.Context, k keeper.Keeper) []abcitypes.ValidatorUpdate {
 			if err := k.SubtractDeposit(ctx, itemOwner, amount); err != nil {
 				log.Error("failed to subtract the deposit", "cause", err)
 			}
-		} else {
-			if item.Status.Equal(hubtypes.StatusActive) {
-				item.Status = hubtypes.StatusInactivePending
-				item.StatusAt = item.Expiry
-
-				k.SetSubscription(ctx, item)
-				k.DeleteInactiveSubscriptionAt(ctx, item.Expiry, item.Id)
-				k.SetInactiveSubscriptionAt(ctx, item.Expiry.Add(inactiveDuration), item.Id)
-
-				return false
-			}
 		}
 
 		k.DeleteInactiveSubscriptionAt(ctx, item.StatusAt.Add(inactiveDuration), item.Id)
-		k.IterateQuotas(ctx, item.Id, func(_ int, quota types.Quota) bool {
-			quotaAddress := quota.GetAddress()
-			k.DeleteActiveSubscriptionForAddress(ctx, quotaAddress, item.Id)
-			k.SetInactiveSubscriptionForAddress(ctx, quotaAddress, item.Id)
-
-			return false
-		})
 
 		item.Status = hubtypes.StatusInactive
 		item.StatusAt = ctx.BlockTime()
diff --git a/x/subscription/keeper/msg_server.go b/x/subscription/keeper/msg_server.go
index 275588b2..a7a1766b 100644
--- a/x/subscription/keeper/msg_server.go
+++ b/x/subscription/keeper/msg_server.go
@@ -229,11 +229,19 @@ func (k *msgServer) MsgCancel(c context.Context, msg *types.MsgCancelRequest) (*
 		k.DeleteInactiveSubscriptionAt(ctx, subscription.Expiry, subscription.Id)
 	}
 
+	k.IterateQuotas(ctx, subscription.Id, func(_ int, quota types.Quota) bool {
+		address := quota.GetAddress()
+		k.DeleteActiveSubscriptionForAddress(ctx, address, subscription.Id)
+		k.SetInactiveSubscriptionForAddress(ctx, address, subscription.Id)
+
+		return false
+	})
+
 	subscription.Status = hubtypes.StatusInactivePending
 	subscription.StatusAt = ctx.BlockTime()
-	k.SetInactiveSubscriptionAt(ctx, ctx.BlockTime().Add(inactiveDuration), subscription.Id)
 
 	k.SetSubscription(ctx, subscription)
+	k.SetInactiveSubscriptionAt(ctx, subscription.StatusAt.Add(inactiveDuration), subscription.Id)
 	ctx.EventManager().EmitTypedEvent(
 		&types.EventCancelSubscription{
 			From: sdk.AccAddress(msgFrom.Bytes()).String(),