[SessionManager] Skip claims creation if supplier operator balance is…

… too low (#817)

## Summary

This PR introduces filters in the `SubmitClaims` step to ensure that the
`SessionManager` only submits claims it can afford to prove.

Before each submission, it checks the `SupplierOperator` balance and
filters the claims based on the available funds.

Claims are sorted by descending profitability to maximize rewards.

## Issue

Since submitting a proof incurs a cost, a `SupplierOperator` without
sufficient funds would be unable to submit the required proof.

This could result in the supplier's stake being slashed if a proof
submission is required.

The `SessionManager` must verify that it can afford a proof submission
before creating the corresponding claim.

## Type of change

Select one or more from the following:

- [x] New feature, functionality or library
- [ ] Consensus breaking; add the `consensus-breaking` label if so. See
#791 for details
- [ ] Bug fix
- [ ] Code health or cleanup
- [ ] Documentation
- [ ] Other (specify)

## Testing

- [x] **Documentation**: `make docusaurus_start`; only needed if you
make doc changes
- [x] **Unit Tests**: `make go_develop_and_test`
- [x] **LocalNet E2E Tests**: `make test_e2e`
- [ ] **DevNet E2E Tests**: Add the `devnet-test-e2e` label to the PR.

## Sanity Checklist

- [x] I have tested my changes using the available tooling
- [x] I have commented my code
- [x] I have performed a self-review of my own code; both comments &
source code
- [ ] I create and reference any new tickets, if applicable
- [ ] I have left TODOs throughout the codebase, if applicable

---------

Co-authored-by: Daniel Olshansky <[email protected]>

pkg/client/interface.go

@@ -11,6 +11,7 @@
 //go:generate mockgen -destination=../../testutil/mockclient/proof_query_client_mock.go -package=mockclient . ProofQueryClient
 //go:generate mockgen -destination=../../testutil/mockclient/tokenomics_query_client_mock.go -package=mockclient . TokenomicsQueryClient
 //go:generate mockgen -destination=../../testutil/mockclient/service_query_client_mock.go -package=mockclient . ServiceQueryClient
+//go:generate mockgen -destination=../../testutil/mockclient/bank_query_client_mock.go -package=mockclient . BankQueryClient
 //go:generate mockgen -destination=../../testutil/mockclient/cosmos_tx_builder_mock.go -package=mockclient github.com/cosmos/cosmos-sdk/client TxBuilder
 //go:generate mockgen -destination=../../testutil/mockclient/cosmos_keyring_mock.go -package=mockclient github.com/cosmos/cosmos-sdk/crypto/keyring Keyring
 //go:generate mockgen -destination=../../testutil/mockclient/cosmos_client_mock.go -package=mockclient github.com/cosmos/cosmos-sdk/client AccountRetriever
@@ -371,3 +372,10 @@ type ServiceQueryClient interface {
 	// GetService queries the chain for the details of the service provided
 	GetService(ctx context.Context, serviceId string) (sharedtypes.Service, error)
 }
+
+// BankQueryClient defines an interface that enables the querying of the
+// on-chain bank information
+type BankQueryClient interface {
+	// GetBalance queries the chain for the uPOKT balance of the account provided
+	GetBalance(ctx context.Context, address string) (*cosmostypes.Coin, error)
+}

pkg/client/query/bankquerier.go

@@ -0,0 +1,57 @@
+package query
+
+import (
+	"context"
+
+	"cosmossdk.io/depinject"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	grpc "github.com/cosmos/gogoproto/grpc"
+
+	"github.com/pokt-network/poktroll/app/volatile"
+	"github.com/pokt-network/poktroll/pkg/client"
+)
+
+var _ client.BankQueryClient = (*bankQuerier)(nil)
+
+// bankQuerier is a wrapper around the banktypes.QueryClient that enables the
+// querying of on-chain balance information.
+type bankQuerier struct {
+	clientConn  grpc.ClientConn
+	bankQuerier banktypes.QueryClient
+}
+
+// NewBankQuerier returns a new instance of a client.BankQueryClient by
+// injecting the dependecies provided by the depinject.Config.
+//
+// Required dependencies:
+// - clientCtx
+func NewBankQuerier(deps depinject.Config) (client.BankQueryClient, error) {
+	bq := &bankQuerier{}
+
+	if err := depinject.Inject(
+		deps,
+		&bq.clientConn,
+	); err != nil {
+		return nil, err
+	}
+
+	bq.bankQuerier = banktypes.NewQueryClient(bq.clientConn)
+
+	return bq, nil
+}
+
+// GetBalance returns the uPOKT balance of a given address
+func (bq *bankQuerier) GetBalance(
+	ctx context.Context,
+	address string,
+) (*sdk.Coin, error) {
+	// Query the blockchain for the balance record
+	req := &banktypes.QueryBalanceRequest{Address: address, Denom: volatile.DenomuPOKT}
+	res, err := bq.bankQuerier.Balance(ctx, req)
+	if err != nil {
+		return nil, ErrQueryBalanceNotFound.Wrapf("address: %s [%s]", address, err)
+	}
+
+	return res.Balance, nil
+}

pkg/client/query/errors.go

@@ -10,4 +10,5 @@ var (
 	ErrQueryPubKeyNotFound             = sdkerrors.Register(codespace, 4, "account pub key not found")
 	ErrQuerySessionParams              = sdkerrors.Register(codespace, 5, "unable to query session params")
 	ErrQueryRetrieveService            = sdkerrors.Register(codespace, 6, "error while trying to retrieve a service")
+	ErrQueryBalanceNotFound            = sdkerrors.Register(codespace, 7, "balance not found")
 )

pkg/deps/config/suppliers.go

@@ -502,6 +502,24 @@ func NewSupplyServiceQueryClientFn() SupplierFn {
 	}
 }
 
+// NewSupplyBankQuerierFn supplies a depinject config with an BankQuerier.
+func NewSupplyBankQuerierFn() SupplierFn {
+	return func(
+		_ context.Context,
+		deps depinject.Config,
+		_ *cobra.Command,
+	) (depinject.Config, error) {
+		// Create the bank querier.
+		bankQuerier, err := query.NewBankQuerier(deps)
+		if err != nil {
+			return nil, err
+		}
+
+		// Supply the bank querier to the provided deps
+		return depinject.Configs(deps, depinject.Supply(bankQuerier)), nil
+	}
+}
+
 // newSupplyTxClientFn returns a new depinject.Config which is supplied with
 // the given deps and the new TxClient.
 func newSupplyTxClientsFn(ctx context.Context, deps depinject.Config, signingKeyName string) (depinject.Config, error) {

pkg/relayer/cmd/cmd.go

@@ -197,6 +197,7 @@ func setupRelayerDependencies(
 		config.NewSupplyTokenomicsQueryClientFn(),
 		supplyMiner,
 		config.NewSupplyAccountQuerierFn(),
+		config.NewSupplyBankQuerierFn(),
 		config.NewSupplyApplicationQuerierFn(),
 		config.NewSupplySupplierQuerierFn(),
 		config.NewSupplySessionQuerierFn(),

pkg/relayer/session/claim.go

@@ -3,6 +3,7 @@ package session
 import (
 	"context"
 	"fmt"
+	"slices"
 
 	"github.com/pokt-network/poktroll/pkg/client"
 	"github.com/pokt-network/poktroll/pkg/either"
@@ -13,6 +14,7 @@ import (
 	"github.com/pokt-network/poktroll/pkg/relayer"
 	prooftypes "github.com/pokt-network/poktroll/x/proof/types"
 	"github.com/pokt-network/poktroll/x/shared"
+	"github.com/pokt-network/smt"
 )
 
 // createClaims maps over the sessionsToClaimObs observable. For each claim batch, it:
@@ -180,13 +182,16 @@ func (rs *relayerSessionsManager) newMapClaimSessionsFn(
 			return either.Success(sessionTrees), false
 		}
 
-		// TODO_FOLLOWUP(@red-0ne): Ensure that the supplier operator account
-		// has enough funds to cover for any potential proof submission in order to
-		// avoid slashing due to missing proofs.
-		// We should order the claimMsgs by reward amount and include claims up to
-		// whatever the supplier can afford to cover.
-		claimMsgs := make([]client.MsgCreateClaim, len(sessionTrees))
-		for idx, sessionTree := range sessionTrees {
+		// Filter out the session trees that the supplier operator can afford to claim.
+		claimableSessionTrees, err := rs.payableProofsSessionTrees(ctx, sessionTrees)
+		if err != nil {
+			failedCreateClaimsSessionsPublishCh <- sessionTrees
+			rs.logger.Error().Err(err).Msg("failed to calculate payable proofs session trees")
+			return either.Error[[]relayer.SessionTree](err), false
+		}
+
+		claimMsgs := make([]client.MsgCreateClaim, len(claimableSessionTrees))
+		for idx, sessionTree := range claimableSessionTrees {
 			claimMsgs[idx] = &prooftypes.MsgCreateClaim{
 				RootHash:                sessionTree.GetClaimRoot(),
 				SessionHeader:           sessionTree.GetSessionHeader(),
@@ -196,12 +201,12 @@ func (rs *relayerSessionsManager) newMapClaimSessionsFn(
 
 		// Create claims for each supplier operator address in `sessionTrees`.
 		if err := supplierClient.CreateClaims(ctx, claimMsgs...); err != nil {
-			failedCreateClaimsSessionsPublishCh <- sessionTrees
+			failedCreateClaimsSessionsPublishCh <- claimableSessionTrees
 			rs.logger.Error().Err(err).Msg("failed to create claims")
 			return either.Error[[]relayer.SessionTree](err), false
 		}
 
-		return either.Success(sessionTrees), false
+		return either.Success(claimableSessionTrees), false
 	}
 }
 
@@ -235,3 +240,85 @@ func (rs *relayerSessionsManager) goCreateClaimRoots(
 	failSubmitProofsSessionsCh <- failedClaims
 	claimsFlushedCh <- flushedClaims
 }
+
+// payableProofsSessionTrees returns the session trees that the supplier operator
+// can afford to claim (i.e. pay the fee for submitting a proof).
+// The session trees are sorted from the most rewarding to the least rewarding to
+// ensure optimal rewards in the case of insufficient funds.
+// Note that all sessionTrees are associated with the same supplier operator address.
+func (rs *relayerSessionsManager) payableProofsSessionTrees(
+	ctx context.Context,
+	sessionTrees []relayer.SessionTree,
+) ([]relayer.SessionTree, error) {
+	supplierOpeartorAddress := sessionTrees[0].GetSupplierOperatorAddress().String()
+	logger := rs.logger.With(
+		"supplier_operator_address", supplierOpeartorAddress,
+	)
+
+	proofParams, err := rs.proofQueryClient.GetParams(ctx)
+	if err != nil {
+		return nil, err
+	}
+	proofSubmissionFeeCoin := proofParams.GetProofSubmissionFee()
+
+	supplierOperatorBalanceCoin, err := rs.bankQueryClient.GetBalance(
+		ctx,
+		sessionTrees[0].GetSupplierOperatorAddress().String(),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Sort the session trees by the sum of the claim root to ensure that the
+	// most rewarding claims are claimed first.
+	slices.SortFunc(sessionTrees, func(a, b relayer.SessionTree) int {
+		rootA := a.GetClaimRoot()
+		sumA, errA := smt.MerkleSumRoot(rootA).Sum()
+		if errA != nil {
+			logger.With(
+				"session_id", a.GetSessionHeader().GetSessionId(),
+				"claim_root", fmt.Sprintf("%x", rootA),
+			).Error().Err(errA).Msg("failed to calculate sum of claim root, assuming 0")
+			sumA = 0
+		}
+
+		rootB := b.GetClaimRoot()
+		sumB, errB := smt.MerkleSumRoot(rootB).Sum()
+		if errB != nil {
+			logger.With(
+				"session_id", a.GetSessionHeader().GetSessionId(),
+				"claim_root", fmt.Sprintf("%x", rootA),
+			).Error().Err(errB).Msg("failed to calculate sum of claim root, assuming 0")
+			sumB = 0
+		}
+
+		// Sort in descending order.
+		return int(sumB - sumA)
+	})
+
+	claimableSessionTrees := []relayer.SessionTree{}
+	for _, sessionTree := range sessionTrees {
+		// If the supplier operator can afford to claim the session, add it to the
+		// claimableSessionTrees slice.
+		if supplierOperatorBalanceCoin.IsGTE(*proofSubmissionFeeCoin) {
+			claimableSessionTrees = append(claimableSessionTrees, sessionTree)
+			newSupplierOperatorBalanceCoin := supplierOperatorBalanceCoin.Sub(*proofSubmissionFeeCoin)
+			supplierOperatorBalanceCoin = &newSupplierOperatorBalanceCoin
+			continue
+		}
+
+		// Log a warning of any session that the supplier operator cannot afford to claim.
+		logger.With(
+			"session_id", sessionTree.GetSessionHeader().GetSessionId(),
+			"supplier_operator_balance", supplierOperatorBalanceCoin,
+			"proof_submission_fee", proofSubmissionFeeCoin,
+		).Warn().Msg("supplier operator cannot afford to submit proof for claim, skipping")
+	}
+
+	logger.Warn().Msgf(
+		"Supplier operator %q can only afford %d out of %d claims",
+		supplierOpeartorAddress, len(claimableSessionTrees), len(sessionTrees),
+	)
+
+	return claimableSessionTrees, nil
+}

pkg/relayer/session/session.go

@@ -67,6 +67,9 @@ type relayerSessionsManager struct {
 
 	// tokenomicsQueryClient is used to query for the tokenomics module parameters.
 	tokenomicsQueryClient client.TokenomicsQueryClient
+
+	// bankQueryClient is used to query for the bank module parameters.
+	bankQueryClient client.BankQueryClient
 }
 
 // NewRelayerSessions creates a new relayerSessions.
@@ -98,6 +101,7 @@ func NewRelayerSessions(
 		&rs.serviceQueryClient,
 		&rs.proofQueryClient,
 		&rs.tokenomicsQueryClient,
+		&rs.bankQueryClient,
 	); err != nil {
 		return nil, err
 	}