-
Notifications
You must be signed in to change notification settings - Fork 7
/
Copy pathapi_rate_limiter.go
60 lines (50 loc) · 1.33 KB
/
api_rate_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
// Copyright 2019 Clivern. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
package hippo
import (
"golang.org/x/time/rate"
"sync"
"time"
)
// caller struct
type caller struct {
limiter *rate.Limiter
lastSeen time.Time
}
// callers list
var callers = make(map[string]*caller)
// mtx mutex
var mtx sync.Mutex
// NewCallerLimiter create a new rate limiter with an identifier
func NewCallerLimiter(identifier string, eventsRate rate.Limit, tokenBurst int) *rate.Limiter {
mtx.Lock()
v, exists := callers[identifier]
if !exists {
mtx.Unlock()
return addCaller(identifier, eventsRate, tokenBurst)
}
// Update the last seen time for the caller.
v.lastSeen = time.Now()
mtx.Unlock()
return v.limiter
}
// addCaller add a caller
func addCaller(identifier string, eventsRate rate.Limit, tokenBurst int) *rate.Limiter {
limiter := rate.NewLimiter(eventsRate, tokenBurst)
mtx.Lock()
// Include the current time when creating a new caller.
callers[identifier] = &caller{limiter, time.Now()}
mtx.Unlock()
return limiter
}
// CleanupCallers cleans old clients
func CleanupCallers(cleanAfter time.Duration) {
mtx.Lock()
for identifier, v := range callers {
if time.Now().Sub(v.lastSeen) > cleanAfter*time.Second {
delete(callers, identifier)
}
}
mtx.Unlock()
}