synapsecns/sanguine

package client

import (
    "context"
    "fmt"
    "github.com/prometheus/client_golang/prometheus"
    "golang.org/x/sync/semaphore"
    "golang.org/x/time/rate"
    "math/big"
    "sync/atomic"
    "time"
)

// MeteredEVMClient is a metered client that exposes a counter metric for evm requests sent through the client.
//
//go:generate go run github.com/vektra/mockery/v2 --name MeteredEVMClient --output ./mocks --case=underscore
type MeteredEVMClient interface {
    EVMClient
    // RequestCount gets the request count from the metered evm client
    RequestCount() int64
    // ConcurrencyCount gets the request concurrency on the meterd client
    ConcurrencyCount() int32
    // ClientID is a unique identifier for the client.
    //
    // note: this is not guaranteed to be unique - it's set by the caller.
    ClientID() string
    // AttemptReconnect attempts to reconnect
    // TODO: replace with https://github.com/ethereum/go-ethereum/issues/22266
    AttemptReconnect() bool
}

// meteredEVMClientImpl is an instrumented client with a rate limiter
// it wraps EVMClient and takes the keepRateLimiter as a config
// we can't use the original keep rate limiter since it does not respect
// the context analytics.
type meteredEVMClientImpl struct {
    LifecycleClient
    // semaphore is used to manage concurrency limits. If
    // concurrency limiting is turned off, this is not used
    semaphore *semaphore.Weighted
    // limiter is used to limit the number of requests
    limiter *rate.Limiter
    // acquirePermitTimeout is the max amount of time a metered client
    // will wait to acquire a semaphore and make a request. This will default to 5 minutes
    acquirePermitTimeout time.Duration
    // counter is the counter for total number of requests since client creation
    counter uint64
    // concurrency is the number of ongoing requests at any given time
    concurrency uint32
    // chainID is the chain id used for the client
    chainID *big.Int
    // clientID is a unique identifier used for metrics across a pool.
    // this can be the wsurl in many cases (if metrics are private/authentication on rpc server)
    // is properly secured without a token based url
    clientID string
}

// AttemptReconnect attempts to reconnect
// TODO implement.
func (m meteredEVMClientImpl) AttemptReconnect() bool {
    return true
}

const requestTimeout = time.Minute

// NewMeteredClient wraps an evm client in a keepRate limiter and creates
// a metric handler with some standard metrics. It also implements a ChainConfig()
// method to get the chainconfig for a given chain by id. This will return nil if no chain config is found.
func NewMeteredClient(client EVMClient, chainID *big.Int, clientID string, config *LimiterConfig) MeteredEVMClient {
    meteredClient := getMeteredClientStub(chainID, clientID, config)
    meteredClient.LifecycleClient = NewLifecycleClient(client, chainID, meteredClient, requestTimeout)
    return &meteredClient
}

// getMeteredClientStub stub gets the metered client without setting the lifecylce using the config.
func getMeteredClientStub(chainID *big.Int, clientID string, config *LimiterConfig) meteredEVMClientImpl {
    meteredClient := meteredEVMClientImpl{
        // counter is an atomicically incrementing int used to instrument eth client
        counter:  0,
        chainID:  chainID,
        clientID: clientID,
    }

    if config == nil {
        config = &LimiterConfig{}
    }

    // setup the rate limiter
    if config.RequestsPerSecondLimit > 0 {
        meteredClient.limiter = rate.NewLimiter(
            rate.Limit(config.RequestsPerSecondLimit),
            1,
        )
    }

    if config.ConcurrencyLimit > 0 {
        meteredClient.semaphore = semaphore.NewWeighted(
            int64(config.ConcurrencyLimit),
        )
    }

    if config.AcquirePermitTimeout > 0 {
        meteredClient.acquirePermitTimeout = config.AcquirePermitTimeout
    } else {
        meteredClient.acquirePermitTimeout = 5 * time.Minute
    }

    return meteredClient
}

// ClientID gets the unique client identifier.
func (m meteredEVMClientImpl) ClientID() string {
    return m.clientID
}

// AcquirePermit attempts to acquire a permit from the keepRate limiter.
// (this needs to be called before the request).
func (m meteredEVMClientImpl) AcquirePermit(ctx context.Context) (err error) {
    ctx, cancel := context.WithTimeout(ctx, m.acquirePermitTimeout)
    defer cancel()

    if m.limiter != nil {
        err = m.limiter.Wait(ctx)
        if err != nil {
            return fmt.Errorf("cannot wait for limiter: %w", err)
        }
    }

    if m.semaphore != nil {
        err = m.semaphore.Acquire(ctx, 1)
        if err != nil {
            return fmt.Errorf("could not acquire semaphore: %w", err)
        }
    }

    atomic.AddUint64(&m.counter, 1)
    atomic.AddUint32(&m.concurrency, 1)
    return nil
}

// ReleasePermit releases a permit (this needs to be called after the request).
func (m meteredEVMClientImpl) ReleasePermit() {
    if m.semaphore != nil {
        m.semaphore.Release(1)
    }
    // decrement the concurrency counter (see the atomic.AddUint32())
    atomic.AddUint32(&m.concurrency, ^uint32(0))
}

// RequestCount gets the request count.
func (m meteredEVMClientImpl) RequestCount() int64 {
    return int64(atomic.LoadUint64(&m.counter))
}

// ConcurrencyCount gets the number of requests in progress.
func (m meteredEVMClientImpl) ConcurrencyCount() int32 {
    return int32(atomic.LoadUint32(&m.concurrency))
}

// RequestCountMetricName is the name of the metric which counts requests.
const RequestCountMetricName = "request_count_total"

// ConcurrencyGaugeMetricName is the name of the metric which gauges request concurrency.
const ConcurrencyGaugeMetricName = "request_concurrency"

// GetMetrics gets metrics associated with the network provider.
func (m meteredEVMClientImpl) GetMetrics(labels map[string]string) []prometheus.Collector {
    requestCount := prometheus.NewCounterFunc(prometheus.CounterOpts{
        Name: RequestCountMetricName,
        Help: "the number of requests sent by the client",
    }, func() float64 {
        return float64(m.RequestCount())
    })

    concurrencyCount := prometheus.NewGaugeFunc(prometheus.GaugeOpts{
        Name: ConcurrencyGaugeMetricName,
        Help: "the number of requests sent by the client concurrently",
    }, func() float64 {
        return float64(m.ConcurrencyCount())
    })

    return []prometheus.Collector{requestCount, concurrencyCount}
}

var _ MeteredEVMClient = &meteredEVMClientImpl{}

// LimiterConfig represents the configuration of the rate limiter.
// copied from https://github.com/keep-network/keep-common/blob/v1.7.0/pkg/rate/limiter.go#L19
type LimiterConfig struct {
    // RequestsPerSecondLimit sets the maximum average number of requests
    // per second. It's important to note that in short periods of time
    // the actual average may exceed this limit slightly.
    RequestsPerSecondLimit int

    // ConcurrencyLimit sets the maximum number of concurrent requests which
    // can be executed against the target at the same time.
    ConcurrencyLimit int

    // AcquirePermitTimeout determines how long a request can wait trying
    // to acquire a permit from the rate limiter.
    AcquirePermitTimeout time.Duration
}