mirror of
https://tildegit.org/solderpunk/molly-brown
synced 2024-05-10 02:26:04 +02:00
53 lines
834 B
Go
53 lines
834 B
Go
package main
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type RateLimiter struct {
|
|
mu sync.Mutex
|
|
bucket map[string]int
|
|
rate int
|
|
burst int
|
|
}
|
|
|
|
func newRateLimiter(rate int, burst int) RateLimiter {
|
|
var rl = new(RateLimiter)
|
|
rl.bucket = make(map[string]int)
|
|
rl.rate = rate
|
|
rl.burst = burst
|
|
// Leak periodically
|
|
go func () {
|
|
for(true) {
|
|
rl.mu.Lock()
|
|
for addr, drips := range rl.bucket {
|
|
if drips <= rate {
|
|
delete(rl.bucket, addr)
|
|
} else {
|
|
rl.bucket[addr] = drips - rl.rate
|
|
}
|
|
}
|
|
rl.mu.Unlock()
|
|
time.Sleep(time.Second)
|
|
}
|
|
}()
|
|
return *rl
|
|
}
|
|
|
|
func (rl *RateLimiter) Allowed(addr string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
drips, present := rl.bucket[addr]
|
|
if !present {
|
|
rl.bucket[addr] = 1
|
|
return true
|
|
}
|
|
if drips == rl.burst {
|
|
return false
|
|
}
|
|
rl.bucket[addr] = drips + 1
|
|
return true
|
|
}
|
|
|