github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/alertmanager/rate_limited_notifier.go (about)

     1  package alertmanager
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/prometheus/alertmanager/notify"
     9  	"github.com/prometheus/alertmanager/types"
    10  	"github.com/prometheus/client_golang/prometheus"
    11  	"go.uber.org/atomic"
    12  	"golang.org/x/time/rate"
    13  )
    14  
    15  type rateLimits interface {
    16  	RateLimit() rate.Limit
    17  	Burst() int
    18  }
    19  
    20  type rateLimitedNotifier struct {
    21  	upstream notify.Notifier
    22  	counter  prometheus.Counter
    23  
    24  	limiter *rate.Limiter
    25  	limits  rateLimits
    26  
    27  	recheckInterval time.Duration
    28  	recheckAt       atomic.Int64 // unix nanoseconds timestamp
    29  }
    30  
    31  func newRateLimitedNotifier(upstream notify.Notifier, limits rateLimits, recheckInterval time.Duration, counter prometheus.Counter) *rateLimitedNotifier {
    32  	return &rateLimitedNotifier{
    33  		upstream:        upstream,
    34  		counter:         counter,
    35  		limits:          limits,
    36  		limiter:         rate.NewLimiter(limits.RateLimit(), limits.Burst()),
    37  		recheckInterval: recheckInterval,
    38  	}
    39  }
    40  
    41  var errRateLimited = errors.New("failed to notify due to rate limits")
    42  
    43  func (r *rateLimitedNotifier) Notify(ctx context.Context, alerts ...*types.Alert) (bool, error) {
    44  	now := time.Now()
    45  	if now.UnixNano() >= r.recheckAt.Load() {
    46  		if limit := r.limits.RateLimit(); r.limiter.Limit() != limit {
    47  			r.limiter.SetLimitAt(now, limit)
    48  		}
    49  
    50  		if burst := r.limits.Burst(); r.limiter.Burst() != burst {
    51  			r.limiter.SetBurstAt(now, burst)
    52  		}
    53  
    54  		r.recheckAt.Store(now.UnixNano() + r.recheckInterval.Nanoseconds())
    55  	}
    56  
    57  	// This counts as single notification, no matter how many alerts there are in it.
    58  	if !r.limiter.AllowN(now, 1) {
    59  		r.counter.Inc()
    60  		// Don't retry this notification later.
    61  		return false, errRateLimited
    62  	}
    63  
    64  	return r.upstream.Notify(ctx, alerts...)
    65  }