From 30b4f737242e0a9a86a5a13adfee88111b6df32a Mon Sep 17 00:00:00 2001
From: stana-ethernal <stana.miric@ethernal.tech>
Date: Thu, 13 Jun 2024 16:03:07 +0200
Subject: [PATCH] (Backport)fix(x/authz)!: limit expired authz grant pruning to
 200 per block

---
 CHANGELOG.md                  |  1 +
 x/authz/keeper/keeper.go      | 10 ++++-
 x/authz/keeper/keeper_test.go | 71 ++++++++++++++++++++++++++++++++++-
 x/authz/module/abci.go        |  2 +-
 4 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb90e0006070..387f146b0d5b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -39,6 +39,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ## [Unreleased]
 
 ## [v0.50.7](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.50.7) - 2024-06-04
+* (authz) [#513](https://github.com/osmosis-labs/cosmos-sdk/pull/513) Limit expired authz grant pruning to 200 per block
 
 ### Improvements
 
diff --git a/x/authz/keeper/keeper.go b/x/authz/keeper/keeper.go
index 974c18000e42..aa97f792f654 100644
--- a/x/authz/keeper/keeper.go
+++ b/x/authz/keeper/keeper.go
@@ -413,7 +413,7 @@ func (k Keeper) removeFromGrantQueue(ctx context.Context, grantKey []byte, grant
 }
 
 // DequeueAndDeleteExpiredGrants deletes expired grants from the state and grant queue.
-func (k Keeper) DequeueAndDeleteExpiredGrants(ctx context.Context) error {
+func (k Keeper) DequeueAndDeleteExpiredGrants(ctx context.Context, limit int) error {
 	store := k.storeService.OpenKVStore(ctx)
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
@@ -423,7 +423,13 @@ func (k Keeper) DequeueAndDeleteExpiredGrants(ctx context.Context) error {
 	}
 	defer iterator.Close()
 
+	count := 0
 	for ; iterator.Valid(); iterator.Next() {
+		// limit the amount of iterations to avoid taking too much time
+		if count >= limit {
+			return nil
+		}
+
 		var queueItem authz.GrantQueueItem
 		if err := k.cdc.Unmarshal(iterator.Value(), &queueItem); err != nil {
 			return err
@@ -445,6 +451,8 @@ func (k Keeper) DequeueAndDeleteExpiredGrants(ctx context.Context) error {
 				return err
 			}
 		}
+
+		count++
 	}
 
 	return nil
diff --git a/x/authz/keeper/keeper_test.go b/x/authz/keeper/keeper_test.go
index 94f3f50fa3c0..f28c574bb698 100644
--- a/x/authz/keeper/keeper_test.go
+++ b/x/authz/keeper/keeper_test.go
@@ -378,7 +378,8 @@ func (s *TestSuite) TestDequeueAllGrantsQueue() {
 	require.NoError(err)
 
 	newCtx := s.ctx.WithBlockTime(exp.AddDate(1, 0, 0))
-	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx)
+	// setting a high limit so all grants are dequeued
+	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx, 200)
 	require.NoError(err)
 
 	s.T().Log("verify expired grants are pruned from the state")
@@ -399,6 +400,74 @@ func (s *TestSuite) TestDequeueAllGrantsQueue() {
 	require.Len(authzs, 1)
 }
 
+func (s *TestSuite) TestDequeueGrantsQueueEdgecases() {
+	require := s.Require()
+	addrs := s.addrs
+	granter := addrs[0]
+	grantee := addrs[1]
+	grantee1 := addrs[2]
+	exp := s.ctx.BlockTime().AddDate(0, 0, 1)
+	a := banktypes.SendAuthorization{SpendLimit: coins100}
+
+	// create few authorizations
+	err := s.authzKeeper.SaveGrant(s.ctx, grantee, granter, &a, &exp)
+	require.NoError(err)
+
+	err = s.authzKeeper.SaveGrant(s.ctx, grantee1, granter, &a, &exp)
+	require.NoError(err)
+
+	exp2 := exp.AddDate(0, 1, 0)
+	err = s.authzKeeper.SaveGrant(s.ctx, granter, grantee1, &a, &exp2)
+	require.NoError(err)
+
+	exp2 = exp.AddDate(2, 0, 0)
+	err = s.authzKeeper.SaveGrant(s.ctx, granter, grantee, &a, &exp2)
+	require.NoError(err)
+
+	newCtx := s.ctx.WithBlockTime(exp.AddDate(1, 0, 0))
+	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx, 0)
+	require.NoError(err)
+
+	s.T().Log("verify no pruning happens when limit is 0")
+	authzs, err := s.authzKeeper.GetAuthorizations(newCtx, grantee, granter)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee1)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, grantee1, granter)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	// expecting to prune 1 record when limit is 1
+	newCtx = s.ctx.WithBlockTime(exp.AddDate(1, 0, 0))
+	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx, 1)
+	require.NoError(err)
+
+	s.T().Log("verify 1 record is prunded when limit is 1")
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, grantee, granter)
+	require.NoError(err)
+	require.Len(authzs, 0)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee1)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, grantee1, granter)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee)
+	require.NoError(err)
+	require.Len(authzs, 1)
+}
+
 func (s *TestSuite) TestGetAuthorization() {
 	addr1 := s.addrs[3]
 	addr2 := s.addrs[4]
diff --git a/x/authz/module/abci.go b/x/authz/module/abci.go
index fc405bb8f06f..d5df8ecb20bf 100644
--- a/x/authz/module/abci.go
+++ b/x/authz/module/abci.go
@@ -8,5 +8,5 @@ import (
 // BeginBlocker is called at the beginning of every block
 func BeginBlocker(ctx sdk.Context, keeper keeper.Keeper) error {
 	// delete all the mature grants
-	return keeper.DequeueAndDeleteExpiredGrants(ctx)
+	return keeper.DequeueAndDeleteExpiredGrants(ctx, 200)
 }