[RelayMiner] Supplier rate limiting (#895)

## Summary

This PR adds an optimistic `RelayMeter` to the `RelayMiner` to monitor
the consumed stake and stops servicing when `maxStake` is reached.

It works by
* Intercepting `RelayRequest`s before being served.
* Assume the relay will be mined and deduce the corresponding amount
from the consumed stake.
* Send a `rate-limited` error if app stake share cannot cover of a
minable `RelayRequest`.
* After a relay has been served check if the `Relay` is volume
applicable and revert the initial deduction if it is not.
* Updates the known `Application`s stake if an `EventApplicationStaked`
is observed.
* Resets the `Application`s stakes at the beginning of every session.

## Issue

`RelayMiner`s need to know how much of the `Application`'s stake they
are allowed to consume without over-servicing.

![image](https://github.com/user-attachments/assets/4354c747-23a6-4b44-a63a-62a13617d796)

## Type of change

Select one or more from the following:

- [x] New feature, functionality or library

## Testing

- [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

---------

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

config.yml

@@ -141,13 +141,15 @@ genesis:
             denom: upokt
     bank:
       supply:
-        - amount: "1102000204"
+        - amount: "1202000272"
           denom: upokt
       balances:
         # Application module
         - address: pokt1rl3gjgzexmplmds3tq3r3yk84zlwdl6djzgsvm
           coins:
-            - amount: "100000068" # MUST BE equal to the total of all app stakes below
+            # TODO_MAINNET(@olshansk): Pass config.yml into ChatGPT to build a script
+            # that ensures the amounts line up
+            - amount: "200000136" # MUST BE equal to the total of all app stakes below
               denom: upokt
         # Supplier module
         - address: pokt1j40dzzmn6cn9kxku7a5tjnud6hv37vesr5ccaa

pkg/client/interface.go

@@ -291,6 +291,9 @@ type SessionQueryClient interface {
 		serviceId string,
 		blockHeight int64,
 	) (*sessiontypes.Session, error)
+
+	// GetParams queries the chain for the session module parameters.
+	GetParams(ctx context.Context) (*sessiontypes.Params, error)
 }
 
 // SharedQueryClient defines an interface that enables the querying of the

pkg/client/query/sessionquerier.go

@@ -62,3 +62,13 @@ func (sessq *sessionQuerier) GetSession(
 	}
 	return res.Session, nil
 }
+
+// GetParams queries & returns the session module on-chain parameters.
+func (sessq *sessionQuerier) GetParams(ctx context.Context) (*sessiontypes.Params, error) {
+	req := &sessiontypes.QueryParamsRequest{}
+	res, err := sessq.sessionQuerier.Params(ctx, req)
+	if err != nil {
+		return nil, ErrQuerySessionParams.Wrapf("[%v]", err)
+	}
+	return &res.Params, nil
+}

pkg/relayer/cmd/cmd.go

@@ -198,12 +198,13 @@ func setupRelayerDependencies(
 		config.NewSupplyDelegationClientFn(),                              // leaf
 		config.NewSupplySharedQueryClientFn(),                             // leaf
 		config.NewSupplyServiceQueryClientFn(),
+		config.NewSupplyApplicationQuerierFn(),
+		config.NewSupplySessionQuerierFn(),
+		supplyRelayMeter,
 		supplyMiner,
 		config.NewSupplyAccountQuerierFn(),
 		config.NewSupplyBankQuerierFn(),
-		config.NewSupplyApplicationQuerierFn(),
 		config.NewSupplySupplierQuerierFn(),
-		config.NewSupplySessionQuerierFn(),
 		config.NewSupplyProofQueryClientFn(),
 		config.NewSupplyRingCacheFn(),
 		supplyTxFactory,
@@ -231,6 +232,21 @@ func supplyMiner(
 	return depinject.Configs(deps, depinject.Supply(mnr)), nil
 }
 
+// supplyRelayMeter constructs a RelayMeter instance and returns a new
+// depinject.Config which is supplied with the given deps and the new RelayMeter.
+func supplyRelayMeter(
+	_ context.Context,
+	deps depinject.Config,
+	_ *cobra.Command,
+) (depinject.Config, error) {
+	rm, err := proxy.NewRelayMeter(deps)
+	if err != nil {
+		return nil, err
+	}
+
+	return depinject.Configs(deps, depinject.Supply(rm)), nil
+}
+
 // supplyTxFactory constructs a cosmostx.Factory instance and returns a new
 // depinject.Config which is supplied with the given deps and the new
 // cosmostx.Factory.

pkg/relayer/interface.go

@@ -1,6 +1,7 @@
 //go:generate mockgen -destination=../../testutil/mockrelayer/relayer_proxy_mock.go -package=mockrelayer . RelayerProxy
 //go:generate mockgen -destination=../../testutil/mockrelayer/miner_mock.go -package=mockrelayer . Miner
 //go:generate mockgen -destination=../../testutil/mockrelayer/relayer_sessions_manager_mock.go -package=mockrelayer . RelayerSessionsManager
+//go:generate mockgen -destination=../../testutil/mockrelayer/relay_meter_mock.go -package=mockrelayer . RelayMeter
 
 package relayer
 
@@ -162,3 +163,29 @@ type SessionTree interface {
 	// GetTrieSpec returns the trie spec of the SMST.
 	GetTrieSpec() smt.TrieSpec
 }
+
+// RelayMeter is an interface that keeps track of the amount of stake consumed between
+// a single onchain Application and a single onchain Supplier over the course of a single session.
+// It enables the RelayMiner to rate limit the number of requests handled offchain as a function
+// of the optimistic onchain rate limits.
+type RelayMeter interface {
+	// Start starts the relay meter.
+	Start(ctx context.Context) error
+
+	// AccumulateRelayReward adds the relay reward from the incoming request to session's accumulator.
+	// The relay cost is added optimistically, assuming that the relay WILL be volume / reward applicable.
+	//
+	// The reason why optimistic AccumulateRelayReward + SetNonApplicableRelayReward is used instead of
+	// a simpler AccumulateVolumeApplicableRelayReward is that when the relay is first seen
+	// we don't know if it will be volume / reward applicable until it is served.
+	//
+	// To rate limit or not the current relay, we need to always optimistically account all relays as being
+	// volume / reward applicable.
+	AccumulateRelayReward(ctx context.Context, relayRequestMeta servicetypes.RelayRequestMetadata) error
+
+	// SetNonApplicableRelayReward updates the relay meter for the given relay request as
+	// non-applicable between a single Application and a single Supplier for a single session.
+	// The volume / reward applicability of the relay is unknown to the relay miner
+	// until the relay is served and the relay response signed.
+	SetNonApplicableRelayReward(ctx context.Context, relayRequestMeta servicetypes.RelayRequestMetadata) error
+}

pkg/relayer/miner/miner.go

@@ -27,6 +27,7 @@ type miner struct {
 	// serviceQueryClient is used to query for the relay difficulty target hash of a service.
 	// relay_difficulty is the target hash which a relay hash must be less than to be volume/reward applicable.
 	serviceQueryClient client.ServiceQueryClient
+	relayMeter         relayer.RelayMeter
 }
 
 // NewMiner creates a new miner from the given dependencies and options. It
@@ -43,7 +44,7 @@ func NewMiner(
 ) (*miner, error) {
 	mnr := &miner{}
 
-	if err := depinject.Inject(deps, &mnr.serviceQueryClient); err != nil {
+	if err := depinject.Inject(deps, &mnr.serviceQueryClient, &mnr.relayMeter); err != nil {
 		return nil, err
 	}
 
@@ -88,18 +89,27 @@ func (mnr *miner) mapMineRelay(
 ) (_ either.Either[*relayer.MinedRelay], skip bool) {
 	relayBz, err := relay.Marshal()
 	if err != nil {
+		if relayMeteringResult := mnr.unclaimRelayUPOKT(ctx, *relay); relayMeteringResult.IsError() {
+			return relayMeteringResult, false
+		}
 		return either.Error[*relayer.MinedRelay](err), false
 	}
 	relayHashArr := protocol.GetRelayHashFromBytes(relayBz)
 	relayHash := relayHashArr[:]
 
 	relayDifficultyTargetHash, err := mnr.getServiceRelayDifficultyTargetHash(ctx, relay.Req)
 	if err != nil {
-		return either.Error[*relayer.MinedRelay](err), false
+		if relayMeteringResult := mnr.unclaimRelayUPOKT(ctx, *relay); relayMeteringResult.IsError() {
+			return relayMeteringResult, true
+		}
+		return either.Error[*relayer.MinedRelay](err), true
 	}
 
 	// The relay IS NOT volume / reward applicable
 	if !protocol.IsRelayVolumeApplicable(relayHash, relayDifficultyTargetHash) {
+		if eitherMeteringResult := mnr.unclaimRelayUPOKT(ctx, *relay); eitherMeteringResult.IsError() {
+			return eitherMeteringResult, true
+		}
 		return either.Success[*relayer.MinedRelay](nil), true
 	}
 
@@ -135,3 +145,13 @@ func (mnr *miner) getServiceRelayDifficultyTargetHash(ctx context.Context, req *
 
 	return serviceRelayDifficulty.GetTargetHash(), nil
 }
+
+// unclaimRelayUPOKT unclaims the relay UPOKT reward for the relay.
+// It returns an either.Error if the relay UPOKT reward could not be unclaimed.
+func (mnr *miner) unclaimRelayUPOKT(ctx context.Context, relay servicetypes.Relay) either.Either[*relayer.MinedRelay] {
+	if err := mnr.relayMeter.SetNonApplicableRelayReward(ctx, relay.GetReq().GetMeta()); err != nil {
+		return either.Error[*relayer.MinedRelay](err)
+	}
+
+	return either.Success[*relayer.MinedRelay](nil)
+}

pkg/relayer/miner/miner_test.go

@@ -13,12 +13,14 @@ import (
 	"time"
 
 	"cosmossdk.io/depinject"
+	gomock "github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 
 	"github.com/pokt-network/poktroll/pkg/crypto/protocol"
 	"github.com/pokt-network/poktroll/pkg/observable/channel"
 	"github.com/pokt-network/poktroll/pkg/relayer"
 	"github.com/pokt-network/poktroll/pkg/relayer/miner"
+	"github.com/pokt-network/poktroll/testutil/mockrelayer"
 	"github.com/pokt-network/poktroll/testutil/testclient/testqueryclients"
 	servicetypes "github.com/pokt-network/poktroll/x/service/types"
 )
@@ -48,8 +50,9 @@ func TestMiner_MinedRelays(t *testing.T) {
 
 	testqueryclients.SetServiceRelayDifficultyTargetHash(t, testSvcId, testRelayMiningTargetHash)
 	serviceQueryClientMock := testqueryclients.NewTestServiceQueryClient(t)
+	relayMeterMock := newMockRelayMeter(t)
 
-	deps := depinject.Supply(serviceQueryClientMock)
+	deps := depinject.Supply(serviceQueryClientMock, relayMeterMock)
 	mnr, err := miner.NewMiner(deps)
 	require.NoError(t, err)
 
@@ -154,3 +157,17 @@ func unmarshalHexMinedRelay(
 		Hash:  relayHash,
 	}
 }
+
+// newMockRelayMeter returns a mock RelayMeter that is used by the relay miner to claim and unclaim relays.
+func newMockRelayMeter(t *testing.T) relayer.RelayMeter {
+	t.Helper()
+
+	ctrl := gomock.NewController(t)
+	relayMeter := mockrelayer.NewMockRelayMeter(ctrl)
+
+	relayMeter.EXPECT().Start(gomock.Any()).Return(nil).AnyTimes()
+	relayMeter.EXPECT().AccumulateRelayReward(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	relayMeter.EXPECT().SetNonApplicableRelayReward(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+
+	return relayMeter
+}

pkg/relayer/proxy/errors.go

@@ -17,4 +17,7 @@ var (
 	ErrRelayerProxyUnsupportedTransportType       = sdkerrors.Register(codespace, 9, "unsupported proxy transport type")
 	ErrRelayerProxyInternalError                  = sdkerrors.Register(codespace, 10, "internal error")
 	ErrRelayerProxyMissingSupplierOperatorAddress = sdkerrors.Register(codespace, 11, "supplier operator address is missing")
+	ErrRelayerProxyUnknownSession                 = sdkerrors.Register(codespace, 12, "relayer proxy encountered unknown session")
+	ErrRelayerProxyRateLimited                    = sdkerrors.Register(codespace, 13, "offchain rate limit hit by relayer proxy")
+	ErrRelayerProxyUnclaimRelayPrice              = sdkerrors.Register(codespace, 14, "failed to unclaim relay price")
 )

pkg/relayer/proxy/proxy.go

@@ -78,6 +78,11 @@ type relayerProxy struct {
 	// 1. Relay verification to check if the incoming relay matches the supplier hosted by the relay miner;
 	// 2. Relay signing to resolve which keyring key name to use for signing;
 	OperatorAddressToSigningKeyNameMap map[string]string
+
+	// relayMeter keeps track of the total amount of stake an onchhain Application
+	// will owe an onchain Supplier (backed by this RelayMiner) once the session settles.
+	// It also configures application over-servicing allowance.
+	relayMeter relayer.RelayMeter
 }
 
 // NewRelayerProxy creates a new relayer proxy with the given dependencies or returns
@@ -86,9 +91,15 @@ type relayerProxy struct {
 // Required dependencies:
 //   - cosmosclient.Context
 //   - client.BlockClient
+//   - crypto.RingCache
+//   - client.SupplierQueryClient
 //   - client.SessionQueryClient
 //   - client.SharedQueryClient
-//   - client.SupplierQueryClient
+//   - keyring.Keyring
+//   - client.SharedQueryClient
+//   - client.ApplicationQueryClient
+//   - client.ServiceQueryClient
+//   - client.EventsQueryClient
 //
 // Available options:
 //   - WithSigningKeyNames
@@ -108,6 +119,7 @@ func NewRelayerProxy(
 		&rp.sessionQuerier,
 		&rp.sharedQuerier,
 		&rp.keyring,
+		&rp.relayMeter,
 	); err != nil {
 		return nil, err
 	}
@@ -143,6 +155,13 @@ func (rp *relayerProxy) Start(ctx context.Context) error {
 	// Start the ring cache.
 	rp.ringCache.Start(ctx)
 
+	// Start the relay meter by subscribing to the on-chain events.
+	// This function is non-blocking and the subscription cancellation is handled
+	// by the context passed to the Start method.
+	if err := rp.relayMeter.Start(ctx); err != nil {
+		return err
+	}
+
 	startGroup, ctx := errgroup.WithContext(ctx)
 
 	for _, relayServer := range rp.servers {