ripple-api/app/rate_limiter.go

103 lines
1.8 KiB
Go
Raw Normal View History

package app
import (
"fmt"
"strconv"
"sync"
"time"
)
const reqsPerSecond = 5000
const sleepTime = time.Second / reqsPerSecond
var limiter = make(chan struct{}, reqsPerSecond)
func setUpLimiter() {
for i := 0; i < reqsPerSecond; i++ {
limiter <- struct{}{}
}
go func() {
for {
limiter <- struct{}{}
time.Sleep(sleepTime)
}
}()
}
func rateLimiter() {
<-limiter
}
func perUserRequestLimiter(uid int, ip string) {
if uid == 0 {
defaultLimiter.Request("ip:"+ip, 60)
} else {
defaultLimiter.Request("user:"+strconv.Itoa(uid), 2000)
}
}
var defaultLimiter = &specificRateLimiter{
Map: make(map[string]chan struct{}),
Mutex: &sync.RWMutex{},
}
type specificRateLimiter struct {
Map map[string]chan struct{}
Mutex *sync.RWMutex
}
func (s *specificRateLimiter) Request(u string, perMinute int) {
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)
}
}
<-c
}
func (s *specificRateLimiter) 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
}
}
}
func makePrefilledChan(l int) chan struct{} {
c := make(chan struct{}, l)
for i := 0; i < l; i++ {
c <- struct{}{}
}
return c
}