-
Notifications
You must be signed in to change notification settings - Fork 2k
/
limiter.go
59 lines (49 loc) · 1.81 KB
/
limiter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
// Copyright (c) The Thanos Authors.
// Licensed under the Apache License 2.0.
package store
import (
"sync"
"github.com/pkg/errors"
"github.com/prometheus/client_golang/prometheus"
"go.uber.org/atomic"
)
type ChunksLimiter interface {
// Reserve num chunks out of the total number of chunks enforced by the limiter.
// Returns an error if the limit has been exceeded. This function must be
// goroutine safe.
Reserve(num uint64) error
}
// ChunksLimiterFactory is used to create a new ChunksLimiter. The factory is useful for
// projects depending on Thanos (eg. Cortex) which have dynamic limits.
type ChunksLimiterFactory func(failedCounter prometheus.Counter) ChunksLimiter
// Limiter is a simple mechanism for checking if something has passed a certain threshold.
type Limiter struct {
limit uint64
reserved atomic.Uint64
// Counter metric which we will increase if limit is exceeded.
failedCounter prometheus.Counter
failedOnce sync.Once
}
// NewLimiter returns a new limiter with a specified limit. 0 disables the limit.
func NewLimiter(limit uint64, ctr prometheus.Counter) *Limiter {
return &Limiter{limit: limit, failedCounter: ctr}
}
// Reserve implements ChunksLimiter.
func (l *Limiter) Reserve(num uint64) error {
if l.limit == 0 {
return nil
}
if reserved := l.reserved.Add(num); reserved > l.limit {
// We need to protect from the counter being incremented twice due to concurrency
// while calling Reserve().
l.failedOnce.Do(l.failedCounter.Inc)
return errors.Errorf("limit %v violated (got %v)", l.limit, reserved)
}
return nil
}
// NewChunksLimiterFactory makes a new ChunksLimiterFactory with a static limit.
func NewChunksLimiterFactory(limit uint64) ChunksLimiterFactory {
return func(failedCounter prometheus.Counter) ChunksLimiter {
return NewLimiter(limit, failedCounter)
}
}