File: ratelimit.go

package info (click to toggle)
relic 7.6.1-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 3,108 kB
  • sloc: sh: 230; makefile: 10
file content (78 lines) | stat: -rw-r--r-- 1,936 bytes parent folder | download
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package tokencache

import (
	"context"
	"crypto"
	"io"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/sassoftware/relic/v7/token"
	"golang.org/x/time/rate"
)

var metricRateLimited = promauto.NewCounter(prometheus.CounterOpts{
	Name: "token_operation_limited_seconds",
	Help: "Cumulative number of seconds waiting for rate limits",
})

type RateLimited struct {
	token.Token
	limit *rate.Limiter
}

func NewLimiter(base token.Token, limit float64, burst int) *RateLimited {
	if burst < 1 {
		burst = 1
	}
	return &RateLimited{
		Token: base,
		limit: rate.NewLimiter(rate.Limit(limit), burst),
	}
}

type rateLimitedKey struct {
	token.Key
	limit *rate.Limiter
}

func (r *RateLimited) GetKey(ctx context.Context, keyName string) (token.Key, error) {
	start := time.Now()
	if err := r.limit.Wait(ctx); err != nil {
		return nil, err
	}
	if waited := time.Since(start); waited > 1*time.Millisecond {
		metricRateLimited.Add(time.Since(start).Seconds())
	}
	key, err := r.Token.GetKey(ctx, keyName)
	if err != nil {
		return nil, err
	}
	return &rateLimitedKey{
		Key:   key,
		limit: r.limit,
	}, nil
}

func (k *rateLimitedKey) Sign(rand io.Reader, digest []byte, opts crypto.SignerOpts) (sig []byte, err error) {
	start := time.Now()
	if err := k.limit.Wait(context.Background()); err != nil {
		return nil, err
	}
	if waited := time.Since(start); waited > 1*time.Millisecond {
		metricRateLimited.Add(time.Since(start).Seconds())
	}
	return k.Key.Sign(rand, digest, opts)
}

func (k *rateLimitedKey) SignContext(ctx context.Context, digest []byte, opts crypto.SignerOpts) (sig []byte, err error) {
	start := time.Now()
	if err := k.limit.Wait(ctx); err != nil {
		return nil, err
	}
	if waited := time.Since(start); waited > 1*time.Millisecond {
		metricRateLimited.Add(time.Since(start).Seconds())
	}
	return k.Key.SignContext(ctx, digest, opts)
}