refactor(x/protocolpool): Improve claiming funds logic

x/protocolpool/README.md

@@ -125,7 +125,7 @@ https://github.com/cosmos/cosmos-sdk/blob/97724493d792517ba2be8969078b5f92ad04d7
 
 The message will fail under the following conditions:
 
-* The total budget is zero.
+* The budget per tranche is zero.
 * The recipient address is empty or restricted.
 * The start time is less than current block time.
 * The number of tranches is not a positive integer.

x/protocolpool/autocli.go

@@ -42,12 +42,12 @@ func (am AppModule) AutoCLIOptions() *autocliv1.ModuleOptions {
 				},
 				{
 					RpcMethod: "SubmitBudgetProposal",
-					Use:       "submit-budget-proposal [recipient] [total-budget] [start-time] [tranches] [period]",
+					Use:       "submit-budget-proposal [recipient] [budget-per-tranche] [start-time] [tranches] [period]",
 					Short:     "Submit a budget proposal",
 					Example:   fmt.Sprintf(`$ %s tx protocolpool submit-budget-proposal cosmos1... 1000000uatom 2023-10-31T12:34:56.789Z 10 1000 --from mykey`, version.AppName),
 					PositionalArgs: []*autocliv1.PositionalArgDescriptor{
 						{ProtoField: "recipient_address"},
-						{ProtoField: "total_budget"},
+						{ProtoField: "budget_per_tranche"},
 						{ProtoField: "start_time"},
 						{ProtoField: "tranches"},
 						{ProtoField: "period"},

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 || nil
+		if budget.Period == nil || (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
 		}
 
@@ -77,13 +82,11 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error)
 		}
 		budget = append(budget, &types.Budget{
 			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

@@ -50,28 +50,25 @@ func (k Querier) UnclaimedBudget(ctx context.Context, req *types.QueryUnclaimedB
 		}
 		return nil, err
 	}
+
+	totalBudgetAmountLeftToDistribute := budget.BudgetPerTranche.Amount.Mul(math.NewIntFromUint64(budget.TranchesLeft))
+	totalBudgetAmountLeft := sdk.NewCoin(budget.BudgetPerTranche.Denom, totalBudgetAmountLeftToDistribute)
+
 	var unclaimedBudget sdk.Coin
 	if budget.ClaimedAmount == nil {
-		unclaimedBudget = *budget.TotalBudget
-		zeroCoin := sdk.NewCoin(budget.TotalBudget.Denom, math.ZeroInt())
+		unclaimedBudget = totalBudgetAmountLeft
+		zeroCoin := sdk.NewCoin(budget.BudgetPerTranche.Denom, math.ZeroInt())
 		budget.ClaimedAmount = &zeroCoin
 	} else {
-		unclaimedBudget = budget.TotalBudget.Sub(*budget.ClaimedAmount)
+		unclaimedBudget = totalBudgetAmountLeft
 	}
 
-	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 {
@@ -48,10 +49,10 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				// Prepare a valid budget proposal
 				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)
@@ -62,10 +63,9 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 			expErr:         false,
 			unclaimedFunds: &fooCoin,
 			resp: &types.QueryUnclaimedBudgetResponse{
-				TotalBudget:     &fooCoin,
 				ClaimedAmount:   &zeroCoin,
 				UnclaimedAmount: &fooCoin,
-				NextClaimFrom:   &startTime,
+				NextClaimFrom:   &nextClaimFrom,
 				Period:          &period,
 				TranchesLeft:    2,
 			},
@@ -76,10 +76,10 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				// Prepare a valid budget proposal
 				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)
@@ -99,10 +99,9 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 			expErr:         false,
 			unclaimedFunds: &fooCoin2,
 			resp: &types.QueryUnclaimedBudgetResponse{
-				TotalBudget:     &fooCoin,
 				ClaimedAmount:   &fooCoin2,
 				UnclaimedAmount: &fooCoin2,
-				NextClaimFrom:   &nextClaimFrom,
+				NextClaimFrom:   &secondClaimFrom,
 				Period:          &period,
 				TranchesLeft:    1,
 			},

x/protocolpool/keeper/keeper.go

@@ -350,10 +350,14 @@ func (k Keeper) getClaimableFunds(ctx context.Context, recipientAddr string) (am
 		return sdk.Coin{}, err
 	}
 
+	totalBudgetAmountLeftToDistribute := budget.BudgetPerTranche.Amount.Mul(math.NewIntFromUint64(budget.TranchesLeft))
+	totalBudgetAmountLeft := sdk.NewCoin(budget.BudgetPerTranche.Denom, totalBudgetAmountLeftToDistribute)
+	zeroAmount := sdk.NewCoin(totalBudgetAmountLeft.Denom, math.ZeroInt())
+
 	// check if the distribution is completed
 	if budget.TranchesLeft == 0 && budget.ClaimedAmount != nil {
-		// check that claimed amount is equal to total budget
-		if budget.ClaimedAmount.Equal(budget.TotalBudget) {
+		// check that total budget amount left to distribute equals zero
+		if totalBudgetAmountLeft.Equal(zeroAmount) {
 			// remove the entry of budget ended recipient
 			if err := k.BudgetProposal.Remove(ctx, recipient); err != nil {
 				return sdk.Coin{}, err
@@ -364,20 +368,16 @@ 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
-		zeroCoin := sdk.NewCoin(budget.TotalBudget.Denom, math.ZeroInt())
+	if budget.TranchesLeft != 0 && budget.ClaimedAmount == nil {
+		zeroCoin := sdk.NewCoin(budget.BudgetPerTranche.Denom, math.ZeroInt())
 		budget.ClaimedAmount = &zeroCoin
 	}
 
@@ -386,7 +386,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 +396,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)))
-	amount = sdk.NewCoin(budget.TotalBudget.Denom, coinsToDistribute)
+	coinsToDistribute := math.NewInt(periodsPassed).Mul(budget.BudgetPerTranche.Amount)
+	amount = sdk.NewCoin(budget.BudgetPerTranche.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()))
 
@@ -422,11 +430,11 @@ func (k Keeper) calculateClaimableFunds(ctx context.Context, recipient sdk.AccAd
 }
 
 func (k Keeper) validateAndUpdateBudgetProposal(ctx context.Context, bp types.MsgSubmitBudgetProposal) (*types.Budget, error) {
-	if bp.TotalBudget.IsZero() {
-		return nil, fmt.Errorf("invalid budget proposal: total budget cannot be zero")
+	if bp.BudgetPerTranche.IsZero() {
+		return nil, fmt.Errorf("invalid budget proposal: budget per tranche cannot be zero")
 	}
 
-	if err := validateAmount(sdk.NewCoins(*bp.TotalBudget)); err != nil {
+	if err := validateAmount(sdk.NewCoins(*bp.BudgetPerTranche)); err != nil {
 		return nil, fmt.Errorf("invalid budget proposal: %w", err)
 	}
 
@@ -450,9 +458,9 @@ func (k Keeper) validateAndUpdateBudgetProposal(ctx context.Context, bp types.Ms
 	// Create and return an updated budget proposal
 	updatedBudget := types.Budget{
 		RecipientAddress: bp.RecipientAddress,
-		TotalBudget:      bp.TotalBudget,
-		StartTime:        bp.StartTime,
-		Tranches:         bp.Tranches,
+		BudgetPerTranche: bp.BudgetPerTranche,
+		LastClaimedAt:    bp.StartTime,
+		TranchesLeft:     bp.Tranches,
 		Period:           bp.Period,
 	}