-
-
Notifications
You must be signed in to change notification settings - Fork 49
/
limiter.go
64 lines (51 loc) · 1.23 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
60
61
62
63
64
package bernard
import (
"context"
"sync"
"time"
"golang.org/x/sync/semaphore"
"golang.org/x/time/rate"
)
const (
// how many requests can be sent per second by all drives using the same account file
requestLimit = 8
// how many drives can run at once (at the trigger level), e.g. 2 triggers, with 5 drives each.
syncLimit = 5
)
type rateLimiter struct {
ctx context.Context
rl *rate.Limiter
sem *semaphore.Weighted
}
func (r *rateLimiter) Wait() {
_ = r.rl.Wait(r.ctx)
}
func (r *rateLimiter) Acquire(n int64) error {
return r.sem.Acquire(r.ctx, n)
}
func (r *rateLimiter) Release(n int64) {
r.sem.Release(n)
}
func newRateLimiter() *rateLimiter {
return &rateLimiter{
ctx: context.Background(),
rl: rate.NewLimiter(rate.Every(time.Second/time.Duration(requestLimit)), requestLimit),
sem: semaphore.NewWeighted(int64(syncLimit)),
}
}
var (
limiters = make(map[string]*rateLimiter)
lock = &sync.Mutex{}
)
func getRateLimiter(account string) (*rateLimiter, error) {
lock.Lock()
defer lock.Unlock()
// return existing limiter for the account
if limiter, ok := limiters[account]; ok {
return limiter, nil
}
// add limiter to map
limiter := newRateLimiter()
limiters[account] = limiter
return limiter, nil
}