refactor: use branch service in 29-fee (#7732)

* refactor: use branch service in 29-fee

* chore: move comment for readability

* refactor: use branch service in timeout fee distribution

* refactor: complete refactor to branch service and remove cached context

* chore: make lint-fix

Author: damiannolan

modules/apps/29-fee/keeper/escrow.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 
 	errorsmod "cosmossdk.io/errors"
@@ -11,6 +12,7 @@ import (
 
 	"github.com/cosmos/ibc-go/v9/modules/apps/29-fee/types"
 	channeltypes "github.com/cosmos/ibc-go/v9/modules/core/04-channel/types"
+	ibcerrors "github.com/cosmos/ibc-go/v9/modules/core/errors"
 )
 
 // escrowPacketFee sends the packet fee to the 29-fee module account to hold in escrow
@@ -46,38 +48,39 @@ func (k Keeper) escrowPacketFee(ctx context.Context, packetID channeltypes.Packe
 
 // DistributePacketFeesOnAcknowledgement pays all the acknowledgement & receive fees for a given packetID while refunding the timeout fees to the refund account.
 func (k Keeper) DistributePacketFeesOnAcknowledgement(ctx context.Context, forwardRelayer string, reverseRelayer sdk.AccAddress, packetFees []types.PacketFee, packetID channeltypes.PacketId) {
-	// cache context before trying to distribute fees
+	// use branched multistore for distribution of fees.
 	// if the escrow account has insufficient balance then we want to avoid partially distributing fees
-	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO: https://github.com/cosmos/ibc-go/issues/5917
-	cacheCtx, writeFn := sdkCtx.CacheContext()
+	if err := k.BranchService.Execute(ctx, func(ctx context.Context) error {
+		// forward relayer address will be empty if conversion fails
+		forwardAddr, _ := sdk.AccAddressFromBech32(forwardRelayer)
+
+		for _, packetFee := range packetFees {
+			if !k.EscrowAccountHasBalance(ctx, packetFee.Fee.Total()) {
+				// NOTE: we lock the fee module on error return so that the state changes are persisted
+				return ibcerrors.ErrInsufficientFunds
+			}
+
+			// check if refundAcc address works
+			refundAddr, err := sdk.AccAddressFromBech32(packetFee.RefundAddress)
+			if err != nil {
+				panic(fmt.Errorf("could not parse refundAcc %s to sdk.AccAddress", packetFee.RefundAddress))
+			}
 
-	// forward relayer address will be empty if conversion fails
-	forwardAddr, _ := sdk.AccAddressFromBech32(forwardRelayer)
+			k.distributePacketFeeOnAcknowledgement(ctx, refundAddr, forwardAddr, reverseRelayer, packetFee)
+		}
 
-	for _, packetFee := range packetFees {
-		if !k.EscrowAccountHasBalance(cacheCtx, packetFee.Fee.Total()) {
+		return nil
+	}); err != nil {
+		if errors.Is(err, ibcerrors.ErrInsufficientFunds) {
 			// if the escrow account does not have sufficient funds then there must exist a severe bug
 			// the fee module should be locked until manual intervention fixes the issue
 			// a locked fee module will simply skip fee logic, all channels will temporarily function as
 			// fee disabled channels
-			// NOTE: we use the uncached context to lock the fee module so that the state changes from
-			// locking the fee module are persisted
 			k.lockFeeModule(ctx)
 			return
 		}
-
-		// check if refundAcc address works
-		refundAddr, err := sdk.AccAddressFromBech32(packetFee.RefundAddress)
-		if err != nil {
-			panic(fmt.Errorf("could not parse refundAcc %s to sdk.AccAddress", packetFee.RefundAddress))
-		}
-
-		k.distributePacketFeeOnAcknowledgement(cacheCtx, refundAddr, forwardAddr, reverseRelayer, packetFee)
 	}
 
-	// write the cache
-	writeFn()
-
 	// removes the fees from the store as fees are now paid
 	k.DeleteFeesInEscrow(ctx, packetID)
 }
@@ -104,35 +107,36 @@ func (k Keeper) distributePacketFeeOnAcknowledgement(ctx context.Context, refund
 
 // DistributePacketFeesOnTimeout pays all the timeout fees for a given packetID while refunding the acknowledgement & receive fees to the refund account.
 func (k Keeper) DistributePacketFeesOnTimeout(ctx context.Context, timeoutRelayer sdk.AccAddress, packetFees []types.PacketFee, packetID channeltypes.PacketId) {
-	// cache context before trying to distribute fees
+	// use branched multistore for distribution of fees.
 	// if the escrow account has insufficient balance then we want to avoid partially distributing fees
-	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO: https://github.com/cosmos/ibc-go/issues/5917
-	cacheCtx, writeFn := sdkCtx.CacheContext()
+	if err := k.BranchService.Execute(ctx, func(ctx context.Context) error {
+		for _, packetFee := range packetFees {
+			if !k.EscrowAccountHasBalance(ctx, packetFee.Fee.Total()) {
+				// NOTE: we lock the fee module on error return so that the state changes are persisted
+				return ibcerrors.ErrInsufficientFunds
+			}
 
-	for _, packetFee := range packetFees {
-		if !k.EscrowAccountHasBalance(cacheCtx, packetFee.Fee.Total()) {
+			// check if refundAcc address works
+			refundAddr, err := sdk.AccAddressFromBech32(packetFee.RefundAddress)
+			if err != nil {
+				panic(fmt.Errorf("could not parse refundAcc %s to sdk.AccAddress", packetFee.RefundAddress))
+			}
+
+			k.distributePacketFeeOnTimeout(ctx, refundAddr, timeoutRelayer, packetFee)
+		}
+
+		return nil
+	}); err != nil {
+		if errors.Is(err, ibcerrors.ErrInsufficientFunds) {
 			// if the escrow account does not have sufficient funds then there must exist a severe bug
 			// the fee module should be locked until manual intervention fixes the issue
 			// a locked fee module will simply skip fee logic, all channels will temporarily function as
 			// fee disabled channels
-			// NOTE: we use the uncached context to lock the fee module so that the state changes from
-			// locking the fee module are persisted
 			k.lockFeeModule(ctx)
 			return
 		}
-
-		// check if refundAcc address works
-		refundAddr, err := sdk.AccAddressFromBech32(packetFee.RefundAddress)
-		if err != nil {
-			panic(fmt.Errorf("could not parse refundAcc %s to sdk.AccAddress", packetFee.RefundAddress))
-		}
-
-		k.distributePacketFeeOnTimeout(cacheCtx, refundAddr, timeoutRelayer, packetFee)
 	}
 
-	// write the cache
-	writeFn()
-
 	// removing the fee from the store as the fee is now paid
 	k.DeleteFeesInEscrow(ctx, packetID)
 }
@@ -151,36 +155,36 @@ func (k Keeper) distributePacketFeeOnTimeout(ctx context.Context, refundAddr, ti
 // If the distribution fails for any reason (such as the receiving address being blocked),
 // the state changes will be discarded.
 func (k Keeper) distributeFee(ctx context.Context, receiver, refundAccAddress sdk.AccAddress, fee sdk.Coins) {
-	// cache context before trying to distribute fees
-	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO: https://github.com/cosmos/ibc-go/issues/7223
-	cacheCtx, writeFn := sdkCtx.CacheContext()
+	// use branched multistore before trying to distribute fees
+	if err := k.BranchService.Execute(ctx, func(ctx context.Context) error {
+		err := k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, receiver, fee)
+		if err != nil {
+			if bytes.Equal(receiver, refundAccAddress) {
+				// if sending to the refund address already failed, then return (no-op)
+				return errorsmod.Wrapf(types.ErrRefundDistributionFailed, "receiver address: %s", receiver)
+			}
 
-	err := k.bankKeeper.SendCoinsFromModuleToAccount(cacheCtx, types.ModuleName, receiver, fee)
-	if err != nil {
-		if bytes.Equal(receiver, refundAccAddress) {
-			k.Logger.Error("error distributing fee", "receiver address", receiver, "fee", fee)
-			return // if sending to the refund address already failed, then return (no-op)
-		}
+			// if an error is returned from x/bank and the receiver is not the refundAccAddress
+			// then attempt to refund the fee to the original sender
+			err := k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, refundAccAddress, fee)
+			if err != nil {
+				// if sending to the refund address fails, no-op
+				return errorsmod.Wrapf(types.ErrRefundDistributionFailed, "receiver address: %s", refundAccAddress)
+			}
 
-		// if an error is returned from x/bank and the receiver is not the refundAccAddress
-		// then attempt to refund the fee to the original sender
-		err := k.bankKeeper.SendCoinsFromModuleToAccount(cacheCtx, types.ModuleName, refundAccAddress, fee)
-		if err != nil {
-			k.Logger.Error("error refunding fee to the original sender", "refund address", refundAccAddress, "fee", fee)
-			return // if sending to the refund address fails, no-op
+			if err := k.emitDistributeFeeEvent(ctx, refundAccAddress.String(), fee); err != nil {
+				panic(err)
+			}
 		}
 
-		if err := k.emitDistributeFeeEvent(ctx, refundAccAddress.String(), fee); err != nil {
-			panic(err)
-		}
-	} else {
-		if err := k.emitDistributeFeeEvent(ctx, receiver.String(), fee); err != nil {
-			panic(err)
-		}
+		return nil
+	}); err != nil {
+		k.Logger.Error("error distributing fee", "error", err.Error())
 	}
 
-	// write the cache
-	writeFn()
+	if err := k.emitDistributeFeeEvent(ctx, receiver.String(), fee); err != nil {
+		panic(err)
+	}
 }
 
 // RefundFeesOnChannelClosure will refund all fees associated with the given port and channel identifiers.
@@ -190,52 +194,52 @@ func (k Keeper) distributeFee(ctx context.Context, receiver, refundAccAddress sd
 func (k Keeper) RefundFeesOnChannelClosure(ctx context.Context, portID, channelID string) error {
 	identifiedPacketFees := k.GetIdentifiedPacketFeesForChannel(ctx, portID, channelID)
 
-	// cache context before trying to distribute fees
+	// use branched multistore for distribution of fees.
 	// if the escrow account has insufficient balance then we want to avoid partially distributing fees
-	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO: https://github.com/cosmos/ibc-go/issues/5917
-	cacheCtx, writeFn := sdkCtx.CacheContext()
-
-	for _, identifiedPacketFee := range identifiedPacketFees {
-		var unRefundedFees []types.PacketFee
-		for _, packetFee := range identifiedPacketFee.PacketFees {
-
-			if !k.EscrowAccountHasBalance(cacheCtx, packetFee.Fee.Total()) {
-				// if the escrow account does not have sufficient funds then there must exist a severe bug
-				// the fee module should be locked until manual intervention fixes the issue
-				// a locked fee module will simply skip fee logic, all channels will temporarily function as
-				// fee disabled channels
-				// NOTE: we use the uncached context to lock the fee module so that the state changes from
-				// locking the fee module are persisted
-				k.lockFeeModule(ctx)
-
-				// return a nil error so state changes are committed but distribution stops
-				return nil
+	if err := k.BranchService.Execute(ctx, func(ctx context.Context) error {
+		for _, identifiedPacketFee := range identifiedPacketFees {
+			var unRefundedFees []types.PacketFee
+			for _, packetFee := range identifiedPacketFee.PacketFees {
+
+				if !k.EscrowAccountHasBalance(ctx, packetFee.Fee.Total()) {
+					// NOTE: we lock the fee module on error return so that the state changes are persisted
+					return ibcerrors.ErrInsufficientFunds
+				}
+
+				refundAddr, err := sdk.AccAddressFromBech32(packetFee.RefundAddress)
+				if err != nil {
+					unRefundedFees = append(unRefundedFees, packetFee)
+					continue
+				}
+
+				// refund all fees to refund address
+				if err = k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, refundAddr, packetFee.Fee.Total()); err != nil {
+					unRefundedFees = append(unRefundedFees, packetFee)
+					continue
+				}
 			}
 
-			refundAddr, err := sdk.AccAddressFromBech32(packetFee.RefundAddress)
-			if err != nil {
-				unRefundedFees = append(unRefundedFees, packetFee)
-				continue
-			}
-
-			// refund all fees to refund address
-			if err = k.bankKeeper.SendCoinsFromModuleToAccount(cacheCtx, types.ModuleName, refundAddr, packetFee.Fee.Total()); err != nil {
-				unRefundedFees = append(unRefundedFees, packetFee)
-				continue
+			if len(unRefundedFees) > 0 {
+				// update packet fees to keep only the unrefunded fees
+				packetFees := types.NewPacketFees(unRefundedFees)
+				k.SetFeesInEscrow(ctx, identifiedPacketFee.PacketId, packetFees)
+			} else {
+				k.DeleteFeesInEscrow(ctx, identifiedPacketFee.PacketId)
 			}
 		}
 
-		if len(unRefundedFees) > 0 {
-			// update packet fees to keep only the unrefunded fees
-			packetFees := types.NewPacketFees(unRefundedFees)
-			k.SetFeesInEscrow(cacheCtx, identifiedPacketFee.PacketId, packetFees)
-		} else {
-			k.DeleteFeesInEscrow(cacheCtx, identifiedPacketFee.PacketId)
+		return nil
+	}); err != nil {
+		if errors.Is(err, ibcerrors.ErrInsufficientFunds) {
+			// if the escrow account does not have sufficient funds then there must exist a severe bug
+			// the fee module should be locked until manual intervention fixes the issue
+			// a locked fee module will simply skip fee logic, all channels will temporarily function as
+			// fee disabled channels
+			k.lockFeeModule(ctx)
+
+			return nil // commit state changes to lock module and stop fee distribution
 		}
 	}
 
-	// write the cache
-	writeFn()
-
 	return nil
 }

modules/apps/29-fee/types/errors.go

@@ -17,4 +17,5 @@ var (
 	ErrRelayerNotFoundForAsyncAck    = errorsmod.Register(ModuleName, 10, "relayer address must be stored for async WriteAcknowledgement")
 	ErrFeeModuleLocked               = errorsmod.Register(ModuleName, 11, "the fee module is currently locked, a severe bug has been detected")
 	ErrUnsupportedAction             = errorsmod.Register(ModuleName, 12, "unsupported action")
+	ErrRefundDistributionFailed      = errorsmod.Register(ModuleName, 13, "refund distribution failed")
 )