bench-forgejo/vendor/gopkg.in/redis.v2/rate_limit.go

54 lines
864 B
Go
Raw Normal View History

2016-11-04 03:46:01 +05:30
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
}