waku-org/go-waku

package publish

import (
    "context"
    "errors"

    "github.com/waku-org/go-waku/waku/v2/protocol"
    "go.uber.org/zap"
    "golang.org/x/time/rate"
)

// PublishRateLimiter is used to decorate publish functions to limit the
// number of messages per second that can be published
type PublishRateLimiter struct {
    limiter *rate.Limiter
}

// NewPublishRateLimiter will create a new instance of PublishRateLimiter.
// You can specify an rate.Inf value to in practice ignore the rate limiting
func NewPublishRateLimiter(r rate.Limit, b int) *PublishRateLimiter {
    return &PublishRateLimiter{
        limiter: rate.NewLimiter(r, b),
    }
}

// ThrottlePublishFn is used to decorate a PublishFn so rate limiting is applied
func (p *PublishRateLimiter) ThrottlePublishFn(ctx context.Context, publishFn PublishFn) PublishFn {
    return func(envelope *protocol.Envelope, logger *zap.Logger) error {
        if err := p.Check(ctx, logger); err != nil {
            return err
        }
        return publishFn(envelope, logger)
    }
}

func (p *PublishRateLimiter) Check(ctx context.Context, logger *zap.Logger) error {
    if err := p.limiter.Wait(ctx); err != nil {
        if !errors.Is(err, context.Canceled) {
            logger.Error("could not send message (limiter)", zap.Error(err))
        }
        return err
    }
    return nil
}