forked from RealistikOsu/RealistikAPI
-
Notifications
You must be signed in to change notification settings - Fork 0
/
limit.go
122 lines (109 loc) · 2.67 KB
/
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
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package limit
import (
"fmt"
"sync"
"time"
)
// Request is a Request with DefaultLimiter.
func Request(u string, perMinute int) { DefaultLimiter.Request(u, perMinute) }
// NonBlockingRequest is a NonBlockingRequest with DefaultLimiter.
func NonBlockingRequest(u string, perMinute int) bool {
return DefaultLimiter.NonBlockingRequest(u, perMinute)
}
// DefaultLimiter is the RateLimiter used by the package-level
// Request and NonBlockingRequest.
var DefaultLimiter = &RateLimiter{
Map: make(map[string]chan struct{}),
Mutex: &sync.RWMutex{},
}
// RateLimiter is a simple rate limiter.
type RateLimiter struct {
Map map[string]chan struct{}
Mutex *sync.RWMutex
}
// Request is a simple request. Blocks until it can make the request.
func (s *RateLimiter) Request(u string, perMinute int) {
s.request(u, perMinute, true)
}
// NonBlockingRequest checks if it can do a request. If it can't, it returns
// false, else it returns true if the request succeded.
func (s *RateLimiter) NonBlockingRequest(u string, perMinute int) bool {
return s.request(u, perMinute, false)
}
func (s *RateLimiter) request(u string, perMinute int, blocking bool) bool {
s.check()
s.Mutex.RLock()
c, exists := s.Map[u]
s.Mutex.RUnlock()
if !exists {
c = makePrefilledChan(perMinute)
s.Mutex.Lock()
// Now that we have exclusive read and write-access, we want to
// make sure we don't overwrite an existing channel. Otherwise,
// race conditions and panic happen.
if cNew, exists := s.Map[u]; exists {
c = cNew
s.Mutex.Unlock()
} else {
s.Map[u] = c
s.Mutex.Unlock()
<-c
go s.filler(u, perMinute)
}
}
return rcv(c, blocking)
}
// rcv receives from a channel, but if blocking is true it waits til something
// is received and always returns true, otherwise if it can't receive it
// returns false.
func rcv(c chan struct{}, blocking bool) bool {
if blocking {
<-c
return true
}
select {
case <-c:
return true
default:
return false
}
}
func (s *RateLimiter) filler(el string, perMinute int) {
defer func() {
r := recover()
if r != nil {
fmt.Println(r)
}
}()
s.Mutex.RLock()
c := s.Map[el]
s.Mutex.RUnlock()
for {
select {
case c <- struct{}{}:
time.Sleep(time.Minute / time.Duration(perMinute))
default: // c is full
s.Mutex.Lock()
close(c)
delete(s.Map, el)
s.Mutex.Unlock()
return
}
}
}
// check makes sure the map and the mutex are properly initialised.
func (s *RateLimiter) check() {
if s.Map == nil {
s.Map = make(map[string]chan struct{})
}
if s.Mutex == nil {
s.Mutex = new(sync.RWMutex)
}
}
func makePrefilledChan(l int) chan struct{} {
c := make(chan struct{}, l)
for i := 0; i < l; i++ {
c <- struct{}{}
}
return c
}