175 lines
2.9 KiB
Go
175 lines
2.9 KiB
Go
package http
|
|
|
|
import (
|
|
"errors"
|
|
"log"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type RateLimiter interface {
|
|
Acquire() (*Token, error)
|
|
Release(*Token)
|
|
LimitExceededPerc() float32
|
|
ReleaseAll()
|
|
}
|
|
|
|
type Token struct {
|
|
ID string
|
|
CreatedAt time.Time
|
|
}
|
|
|
|
func NewMaxConcurrencyRateLimiter(limit int) (RateLimiter, error) {
|
|
if limit <= 0 {
|
|
return nil, errors.New("invalid rate limit")
|
|
}
|
|
|
|
m := newManager(limit)
|
|
|
|
// Await.
|
|
func() {
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-m.inChan:
|
|
m.tryGenerateToken()
|
|
case t := <-m.releaseChan:
|
|
m.releaseToken(t)
|
|
}
|
|
}
|
|
}()
|
|
}()
|
|
return m, nil
|
|
}
|
|
|
|
func newManager(limit int) *Manager {
|
|
return &Manager{
|
|
errorChan: make(chan error),
|
|
outChan: make(chan *Token),
|
|
inChan: make(chan struct{}),
|
|
activeTokens: make(map[string]*Token),
|
|
releaseChan: make(chan *Token),
|
|
needToken: 0,
|
|
limit: limit,
|
|
makeToken: func() *Token {
|
|
return &Token{
|
|
ID: uuid.New().String(),
|
|
CreatedAt: time.Now().UTC(),
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
type Manager struct {
|
|
errorChan chan error
|
|
releaseChan chan *Token
|
|
outChan chan *Token
|
|
inChan chan struct{}
|
|
needToken int64
|
|
activeTokens map[string]*Token
|
|
tokensMutex sync.RWMutex
|
|
limit int
|
|
makeToken func() *Token
|
|
}
|
|
|
|
func (m *Manager) Acquire() (*Token, error) {
|
|
go func() {
|
|
m.inChan <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case t := <-m.outChan:
|
|
return t, nil
|
|
case err := <-m.errorChan:
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
func (m *Manager) Release(t *Token) {
|
|
go func() {
|
|
m.releaseChan <- t
|
|
}()
|
|
}
|
|
|
|
func (m *Manager) ReleaseAll() {
|
|
m.tokensMutex.Lock()
|
|
for _, t := range m.activeTokens {
|
|
m.Release(t)
|
|
}
|
|
m.activeTokens = make(map[string]*Token)
|
|
m.tokensMutex.Unlock()
|
|
}
|
|
|
|
func (m *Manager) LimitExceededPerc() float32 {
|
|
return m.limitExceededPerc()
|
|
}
|
|
|
|
func (m *Manager) incNeedToken() {
|
|
atomic.AddInt64(&m.needToken, 1)
|
|
}
|
|
|
|
func (m *Manager) decNeedToken() {
|
|
atomic.AddInt64(&m.needToken, -1)
|
|
}
|
|
|
|
func (m *Manager) awaitingToken() bool {
|
|
return atomic.LoadInt64(&m.needToken) > 0
|
|
}
|
|
|
|
func (m *Manager) limitExceededPerc() float32 {
|
|
m.tokensMutex.RLock()
|
|
defer m.tokensMutex.RUnlock()
|
|
return float32(atomic.LoadInt64(&m.needToken)) / float32(m.limit)
|
|
}
|
|
|
|
func (m *Manager) isLimitExceeded() bool {
|
|
m.tokensMutex.RLock()
|
|
defer m.tokensMutex.RUnlock()
|
|
|
|
if len(m.activeTokens) >= m.limit {
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (m *Manager) releaseToken(t *Token) {
|
|
if t == nil {
|
|
log.Println("cannot release a nil token")
|
|
return
|
|
}
|
|
|
|
m.tokensMutex.Lock()
|
|
if _, ok := m.activeTokens[t.ID]; !ok {
|
|
log.Printf("unable to release token %s - not in use", t)
|
|
return
|
|
}
|
|
|
|
delete(m.activeTokens, t.ID)
|
|
m.tokensMutex.Unlock()
|
|
|
|
if m.awaitingToken() {
|
|
m.decNeedToken()
|
|
go m.tryGenerateToken()
|
|
}
|
|
}
|
|
|
|
func (m *Manager) tryGenerateToken() {
|
|
if m.isLimitExceeded() {
|
|
m.incNeedToken()
|
|
return
|
|
}
|
|
|
|
t := m.makeToken()
|
|
|
|
m.tokensMutex.Lock()
|
|
m.activeTokens[t.ID] = t
|
|
m.tokensMutex.Unlock()
|
|
|
|
go func() {
|
|
m.outChan <- t
|
|
}()
|
|
}
|