-
Notifications
You must be signed in to change notification settings - Fork 2
/
limiter.go
45 lines (35 loc) · 910 Bytes
/
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
package limiter
import (
"context"
"errors"
ge "github.com/alexfalkowski/go-service/errors"
"github.com/alexfalkowski/go-service/meta"
"github.com/ulule/limiter/v3"
"github.com/ulule/limiter/v3/drivers/store/memory"
)
var (
// ErrMissingKey for limiter.
ErrMissingKey = errors.New("missing key")
keys = map[string]KeyFunc{}
)
// RegisterKey with name and fn. Last register wins.
func RegisterKey(name string, fn KeyFunc) {
keys[name] = fn
}
// New limiter.
func New(cfg *Config) (*limiter.Limiter, KeyFunc, error) {
if !IsEnabled(cfg) {
return nil, nil, nil
}
k, ok := keys[cfg.Kind]
if !ok {
return nil, nil, ErrMissingKey
}
rate, err := limiter.NewRateFromFormatted(cfg.Pattern)
if err != nil {
return nil, nil, ge.Prefix("new limiter", err)
}
return limiter.New(memory.NewStore(), rate), k, nil
}
// KeyFunc for the limiter.
type KeyFunc func(context.Context) meta.Valuer