mirror of https://github.com/status-im/go-waku.git
38 lines
1.1 KiB
Go
38 lines
1.1 KiB
Go
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.limiter.Wait(ctx); err != nil {
|
|
if !errors.Is(err, context.Canceled) {
|
|
logger.Error("could not send message (limiter)", zap.Error(err))
|
|
}
|
|
return err
|
|
}
|
|
return publishFn(envelope, logger)
|
|
}
|
|
}
|