feat(x/protocolpool,x/distribution)!: remove dependency + fix continuous fund bug (#20790)

x/protocolpool/keeper/genesis.go

@@ -3,7 +3,9 @@ package keeper
 import (
 	"context"
 	"fmt"
+	"time"
 
+	"cosmossdk.io/math"
 	"cosmossdk.io/x/protocolpool/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -49,8 +51,21 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) error
 		}
 	}
 
-	if err := k.ToDistribute.Set(ctx, data.ToDistribute); err != nil {
-		return fmt.Errorf("failed to set to distribute: %w", err)
+	if err := k.LastBalance.Set(ctx, data.LastBalance); err != nil {
+		return fmt.Errorf("failed to set last balance: %w", err)
+	}
+
+	totalToBeDistributed := math.ZeroInt()
+	for _, distribution := range data.Distributions {
+		totalToBeDistributed = totalToBeDistributed.Add(distribution.Amount)
+		if err := k.Distributions.Set(ctx, *distribution.Time, distribution.Amount); err != nil {
+			return fmt.Errorf("failed to set distribution: %w", err)
+		}
+	}
+
+	// sanity check to avoid trying to distribute more than what is available
+	if data.LastBalance.LT(totalToBeDistributed) {
+		return fmt.Errorf("total to be distributed is greater than the last balance")
 	}
 
 	return nil
@@ -96,7 +111,19 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error)
 
 	genState := types.NewGenesisState(cf, budget)
 
-	genState.ToDistribute, err = k.ToDistribute.Get(ctx)
+	genState.LastBalance, err = k.LastBalance.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	err = k.Distributions.Walk(ctx, nil, func(key time.Time, value math.Int) (stop bool, err error) {
+		genState.Distributions = append(genState.Distributions, &types.Distribution{
+			Time:   &key,
+			Amount: value,
+		})
+
+		return false, nil
+	})
 	if err != nil {
 		return nil, err
 	}

x/protocolpool/keeper/genesis_test.go

@@ -31,7 +31,17 @@ func (suite *KeeperTestSuite) TestInitGenesis() {
 		},
 	)
 
+	gs.Distributions = append(gs.Distributions, &types.Distribution{
+		Amount: math.OneInt(),
+		Time:   &time.Time{},
+	})
+
 	err := suite.poolKeeper.InitGenesis(suite.ctx, gs)
+	suite.Require().ErrorContains(err, "total to be distributed is greater than the last balance")
+
+	// Set last balance
+	gs.LastBalance = math.NewInt(1)
+	err = suite.poolKeeper.InitGenesis(suite.ctx, gs)
 	suite.Require().NoError(err)
 
 	// Export
@@ -39,4 +49,5 @@ func (suite *KeeperTestSuite) TestInitGenesis() {
 	suite.Require().NoError(err)
 	suite.Require().Equal(gs.ContinuousFund, exportedGenState.ContinuousFund)
 	suite.Require().Equal(gs.Budget, exportedGenState.Budget)
+	suite.Require().Equal(math.OneInt(), exportedGenState.LastBalance)
 }

x/protocolpool/keeper/keeper.go

@@ -1,10 +1,10 @@
 package keeper
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
+	"sort"
 	"time"
 
 	"cosmossdk.io/collections"
@@ -35,8 +35,8 @@ type Keeper struct {
 	ContinuousFund collections.Map[sdk.AccAddress, types.ContinuousFund]
 	// RecipientFundDistribution key: RecipientAddr | value: Claimable amount
 	RecipientFundDistribution collections.Map[sdk.AccAddress, math.Int]
-	// ToDistribute is to keep track of funds to be distributed. It gets zeroed out in iterateAndUpdateFundsDistribution.
-	ToDistribute collections.Item[math.Int]
+	Distributions             collections.Map[time.Time, math.Int] // key: time.Time | value: amount
+	LastBalance               collections.Item[math.Int]
 }
 
 func NewKeeper(cdc codec.BinaryCodec, env appmodule.Environment, ak types.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, authority string,
@@ -49,6 +49,10 @@ func NewKeeper(cdc codec.BinaryCodec, env appmodule.Environment, ak types.Accoun
 	if addr := ak.GetModuleAddress(types.StreamAccount); addr == nil {
 		panic(fmt.Sprintf("%s module account has not been set", types.StreamAccount))
 	}
+	// ensure protocol pool distribution account is set
+	if addr := ak.GetModuleAddress(types.ProtocolPoolDistrAccount); addr == nil {
+		panic(fmt.Sprintf("%s module account has not been set", types.ProtocolPoolDistrAccount))
+	}
 
 	sb := collections.NewSchemaBuilder(env.KVStoreService)
 
@@ -62,7 +66,8 @@ func NewKeeper(cdc codec.BinaryCodec, env appmodule.Environment, ak types.Accoun
 		BudgetProposal:            collections.NewMap(sb, types.BudgetKey, "budget", sdk.AccAddressKey, codec.CollValue[types.Budget](cdc)),
 		ContinuousFund:            collections.NewMap(sb, types.ContinuousFundKey, "continuous_fund", sdk.AccAddressKey, codec.CollValue[types.ContinuousFund](cdc)),
 		RecipientFundDistribution: collections.NewMap(sb, types.RecipientFundDistributionKey, "recipient_fund_distribution", sdk.AccAddressKey, sdk.IntValue),
-		ToDistribute:              collections.NewItem(sb, types.ToDistributeKey, "to_distribute", sdk.IntValue),
+		Distributions:             collections.NewMap(sb, types.DistributionsKey, "distributions", sdk.TimeKey, sdk.IntValue),
+		LastBalance:               collections.NewItem(sb, types.LastBalanceKey, "last_balance", sdk.IntValue),
 	}
 
 	schema, err := sb.Build()
@@ -80,19 +85,19 @@ func (k Keeper) GetAuthority() string {
 }
 
 // FundCommunityPool allows an account to directly fund the community fund pool.
-func (k Keeper) FundCommunityPool(ctx context.Context, amount sdk.Coins, sender sdk.AccAddress) error {
+func (k Keeper) FundCommunityPool(ctx context.Context, amount sdk.Coins, sender []byte) error {
 	return k.bankKeeper.SendCoinsFromAccountToModule(ctx, sender, types.ModuleName, amount)
 }
 
 // DistributeFromCommunityPool distributes funds from the protocolpool module account to
 // a receiver address.
-func (k Keeper) DistributeFromCommunityPool(ctx context.Context, amount sdk.Coins, receiveAddr sdk.AccAddress) error {
+func (k Keeper) DistributeFromCommunityPool(ctx context.Context, amount sdk.Coins, receiveAddr []byte) error {
 	return k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, receiveAddr, amount)
 }
 
 // DistributeFromStreamFunds distributes funds from the protocolpool's stream module account to
 // a receiver address.
-func (k Keeper) DistributeFromStreamFunds(ctx context.Context, amount sdk.Coins, receiveAddr sdk.AccAddress) error {
+func (k Keeper) DistributeFromStreamFunds(ctx context.Context, amount sdk.Coins, receiveAddr []byte) error {
 	return k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.StreamAccount, receiveAddr, amount)
 }
 
@@ -100,48 +105,17 @@ func (k Keeper) DistributeFromStreamFunds(ctx context.Context, amount sdk.Coins,
 func (k Keeper) GetCommunityPool(ctx context.Context) (sdk.Coins, error) {
 	moduleAccount := k.authKeeper.GetModuleAccount(ctx, types.ModuleName)
 	if moduleAccount == nil {
-		return nil, errorsmod.Wrapf(sdkerrors.ErrUnknownAddress, "module account %s does not exist", moduleAccount)
+		return nil, errorsmod.Wrapf(sdkerrors.ErrUnknownAddress, "module account %s does not exist", types.ModuleName)
 	}
 	return k.bankKeeper.GetAllBalances(ctx, moduleAccount.GetAddress()), nil
 }
 
-func (k Keeper) withdrawContinuousFund(ctx context.Context, recipientAddr string) (sdk.Coin, error) {
-	recipient, err := k.authKeeper.AddressCodec().StringToBytes(recipientAddr)
-	if err != nil {
-		return sdk.Coin{}, sdkerrors.ErrInvalidAddress.Wrapf("invalid recipient address: %s", err)
-	}
-
-	cf, err := k.ContinuousFund.Get(ctx, recipient)
-	if err != nil {
-		if errors.Is(err, collections.ErrNotFound) {
-			return sdk.Coin{}, fmt.Errorf("no continuous fund found for recipient: %s", recipientAddr)
-		}
-		return sdk.Coin{}, fmt.Errorf("get continuous fund failed for recipient: %s", recipientAddr)
-	}
-	if cf.Expiry != nil && cf.Expiry.Before(k.HeaderService.HeaderInfo(ctx).Time) {
-		return sdk.Coin{}, fmt.Errorf("cannot withdraw continuous funds: continuous fund expired for recipient: %s", recipientAddr)
-	}
-
-	err = k.IterateAndUpdateFundsDistribution(ctx)
-	if err != nil {
-		return sdk.Coin{}, fmt.Errorf("error while iterating all the continuous funds: %w", err)
-	}
-
-	// withdraw continuous fund
-	withdrawnAmount, err := k.withdrawRecipientFunds(ctx, recipient)
-	if err != nil {
-		return sdk.Coin{}, fmt.Errorf("error while withdrawing recipient funds for recipient: %s", recipientAddr)
-	}
-
-	return withdrawnAmount, nil
-}
-
 func (k Keeper) withdrawRecipientFunds(ctx context.Context, recipient []byte) (sdk.Coin, error) {
 	// get allocated continuous fund
 	fundsAllocated, err := k.RecipientFundDistribution.Get(ctx, recipient)
 	if err != nil {
 		if errors.Is(err, collections.ErrNotFound) {
-			return sdk.Coin{}, types.ErrNoRecipientFund
+			return sdk.Coin{}, types.ErrNoRecipientFound
 		}
 		return sdk.Coin{}, err
 	}
@@ -166,20 +140,12 @@ func (k Keeper) withdrawRecipientFunds(ctx context.Context, recipient []byte) (s
 	return withdrawnAmount, nil
 }
 
-// SetToDistribute sets the amount to be distributed among recipients, usually called by x/distribution while allocating
-// reward and fee distribution.
-// This could be only set by the authority address.
-func (k Keeper) SetToDistribute(ctx context.Context, amount sdk.Coins, addr string) error {
-	authAddr, err := k.authKeeper.AddressCodec().StringToBytes(addr)
-	if err != nil {
-		return err
-	}
-	hasPermission, err := k.hasPermission(authAddr)
-	if err != nil {
-		return err
-	}
-	if !hasPermission {
-		return sdkerrors.ErrUnauthorized
+// SetToDistribute sets the amount to be distributed among recipients.
+func (k Keeper) SetToDistribute(ctx context.Context) error {
+	// Get current balance of the intermediary module account
+	moduleAccount := k.authKeeper.GetModuleAccount(ctx, types.ProtocolPoolDistrAccount)
+	if moduleAccount == nil {
+		return errorsmod.Wrapf(sdkerrors.ErrUnknownAddress, "module account %s does not exist", types.ProtocolPoolDistrAccount)
 	}
 
 	denom, err := k.stakingKeeper.BondDenom(ctx)
@@ -187,106 +153,44 @@ func (k Keeper) SetToDistribute(ctx context.Context, amount sdk.Coins, addr stri
 		return err
 	}
 
-	totalStreamFundsPercentage := math.LegacyZeroDec()
-	err = k.ContinuousFund.Walk(ctx, nil, func(key sdk.AccAddress, cf types.ContinuousFund) (stop bool, err error) {
-		// Check if the continuous fund has expired
-		if cf.Expiry != nil && cf.Expiry.Before(k.HeaderService.HeaderInfo(ctx).Time) {
-			return false, nil
-		}
+	currentBalance := k.bankKeeper.GetAllBalances(ctx, moduleAccount.GetAddress())
+	distributionBalance := currentBalance.AmountOf(denom)
 
-		totalStreamFundsPercentage = totalStreamFundsPercentage.Add(cf.Percentage)
-		if totalStreamFundsPercentage.GT(math.LegacyOneDec()) {
-			return true, errors.New("total funds percentage cannot exceed 100")
-		}
-
-		return false, nil
-	})
-	if err != nil {
-		return err
-	}
-
-	// if percentage is 0 then return early
-	if totalStreamFundsPercentage.IsZero() {
+	// if the balance is zero, return early
+	if distributionBalance.IsZero() {
 		return nil
 	}
 
-	// send streaming funds to the stream module account
-	toDistributeAmt := math.LegacyNewDecFromInt(amount.AmountOf(denom)).Mul(totalStreamFundsPercentage).TruncateInt()
-	streamAmt := sdk.NewCoins(sdk.NewCoin(denom, toDistributeAmt))
-	if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ModuleName, types.StreamAccount, streamAmt); err != nil {
-		return err
-	}
-
-	amountToDistribute, err := k.ToDistribute.Get(ctx)
+	lastBalance, err := k.LastBalance.Get(ctx)
 	if err != nil {
 		if errors.Is(err, collections.ErrNotFound) {
-			amountToDistribute = math.ZeroInt()
+			lastBalance = math.ZeroInt()
 		} else {
 			return err
 		}
 	}
 
-	err = k.ToDistribute.Set(ctx, amountToDistribute.Add(amount.AmountOf(denom)))
-	if err != nil {
-		return fmt.Errorf("error while setting ToDistribute: %w", err)
-	}
-	return nil
-}
+	// Calculate the amount to be distributed
+	amountToDistribute := distributionBalance.Sub(lastBalance)
 
-func (k Keeper) hasPermission(addr []byte) (bool, error) {
-	authority := k.GetAuthority()
-	authAcc, err := k.authKeeper.AddressCodec().StringToBytes(authority)
-	if err != nil {
-		return false, err
+	if err = k.Distributions.Set(ctx, k.HeaderService.HeaderInfo(ctx).Time, amountToDistribute); err != nil {
+		return fmt.Errorf("error while setting Distributions: %w", err)
 	}
 
-	return bytes.Equal(authAcc, addr), nil
+	// Update the last balance
+	return k.LastBalance.Set(ctx, distributionBalance)
 }
 
 func (k Keeper) IterateAndUpdateFundsDistribution(ctx context.Context) error {
-	toDistributeAmount, err := k.ToDistribute.Get(ctx)
-	if err != nil {
-		return err
-	}
+	// first we get all the continuous funds, and keep a list of the ones that expired so we can delete later
+	funds := []types.ContinuousFund{}
+	toDelete := [][]byte{}
+	err := k.ContinuousFund.Walk(ctx, nil, func(key sdk.AccAddress, cf types.ContinuousFund) (stop bool, err error) {
+		funds = append(funds, cf)
 
-	// if there are no funds to distribute, return
-	if toDistributeAmount.IsZero() {
-		return nil
-	}
-
-	totalPercentageToBeDistributed := math.LegacyZeroDec()
-
-	denom, err := k.stakingKeeper.BondDenom(ctx)
-	if err != nil {
-		return err
-	}
-	toDistributeDec := sdk.NewDecCoin(denom, toDistributeAmount)
-
-	// Calculate totalPercentageToBeDistributed and store values
-	err = k.ContinuousFund.Walk(ctx, nil, func(key sdk.AccAddress, cf types.ContinuousFund) (stop bool, err error) {
-		// Check if the continuous fund has expired
+		// check if the continuous fund has expired, and add it to the list of funds to delete
 		if cf.Expiry != nil && cf.Expiry.Before(k.HeaderService.HeaderInfo(ctx).Time) {
-			return false, nil
-		}
-
-		// sanity check for max percentage
-		totalPercentageToBeDistributed = totalPercentageToBeDistributed.Add(cf.Percentage)
-		if totalPercentageToBeDistributed.GT(math.LegacyOneDec()) {
-			return true, errors.New("total funds percentage cannot exceed 100")
-		}
-
-		// Calculate the funds to be distributed based on the percentage
-		recipientAmount := toDistributeDec.Amount.Mul(cf.Percentage).TruncateInt()
-
-		// Set funds to be claimed
-		toClaim, err := k.RecipientFundDistribution.Get(ctx, key)
-		if err != nil {
-			return true, err
-		}
-		amount := toClaim.Add(recipientAmount)
-		err = k.RecipientFundDistribution.Set(ctx, key, amount)
-		if err != nil {
-			return true, err
+			toDelete = append(toDelete, key)
 		}
 
 		return false, nil
@@ -295,8 +199,108 @@ func (k Keeper) IterateAndUpdateFundsDistribution(ctx context.Context) error {
 		return err
 	}
 
-	// Set the coins to be distributed from toDistribute to 0
-	return k.ToDistribute.Set(ctx, math.ZeroInt())
+	// next we iterate over the distributions, calculate each recipient's share and the remaining pool funds
+	toDistribute := map[string]math.Int{}
+	poolFunds := math.ZeroInt()
+	fullAmountToDistribute := math.ZeroInt()
+
+	if err = k.Distributions.Walk(ctx, nil, func(key time.Time, amount math.Int) (stop bool, err error) {
+		percentageToDistribute := math.LegacyZeroDec()
+		for _, f := range funds {
+			if f.Expiry != nil && f.Expiry.Before(key) {
+				continue
+			}
+
+			percentageToDistribute = percentageToDistribute.Add(f.Percentage)
+
+			_, ok := toDistribute[f.Recipient]
+			if !ok {
+				toDistribute[f.Recipient] = math.ZeroInt()
+			}
+			amountToDistribute := f.Percentage.MulInt(amount).TruncateInt()
+			toDistribute[f.Recipient] = toDistribute[f.Recipient].Add(amountToDistribute)
+			fullAmountToDistribute = fullAmountToDistribute.Add(amountToDistribute)
+		}
+
+		// sanity check for max percentage
+		if percentageToDistribute.GT(math.LegacyOneDec()) {
+			return true, errors.New("total funds percentage cannot exceed 100")
+		}
+
+		remaining := math.LegacyOneDec().Sub(percentageToDistribute).MulInt(amount).RoundInt()
+		poolFunds = poolFunds.Add(remaining)
+
+		return false, nil
+	}); err != nil {
+		return err
+	}
+
+	// clear the distributions and reset the last balance
+	if err = k.Distributions.Clear(ctx, nil); err != nil {
+		return err
+	}
+
+	if err = k.LastBalance.Set(ctx, math.ZeroInt()); err != nil {
+		return err
+	}
+
+	// send the funds to the stream account to be distributed later, and the remaining to the community pool
+	bondDenom, err := k.stakingKeeper.BondDenom(ctx)
+	if err != nil {
+		return err
+	}
+
+	streamAmt := sdk.NewCoins(sdk.NewCoin(bondDenom, fullAmountToDistribute))
+	if !streamAmt.IsZero() {
+		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.StreamAccount, streamAmt); err != nil {
+			return err
+		}
+	}
+
+	if !poolFunds.IsZero() {
+		poolCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, poolFunds))
+		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.ModuleName, poolCoins); err != nil {
+			return err
+		}
+	}
+
+	// update the recipient fund distribution, first get the keys and sort them
+	recipients := make([]string, 0, len(toDistribute))
+	for k2 := range toDistribute {
+		recipients = append(recipients, k2)
+	}
+	sort.Strings(recipients)
+
+	for _, recipient := range recipients {
+		// Set funds to be claimed
+		bzAddr, err := k.authKeeper.AddressCodec().StringToBytes(recipient)
+		if err != nil {
+			return err
+		}
+
+		toClaim, err := k.RecipientFundDistribution.Get(ctx, bzAddr)
+		if err != nil {
+			if errors.Is(err, collections.ErrNotFound) {
+				toClaim = math.ZeroInt()
+			} else {
+				return err
+			}
+		}
+
+		amount := toClaim.Add(toDistribute[recipient])
+		if err = k.RecipientFundDistribution.Set(ctx, bzAddr, amount); err != nil {
+			return err
+		}
+	}
+
+	// delete expired continuous funds
+	for _, recipient := range toDelete {
+		if err = k.ContinuousFund.Remove(ctx, recipient); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 func (k Keeper) claimFunds(ctx context.Context, recipientAddr string) (amount sdk.Coin, err error) {
@@ -472,3 +476,7 @@ func (k Keeper) validateContinuousFund(ctx context.Context, msg types.MsgCreateC
 
 	return nil
 }
+
+func (k Keeper) BeginBlocker(ctx context.Context) error {
+	return k.SetToDistribute(ctx)
+}

x/protocolpool/keeper/keeper_test.go

@@ -27,8 +27,9 @@ import (
 )
 
 var (
-	poolAcc   = authtypes.NewEmptyModuleAccount(types.ModuleName)
-	streamAcc = authtypes.NewEmptyModuleAccount(types.StreamAccount)
+	poolAcc      = authtypes.NewEmptyModuleAccount(types.ModuleName)
+	streamAcc    = authtypes.NewEmptyModuleAccount(types.StreamAccount)
+	poolDistrAcc = authtypes.NewEmptyModuleAccount(types.ProtocolPoolDistrAccount)
 )
 
 type KeeperTestSuite struct {
@@ -57,6 +58,7 @@ func (s *KeeperTestSuite) SetupTest() {
 	ctrl := gomock.NewController(s.T())
 	accountKeeper := pooltestutil.NewMockAccountKeeper(ctrl)
 	accountKeeper.EXPECT().GetModuleAddress(types.ModuleName).Return(poolAcc.GetAddress())
+	accountKeeper.EXPECT().GetModuleAddress(types.ProtocolPoolDistrAccount).Return(poolDistrAcc.GetAddress())
 	accountKeeper.EXPECT().AddressCodec().Return(address.NewBech32Codec("cosmos")).AnyTimes()
 	accountKeeper.EXPECT().GetModuleAddress(types.StreamAccount).Return(streamAcc.GetAddress())
 	s.authKeeper = accountKeeper
@@ -102,12 +104,14 @@ func (s *KeeperTestSuite) mockWithdrawContinuousFund() {
 	s.stakingKeeper.EXPECT().BondDenom(gomock.Any()).Return("stake", nil).AnyTimes()
 }
 
-func (s *KeeperTestSuite) mockStreamFunds() {
+func (s *KeeperTestSuite) mockStreamFunds(distributed math.Int) {
 	s.authKeeper.EXPECT().GetModuleAccount(s.ctx, types.ModuleName).Return(poolAcc).AnyTimes()
+	s.authKeeper.EXPECT().GetModuleAccount(s.ctx, types.ProtocolPoolDistrAccount).Return(poolDistrAcc).AnyTimes()
 	s.authKeeper.EXPECT().GetModuleAddress(types.StreamAccount).Return(streamAcc.GetAddress()).AnyTimes()
-	distrBal := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(100000)))
-	s.bankKeeper.EXPECT().GetAllBalances(s.ctx, poolAcc.GetAddress()).Return(distrBal).AnyTimes()
-	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(s.ctx, poolAcc.GetName(), streamAcc.GetName(), gomock.Any()).AnyTimes()
+	distrBal := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, distributed))
+	s.bankKeeper.EXPECT().GetAllBalances(s.ctx, poolDistrAcc.GetAddress()).Return(distrBal).AnyTimes()
+	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(s.ctx, poolDistrAcc.GetName(), streamAcc.GetName(), gomock.Any()).AnyTimes()
+	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(s.ctx, poolDistrAcc.GetName(), poolAcc.GetName(), gomock.Any()).AnyTimes()
 }
 
 func TestKeeperTestSuite(t *testing.T) {
@@ -118,10 +122,11 @@ func (s *KeeperTestSuite) TestIterateAndUpdateFundsDistribution() {
 	// We'll create 2 continuous funds of 30% each, and the total pool is 1000000, meaning each fund should get 300000
 
 	s.SetupTest()
-	s.authKeeper.EXPECT().GetModuleAccount(s.ctx, types.ModuleName).Return(poolAcc).AnyTimes()
+	s.authKeeper.EXPECT().GetModuleAccount(s.ctx, types.ProtocolPoolDistrAccount).Return(poolAcc).AnyTimes()
 	distrBal := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(1000000)))
 	s.bankKeeper.EXPECT().GetAllBalances(s.ctx, poolAcc.GetAddress()).Return(distrBal).AnyTimes()
-	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(s.ctx, poolAcc.GetName(), streamAcc.GetName(), sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(600000)))).AnyTimes()
+	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(s.ctx, poolDistrAcc.GetName(), streamAcc.GetName(), sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(600000))))
+	s.bankKeeper.EXPECT().SendCoinsFromModuleToModule(s.ctx, poolDistrAcc.GetName(), poolAcc.GetName(), sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(400000))))
 
 	_, err := s.msgServer.CreateContinuousFund(s.ctx, &types.MsgCreateContinuousFund{
 		Authority:  s.poolKeeper.GetAuthority(),
@@ -137,7 +142,7 @@ func (s *KeeperTestSuite) TestIterateAndUpdateFundsDistribution() {
 	})
 	s.Require().NoError(err)
 
-	_ = s.poolKeeper.SetToDistribute(s.ctx, sdk.NewCoins(sdk.NewCoin("stake", math.NewInt(1000000))), s.poolKeeper.GetAuthority())
+	_ = s.poolKeeper.SetToDistribute(s.ctx)
 
 	err = s.poolKeeper.IterateAndUpdateFundsDistribution(s.ctx)
 	s.Require().NoError(err)

x/protocolpool/keeper/msg_server.go

@@ -139,6 +139,11 @@ func (k MsgServer) CreateContinuousFund(ctx context.Context, msg *types.MsgCreat
 		return nil, fmt.Errorf("cannot set continuous fund proposal\ntotal funds percentage exceeds 100\ncurrent total percentage: %s", totalStreamFundsPercentage.Sub(msg.Percentage).MulInt64(100).TruncateInt().String())
 	}
 
+	// Distribute funds to avoid giving this new fund more than it should get
+	if err := k.IterateAndUpdateFundsDistribution(ctx); err != nil {
+		return nil, err
+	}
+
 	// Create continuous fund proposal
 	cf := types.ContinuousFund{
 		Recipient:  msg.Recipient,
@@ -161,12 +166,23 @@ func (k MsgServer) CreateContinuousFund(ctx context.Context, msg *types.MsgCreat
 }
 
 func (k MsgServer) WithdrawContinuousFund(ctx context.Context, msg *types.MsgWithdrawContinuousFund) (*types.MsgWithdrawContinuousFundResponse, error) {
-	amount, err := k.withdrawContinuousFund(ctx, msg.RecipientAddress)
+	recipient, err := k.authKeeper.AddressCodec().StringToBytes(msg.RecipientAddress)
 	if err != nil {
-		return nil, err
+		return nil, sdkerrors.ErrInvalidAddress.Wrapf("invalid recipient address: %s", err)
 	}
 
-	return &types.MsgWithdrawContinuousFundResponse{Amount: amount}, nil
+	err = k.IterateAndUpdateFundsDistribution(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("error while iterating all the continuous funds: %w", err)
+	}
+
+	// withdraw continuous fund
+	withdrawnAmount, err := k.withdrawRecipientFunds(ctx, recipient)
+	if err != nil {
+		return nil, fmt.Errorf("error while withdrawing recipient funds for recipient: %w", err)
+	}
+
+	return &types.MsgWithdrawContinuousFundResponse{Amount: withdrawnAmount}, nil
 }
 
 func (k MsgServer) CancelContinuousFund(ctx context.Context, msg *types.MsgCancelContinuousFund) (*types.MsgCancelContinuousFundResponse, error) {
@@ -182,17 +198,14 @@ func (k MsgServer) CancelContinuousFund(ctx context.Context, msg *types.MsgCance
 	canceledHeight := k.HeaderService.HeaderInfo(ctx).Height
 	canceledTime := k.HeaderService.HeaderInfo(ctx).Time
 
-	found, err := k.ContinuousFund.Has(ctx, recipient)
-	if !found {
-		return nil, fmt.Errorf("no recipient found to cancel continuous fund: %s", msg.RecipientAddress)
-	}
-	if err != nil {
+	// distribute funds before withdrawing
+	if err = k.IterateAndUpdateFundsDistribution(ctx); err != nil {
 		return nil, err
 	}
 
 	// withdraw funds if any are allocated
 	withdrawnFunds, err := k.withdrawRecipientFunds(ctx, recipient)
-	if err != nil && !errorspkg.Is(err, types.ErrNoRecipientFund) {
+	if err != nil && !errorspkg.Is(err, types.ErrNoRecipientFound) {
 		return nil, fmt.Errorf("error while withdrawing already allocated funds for recipient %s: %w", msg.RecipientAddress, err)
 	}
 

x/protocolpool/keeper/msg_server_test.go

@@ -3,6 +3,8 @@ package keeper_test
 import (
 	"time"
 
+	"github.com/golang/mock/gomock"
+
 	"cosmossdk.io/collections"
 	"cosmossdk.io/core/header"
 	"cosmossdk.io/math"
@@ -403,7 +405,7 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 		"recipient with no continuous fund": {
 			recipientAddress: []sdk.AccAddress{recipient},
 			expErr:           true,
-			expErrMsg:        "no continuous fund found for recipient",
+			expErrMsg:        "error while withdrawing recipient funds for recipient: no recipient found",
 		},
 		"funds percentage > 100": {
 			preRun: func() {
@@ -441,14 +443,14 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 				suite.Require().NoError(err)
 
 				// Set ToDistribute
-				err = suite.poolKeeper.ToDistribute.Set(suite.ctx, math.NewInt(100000))
+				err = suite.poolKeeper.Distributions.Set(suite.ctx, suite.ctx.HeaderInfo().Time, math.NewInt(100000))
 				suite.Require().NoError(err)
 			},
 			recipientAddress: []sdk.AccAddress{recipient},
 			expErr:           true,
 			expErrMsg:        "error while iterating all the continuous funds: total funds percentage cannot exceed 100",
 		},
-		"expired case": {
+		"expired case with no funds left to withdraw": {
 			preRun: func() {
 				percentage, err := math.LegacyNewDecFromStr("0.2")
 				suite.Require().NoError(err)
@@ -464,7 +466,7 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 			},
 			recipientAddress: []sdk.AccAddress{recipient},
 			expErr:           true,
-			expErrMsg:        "cannot withdraw continuous funds: continuous fund expired for recipient",
+			expErrMsg:        "error while withdrawing recipient funds for recipient: no recipient found",
 		},
 		"valid case with ToDistribute amount zero": {
 			preRun: func() {
@@ -483,8 +485,9 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 				// Set recipient fund percentage and recipient fund distribution
 				err = suite.poolKeeper.RecipientFundDistribution.Set(suite.ctx, recipient, math.ZeroInt())
 				suite.Require().NoError(err)
-				err = suite.poolKeeper.ToDistribute.Set(suite.ctx, math.ZeroInt())
+				err = suite.poolKeeper.Distributions.Set(suite.ctx, suite.ctx.HeaderInfo().Time, math.ZeroInt())
 				suite.Require().NoError(err)
+				suite.mockStreamFunds(math.NewInt(0))
 			},
 			recipientAddress: []sdk.AccAddress{recipient},
 			expErr:           false,
@@ -504,9 +507,8 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 				// Set recipient fund percentage and recipient fund distribution
 				err = suite.poolKeeper.RecipientFundDistribution.Set(suite.ctx, recipient, math.ZeroInt())
 				suite.Require().NoError(err)
-				toDistribute := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(100000)))
-				suite.mockStreamFunds()
-				err = suite.poolKeeper.SetToDistribute(suite.ctx, toDistribute, suite.poolKeeper.GetAuthority())
+				suite.mockStreamFunds(math.NewInt(100000))
+				err = suite.poolKeeper.SetToDistribute(suite.ctx)
 				suite.Require().NoError(err)
 			},
 			recipientAddress: []sdk.AccAddress{recipient},
@@ -530,9 +532,8 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 				// Set recipient fund percentage and recipient fund distribution
 				err = suite.poolKeeper.RecipientFundDistribution.Set(suite.ctx, recipient, math.ZeroInt())
 				suite.Require().NoError(err)
-				toDistribute := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(100000)))
-				suite.mockStreamFunds()
-				err = suite.poolKeeper.SetToDistribute(suite.ctx, toDistribute, suite.poolKeeper.GetAuthority())
+				suite.mockStreamFunds(math.NewInt(100000))
+				err = suite.poolKeeper.SetToDistribute(suite.ctx)
 				suite.Require().NoError(err)
 			},
 			recipientAddress: []sdk.AccAddress{recipient},
@@ -588,9 +589,8 @@ func (suite *KeeperTestSuite) TestWithdrawContinuousFund() {
 				err = suite.poolKeeper.RecipientFundDistribution.Set(suite.ctx, recipient3, math.ZeroInt())
 				suite.Require().NoError(err)
 
-				toDistribute := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(100000)))
-				suite.mockStreamFunds()
-				err = suite.poolKeeper.SetToDistribute(suite.ctx, toDistribute, suite.poolKeeper.GetAuthority())
+				suite.mockStreamFunds(math.NewInt(100000))
+				err = suite.poolKeeper.SetToDistribute(suite.ctx)
 				suite.Require().NoError(err)
 			},
 			recipientAddress: []sdk.AccAddress{recipient, recipient2, recipient3},
@@ -813,11 +813,6 @@ func (suite *KeeperTestSuite) TestCancelContinuousFund() {
 			expErr:    true,
 			expErrMsg: "empty address string is not allowed",
 		},
-		"no recipient found": {
-			recipientAddr: recipientAddr,
-			expErr:        true,
-			expErrMsg:     "no recipient found to cancel continuous fund",
-		},
 		"all good with unclaimed funds for recipient": {
 			preRun: func() {
 				// Set fund 1
@@ -853,9 +848,8 @@ func (suite *KeeperTestSuite) TestCancelContinuousFund() {
 				suite.Require().NoError(err)
 
 				// Set ToDistribute
-				toDistribute := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(100000)))
-				suite.mockStreamFunds()
-				err = suite.poolKeeper.SetToDistribute(suite.ctx, toDistribute, suite.poolKeeper.GetAuthority())
+				suite.mockStreamFunds(math.NewInt(100000))
+				err = suite.poolKeeper.SetToDistribute(suite.ctx)
 				suite.Require().NoError(err)
 
 				// withdraw funds for fund request 2
@@ -960,9 +954,8 @@ func (suite *KeeperTestSuite) TestWithdrawExpiredFunds() {
 	})
 	suite.Require().NoError(err)
 
-	toDistribute := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(100000)))
-	suite.mockStreamFunds()
-	err = suite.poolKeeper.SetToDistribute(suite.ctx, toDistribute, suite.poolKeeper.GetAuthority())
+	suite.mockStreamFunds(math.NewInt(100000))
+	err = suite.poolKeeper.SetToDistribute(suite.ctx)
 	suite.Require().NoError(err)
 
 	suite.mockWithdrawContinuousFund()
@@ -973,11 +966,17 @@ func (suite *KeeperTestSuite) TestWithdrawExpiredFunds() {
 	header.Time = expiration.Add(1 * time.Second)
 	suite.ctx = suite.ctx.WithHeaderInfo(header)
 
-	_, err = suite.msgServer.WithdrawContinuousFund(suite.ctx, &types.MsgWithdrawContinuousFund{RecipientAddress: recipientStrAddr})
-	suite.Require().ErrorContains(err, "continuous fund expired for recipient")
+	// If we keep calling WithdrawContinuousFund, it should not error and return always an amount of 0
+	withdrawRes, err := suite.msgServer.WithdrawContinuousFund(suite.ctx, &types.MsgWithdrawContinuousFund{RecipientAddress: recipientStrAddr})
+	suite.Require().True(withdrawRes.Amount.IsZero())
+	suite.Require().NoError(err)
 
-	suite.mockStreamFunds()
-	err = suite.poolKeeper.SetToDistribute(suite.ctx, toDistribute, suite.poolKeeper.GetAuthority())
+	withdrawRes, err = suite.msgServer.WithdrawContinuousFund(suite.ctx, &types.MsgWithdrawContinuousFund{RecipientAddress: recipientStrAddr})
+	suite.Require().True(withdrawRes.Amount.IsZero())
+	suite.Require().NoError(err)
+
+	suite.mockStreamFunds(math.NewInt(100000))
+	err = suite.poolKeeper.SetToDistribute(suite.ctx)
 	suite.Require().NoError(err)
 
 	suite.mockWithdrawContinuousFund()
@@ -985,9 +984,58 @@ func (suite *KeeperTestSuite) TestWithdrawExpiredFunds() {
 	suite.Require().NoError(err)
 
 	res, err := suite.msgServer.CancelContinuousFund(suite.ctx, &types.MsgCancelContinuousFund{
+		Authority:        suite.poolKeeper.GetAuthority(),
+		RecipientAddress: recipient2StrAddr,
+	})
+	suite.Require().NoError(err)
+	suite.Require().Equal(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(0)), res.WithdrawnAllocatedFund)
+
+	// canceling an expired continuous fund, won't error
+	res, err = suite.msgServer.CancelContinuousFund(suite.ctx, &types.MsgCancelContinuousFund{
 		Authority:        suite.poolKeeper.GetAuthority(),
 		RecipientAddress: recipientStrAddr,
 	})
 	suite.Require().NoError(err)
 	suite.Require().Equal(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(0)), res.WithdrawnAllocatedFund)
+
+	// if we try to cancel again the same continuout fund, it won't error, it will still distribute funds if needed.
+	res, err = suite.msgServer.CancelContinuousFund(suite.ctx, &types.MsgCancelContinuousFund{
+		Authority:        suite.poolKeeper.GetAuthority(),
+		RecipientAddress: recipientStrAddr,
+	})
+	suite.Require().NoError(err)
+	suite.Require().True(res.WithdrawnAllocatedFund.IsNil())
+}
+
+func (suite *KeeperTestSuite) TestFundCommunityPool() {
+	sender := []byte("fundingAddr1____________________")
+	addrCodec := codectestutil.CodecOptions{}.GetAddressCodec()
+	senderAddr, err := addrCodec.BytesToString(sender)
+	suite.Require().NoError(err)
+
+	amount := sdk.NewCoins(sdk.NewInt64Coin(sdk.DefaultBondDenom, 1000000))
+	suite.bankKeeper.EXPECT().SendCoinsFromAccountToModule(gomock.Any(), sender, types.ModuleName, amount).Return(nil).Times(1)
+
+	_, err = suite.msgServer.FundCommunityPool(suite.ctx, &types.MsgFundCommunityPool{
+		Amount:    amount,
+		Depositor: senderAddr,
+	})
+	suite.Require().NoError(err)
+}
+
+func (suite *KeeperTestSuite) TestCommunityPoolSpend() {
+	recipient := []byte("fundingAddr1____________________")
+	addrCodec := codectestutil.CodecOptions{}.GetAddressCodec()
+	recipientAddr, err := addrCodec.BytesToString(recipient)
+	suite.Require().NoError(err)
+
+	amount := sdk.NewCoins(sdk.NewInt64Coin(sdk.DefaultBondDenom, 1000000))
+	suite.bankKeeper.EXPECT().SendCoinsFromModuleToAccount(gomock.Any(), types.ModuleName, recipient, amount).Return(nil).Times(1)
+
+	_, err = suite.msgServer.CommunityPoolSpend(suite.ctx, &types.MsgCommunityPoolSpend{
+		Authority: suite.poolKeeper.GetAuthority(),
+		Recipient: recipientAddr,
+		Amount:    amount,
+	})
+	suite.Require().NoError(err)
 }