-
Notifications
You must be signed in to change notification settings - Fork 2.8k
/
rate_limit.go
53 lines (46 loc) · 1.55 KB
/
rate_limit.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
// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium
package helpers
import (
"context"
"time"
"golang.org/x/time/rate"
)
// APILimiter allows to rate limit API calls
type APILimiter struct {
metrics MetricsAPI
limiter *rate.Limiter
}
// MetricsAPI represents the metrics maintained by the API limiter
type MetricsAPI interface {
ObserveRateLimit(operation string, duration time.Duration)
}
// NewAPILimiter returns a new API limiter with the specific rate limit and
// burst configuration. The MetricsAPI interface is called to allow for metrics
// accounting.
func NewAPILimiter(metrics MetricsAPI, rateLimit float64, burst int) *APILimiter {
return &APILimiter{
metrics: metrics,
limiter: rate.NewLimiter(rate.Limit(rateLimit), burst),
}
}
// Limit applies the rate limiting configuration for the given operation
func (l *APILimiter) Limit(ctx context.Context, operation string) {
r := l.limiter.Reserve()
if delay := r.Delay(); delay != time.Duration(0) && delay != rate.InfDuration {
l.metrics.ObserveRateLimit(operation, delay)
// Wait for the required time. We cannot call r.limiter.Wait here, as it
// would request a second reservation, effectively doubling the wait time.
// Instead, the following logic is similar to what r.limiter.Wait(ctx)
// does internally after it successfully obtained/ a reservation.
t := time.NewTimer(delay)
defer t.Stop()
select {
case <-t.C:
// proceed with the operation
case <-ctx.Done():
// cancel the reservation to allow other operations to go through
r.Cancel()
}
}
}