123 lines
2.7 KiB
Go
123 lines
2.7 KiB
Go
|
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
|
||
|
|
}
|