rate_limit.go 864 B

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. package redis
  2. import (
  3. "sync/atomic"
  4. "time"
  5. )
  6. type rateLimiter struct {
  7. v int64
  8. _closed int64
  9. }
  10. func newRateLimiter(limit time.Duration, bucketSize int) *rateLimiter {
  11. rl := &rateLimiter{
  12. v: int64(bucketSize),
  13. }
  14. go rl.loop(limit, int64(bucketSize))
  15. return rl
  16. }
  17. func (rl *rateLimiter) loop(limit time.Duration, bucketSize int64) {
  18. for {
  19. if rl.closed() {
  20. break
  21. }
  22. if v := atomic.LoadInt64(&rl.v); v < bucketSize {
  23. atomic.AddInt64(&rl.v, 1)
  24. }
  25. time.Sleep(limit)
  26. }
  27. }
  28. func (rl *rateLimiter) Check() bool {
  29. for {
  30. if v := atomic.LoadInt64(&rl.v); v > 0 {
  31. if atomic.CompareAndSwapInt64(&rl.v, v, v-1) {
  32. return true
  33. }
  34. } else {
  35. return false
  36. }
  37. }
  38. }
  39. func (rl *rateLimiter) Close() error {
  40. atomic.StoreInt64(&rl._closed, 1)
  41. return nil
  42. }
  43. func (rl *rateLimiter) closed() bool {
  44. return atomic.LoadInt64(&rl._closed) == 1
  45. }