refactor(x/protocolpool): Improve claiming funds logic

x/protocolpool/keeper/genesis.go

@@ -27,11 +27,16 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) error
 	}
 	for _, budget := range data.Budget {
 		// Validate StartTime
-		if budget.StartTime == nil || budget.StartTime.IsZero() {
-			budget.StartTime = &currentTime
+		if budget.LastClaimedAt == nil || budget.LastClaimedAt.IsZero() {
+			budget.LastClaimedAt = &currentTime
 		}
+		// ignore budgets with period <= 0
+		if budget.Period != nil && budget.Period.Seconds() <= 0 {
+			continue
+		}
+
 		// ignore budget with start time < currentTime
-		if budget.StartTime.Before(currentTime) {
+		if budget.LastClaimedAt.Before(currentTime) {
 			continue
 		}
 
@@ -79,11 +84,10 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error)
 			RecipientAddress: recipient,
 			TotalBudget:      value.TotalBudget,
 			ClaimedAmount:    value.ClaimedAmount,
-			StartTime:        value.StartTime,
-			NextClaimFrom:    value.NextClaimFrom,
-			Tranches:         value.Tranches,
+			LastClaimedAt:    value.LastClaimedAt,
 			TranchesLeft:     value.TranchesLeft,
 			Period:           value.Period,
+			BudgetPerTranche: value.BudgetPerTranche,
 		})
 		return false, nil
 	})

x/protocolpool/keeper/grpc_query.go

@@ -59,19 +59,13 @@ func (k Querier) UnclaimedBudget(ctx context.Context, req *types.QueryUnclaimedB
 		unclaimedBudget = budget.TotalBudget.Sub(*budget.ClaimedAmount)
 	}
 
-	if budget.NextClaimFrom == nil {
-		budget.NextClaimFrom = budget.StartTime
-	}
-
-	if budget.TranchesLeft == 0 {
-		budget.TranchesLeft = budget.Tranches
-	}
+	nextClaimFrom := budget.LastClaimedAt.Add(*budget.Period)
 
 	return &types.QueryUnclaimedBudgetResponse{
 		TotalBudget:     budget.TotalBudget,
 		ClaimedAmount:   budget.ClaimedAmount,
 		UnclaimedAmount: &unclaimedBudget,
-		NextClaimFrom:   budget.NextClaimFrom,
+		NextClaimFrom:   &nextClaimFrom,
 		Period:          budget.Period,
 		TranchesLeft:    budget.TranchesLeft,
 	}, nil

x/protocolpool/keeper/grpc_query_test.go

@@ -15,6 +15,7 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 	period := time.Duration(60) * time.Second
 	zeroCoin := sdk.NewCoin("foo", math.ZeroInt())
 	nextClaimFrom := startTime.Add(period)
+	secondClaimFrom := nextClaimFrom.Add(period)
 	recipientStrAddr, err := codectestutil.CodecOptions{}.GetAddressCodec().BytesToString(recipientAddr)
 	suite.Require().NoError(err)
 	testCases := []struct {
@@ -49,9 +50,10 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         2,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     2,
 					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -65,7 +67,7 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				TotalBudget:     &fooCoin,
 				ClaimedAmount:   &zeroCoin,
 				UnclaimedAmount: &fooCoin,
-				NextClaimFrom:   &startTime,
+				NextClaimFrom:   &nextClaimFrom,
 				Period:          &period,
 				TranchesLeft:    2,
 			},
@@ -77,9 +79,10 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         2,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     2,
 					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -102,7 +105,7 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				TotalBudget:     &fooCoin,
 				ClaimedAmount:   &fooCoin2,
 				UnclaimedAmount: &fooCoin2,
-				NextClaimFrom:   &nextClaimFrom,
+				NextClaimFrom:   &secondClaimFrom,
 				Period:          &period,
 				TranchesLeft:    1,
 			},

x/protocolpool/keeper/keeper.go

@@ -364,19 +364,15 @@ func (k Keeper) getClaimableFunds(ctx context.Context, recipientAddr string) (am
 	}
 
 	currentTime := k.HeaderService.HeaderInfo(ctx).Time
-	startTime := budget.StartTime
 
-	// Check if the start time is reached
-	if currentTime.Before(*startTime) {
-		return sdk.Coin{}, fmt.Errorf("distribution has not started yet")
-	}
-
-	if budget.NextClaimFrom == nil || budget.NextClaimFrom.IsZero() {
-		budget.NextClaimFrom = budget.StartTime
+	// Check if the distribution time has not reached
+	if budget.LastClaimedAt != nil {
+		if currentTime.Before(*budget.LastClaimedAt) {
+			return sdk.Coin{}, fmt.Errorf("distribution has not started yet")
+		}
 	}
 
-	if budget.TranchesLeft == 0 && budget.ClaimedAmount == nil {
-		budget.TranchesLeft = budget.Tranches
+	if budget.TranchesLeft != 0 && budget.ClaimedAmount == nil {
 		zeroCoin := sdk.NewCoin(budget.TotalBudget.Denom, math.ZeroInt())
 		budget.ClaimedAmount = &zeroCoin
 	}
@@ -386,7 +382,7 @@ func (k Keeper) getClaimableFunds(ctx context.Context, recipientAddr string) (am
 
 func (k Keeper) calculateClaimableFunds(ctx context.Context, recipient sdk.AccAddress, budget types.Budget, currentTime time.Time) (amount sdk.Coin, err error) {
 	// Calculate the time elapsed since the last claim time
-	timeElapsed := currentTime.Sub(*budget.NextClaimFrom)
+	timeElapsed := currentTime.Sub(*budget.LastClaimedAt)
 
 	// Check the time elapsed has passed period length
 	if timeElapsed < *budget.Period {
@@ -396,20 +392,28 @@ func (k Keeper) calculateClaimableFunds(ctx context.Context, recipient sdk.AccAd
 	// Calculate how many periods have passed
 	periodsPassed := int64(timeElapsed) / int64(*budget.Period)
 
+	if periodsPassed > int64(budget.TranchesLeft) {
+		periodsPassed = int64(budget.TranchesLeft)
+	}
+
 	// Calculate the amount to distribute for all passed periods
-	coinsToDistribute := math.NewInt(periodsPassed).Mul(budget.TotalBudget.Amount.QuoRaw(int64(budget.Tranches)))
+	coinsToDistribute := math.NewInt(periodsPassed).Mul(budget.BudgetPerTranche.Amount)
 	amount = sdk.NewCoin(budget.TotalBudget.Denom, coinsToDistribute)
 
 	// update the budget's remaining tranches
-	budget.TranchesLeft -= uint64(periodsPassed)
+	if budget.TranchesLeft > uint64(periodsPassed) {
+		budget.TranchesLeft -= uint64(periodsPassed)
+	} else {
+		budget.TranchesLeft = 0
+	}
 
 	// update the ClaimedAmount
 	claimedAmount := budget.ClaimedAmount.Add(amount)
 	budget.ClaimedAmount = &claimedAmount
 
 	// Update the last claim time for the budget
-	nextClaimFrom := budget.NextClaimFrom.Add(*budget.Period)
-	budget.NextClaimFrom = &nextClaimFrom
+	nextClaimFrom := budget.LastClaimedAt.Add(*budget.Period * time.Duration(periodsPassed))
+	budget.LastClaimedAt = &nextClaimFrom
 
 	k.Logger.Debug(fmt.Sprintf("Processing budget for recipient: %s. Amount: %s", budget.RecipientAddress, coinsToDistribute.String()))
 
@@ -447,13 +451,17 @@ func (k Keeper) validateAndUpdateBudgetProposal(ctx context.Context, bp types.Ms
 		return nil, fmt.Errorf("invalid budget proposal: period length should be greater than zero")
 	}
 
+	budgetPerTrancheAmount := bp.TotalBudget.Amount.Quo(math.NewIntFromUint64(bp.Tranches))
+	budgetPerTranche := sdk.NewCoin(bp.TotalBudget.Denom, budgetPerTrancheAmount)
+
 	// Create and return an updated budget proposal
 	updatedBudget := types.Budget{
 		RecipientAddress: bp.RecipientAddress,
 		TotalBudget:      bp.TotalBudget,
-		StartTime:        bp.StartTime,
-		Tranches:         bp.Tranches,
+		LastClaimedAt:    bp.StartTime,
+		TranchesLeft:     bp.Tranches,
 		Period:           bp.Period,
+		BudgetPerTranche: &budgetPerTranche,
 	}
 
 	return &updatedBudget, nil

x/protocolpool/keeper/msg_server_test.go

@@ -152,6 +152,7 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 
 	testCases := map[string]struct {
 		preRun           func()
+		postRun          func()
 		recipientAddress sdk.AccAddress
 		expErr           bool
 		expErrMsg        string
@@ -167,16 +168,17 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 			expErr:           true,
 			expErrMsg:        "no budget found for recipient",
 		},
-		"claiming before start time": {
+		"claiming before last claimed at": {
 			preRun: func() {
-				startTime := suite.environment.HeaderService.HeaderInfo(suite.ctx).Time.Add(3600 * time.Second)
-				// Prepare the budget proposal with a future start time
+				startTime := startTime.Add(3600 * time.Second)
+				// Prepare the budget proposal with a future last claimed at time
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         2,
+					TranchesLeft:     2,
 					Period:           &period,
+					LastClaimedAt:    &startTime,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -192,9 +194,10 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         1,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     1,
 					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -209,9 +212,10 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         2,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     2,
 					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -220,15 +224,46 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 			expErr:           false,
 			claimableFunds:   sdk.NewInt64Coin("foo", 50),
 		},
+		"claiming budget after a long time": {
+			preRun: func() {
+				// Prepare the budget proposal with valid start time and period
+				budget := types.Budget{
+					RecipientAddress: recipientStrAddr,
+					TotalBudget:      &fooCoin,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     2,
+					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
+				}
+				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
+				suite.Require().NoError(err)
+
+				// fast forward the block time by 240 hours
+				hinfo := suite.environment.HeaderService.HeaderInfo(suite.ctx)
+				hinfo.Time = hinfo.Time.Add(240 * time.Hour)
+				suite.ctx = suite.ctx.WithHeaderInfo(hinfo)
+
+			},
+			recipientAddress: recipientAddr,
+			claimableFunds:   sdk.NewInt64Coin("foo", 100), // claiming the whole budget, 2 * 50foo = 100foo
+			postRun: func() {
+				prop, err := suite.poolKeeper.BudgetProposal.Get(suite.ctx, recipientAddr)
+				suite.Require().NoError(err)
+				suite.Require().Equal(uint64(0), prop.TranchesLeft)
+				// check if the lastClaimedAt is correct (in this case 2 periods after the start time)
+				suite.Require().Equal(startTime.Add(period*time.Duration(2)), *prop.LastClaimedAt)
+			},
+		},
 		"double claim attempt with budget period not passed": {
 			preRun: func() {
 				// Prepare the budget proposal with valid start time and period
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         2,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     2,
 					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -254,9 +289,10 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTimeBeforeMonth,
-					Tranches:         2,
+					LastClaimedAt:    &startTimeBeforeMonth,
+					TranchesLeft:     2,
 					Period:           &oneMonthPeriod,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -285,9 +321,10 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientStrAddr,
 					TotalBudget:      &fooCoin,
-					StartTime:        &startTime,
-					Tranches:         2,
+					LastClaimedAt:    &startTime,
+					TranchesLeft:     2,
 					Period:           &period,
+					BudgetPerTranche: &fooCoin2,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -340,6 +377,10 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				suite.Require().NoError(err)
 				suite.Require().Equal(tc.claimableFunds, resp.Amount)
 			}
+
+			if tc.postRun != nil {
+				tc.postRun()
+			}
 		})
 	}
 }

x/protocolpool/proto/cosmos/protocolpool/v1/types.proto

@@ -17,19 +17,17 @@ message Budget {
   cosmos.base.v1beta1.Coin total_budget = 2;
   // claimed_amount is the total amount claimed from the total budget amount requested.
   cosmos.base.v1beta1.Coin claimed_amount = 3;
-  // start_time is the time when the budget becomes claimable.
-  google.protobuf.Timestamp start_time = 4 [(gogoproto.stdtime) = true];
-  // next_claim_from is the time when the budget was last successfully claimed or distributed.
-  // It is used to track the next starting claim time for fund distribution. If set, it cannot be less than start_time.
-  google.protobuf.Timestamp next_claim_from = 5 [(gogoproto.stdtime) = true];
-  // tranches is the number of times the total budget amount is to be distributed.
-  uint64 tranches = 6;
+  // last_claimed_at is the time when the budget was last successfully claimed or distributed.
+  // It is used to track the next starting claim time for fund distribution.
+  google.protobuf.Timestamp last_claimed_at = 4 [(gogoproto.stdtime) = true];
   // tranches_left is the number of tranches left for the amount to be distributed.
-  uint64 tranches_left = 7;
+  uint64 tranches_left = 5;
+  // budget_per_tranche is the amount allocated per tranche.
+  cosmos.base.v1beta1.Coin budget_per_tranche = 6;
   // Period is the time interval(number of seconds) at which funds distribution should be performed.
   // For example, if a period is set to 3600, it represents an action that
   // should occur every hour (3600 seconds).
-  google.protobuf.Duration period = 8 [(gogoproto.stdduration) = true];
+  google.protobuf.Duration period = 7 [(gogoproto.stdduration) = true];
 }
 
 // ContinuousFund defines the fields of continuous fund proposal.

x/protocolpool/types/genesis.go

@@ -53,7 +53,7 @@ func validateBudget(bp Budget) error {
 		return errors.Wrap(sdkerrors.ErrInvalidCoins, amount.String())
 	}
 
-	if bp.Tranches == 0 {
+	if bp.TranchesLeft == 0 {
 		return fmt.Errorf("invalid budget proposal: tranches must be greater than zero")
 	}