54 lines
864 B
Go
54 lines
864 B
Go
|
package redis
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync/atomic"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type rateLimiter struct {
|
||
|
|
v int64
|
||
|
|
|
||
|
|
_closed int64
|
||
|
|
}
|
||
|
|
|
||
|
|
func newRateLimiter(limit time.Duration, bucketSize int) *rateLimiter {
|
||
|
|
rl := &rateLimiter{
|
||
|
|
v: int64(bucketSize),
|
||
|
|
}
|
||
|
|
go rl.loop(limit, int64(bucketSize))
|
||
|
|
return rl
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *rateLimiter) loop(limit time.Duration, bucketSize int64) {
|
||
|
|
for {
|
||
|
|
if rl.closed() {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
if v := atomic.LoadInt64(&rl.v); v < bucketSize {
|
||
|
|
atomic.AddInt64(&rl.v, 1)
|
||
|
|
}
|
||
|
|
time.Sleep(limit)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *rateLimiter) Check() bool {
|
||
|
|
for {
|
||
|
|
if v := atomic.LoadInt64(&rl.v); v > 0 {
|
||
|
|
if atomic.CompareAndSwapInt64(&rl.v, v, v-1) {
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *rateLimiter) Close() error {
|
||
|
|
atomic.StoreInt64(&rl._closed, 1)
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *rateLimiter) closed() bool {
|
||
|
|
return atomic.LoadInt64(&rl._closed) == 1
|
||
|
|
}
|