Advertisement
Guest User

Untitled

a guest
Nov 19th, 2018
90
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Go 1.38 KB | None | 0 0
  1.  
  2.  
  3. func newRatelimiter() ratelimiter {
  4.     rl := ratelimiter{
  5.         buckets: map[string]rlBucket{},
  6.         global:  newRatelimitBucket(120, 60),
  7.     }
  8.     rl.buckets[cmd.UpdateStatus] = newRatelimitBucket(5, 60)
  9.  
  10.     return rl
  11. }
  12.  
  13. type rlEntry struct {
  14.     unix int64
  15.     cmd  string
  16. }
  17.  
  18. func newRatelimitBucket(requests, seconds int) rlBucket {
  19.     return rlBucket{
  20.         entries:  make([]rlEntry, requests),
  21.         duration: (time.Duration(seconds) * time.Second).Nanoseconds(),
  22.     }
  23. }
  24.  
  25. type rlBucket struct {
  26.     entries  []rlEntry
  27.     duration int64
  28. }
  29.  
  30. func (b *rlBucket) Blocked() bool {
  31.     last := b.entries[len(b.entries)-1]
  32.     return time.Now().UnixNano()-last.unix <= b.duration
  33. }
  34.  
  35. func (b *rlBucket) Insert(cmd string) {
  36.     // TODO: we could shift the last valid element to the bottom and then not shift on every insert
  37.     b.entries = append(b.entries[1:], b.entries[:len(b.entries)-2])
  38.     b.entries[0] = rlEntry{
  39.         unix: time.Now().UnixNano(),
  40.         cmd:  cmd,
  41.     }
  42. }
  43.  
  44. type ratelimiter struct {
  45.     sync.RWMutex
  46.     buckets map[string]rlBucket
  47.     global  rlBucket
  48. }
  49.  
  50. func (rl *ratelimiter) Request(command string) (accepted bool) {
  51.     rl.Lock()
  52.     defer rl.Unlock()
  53.  
  54.     // global
  55.     if rl.global.Blocked() {
  56.         return false
  57.     }
  58.     rl.global.Insert(command)
  59.  
  60.     // bucket specific
  61.     if bucket, exists := rl.buckets[command]; exists {
  62.         if bucket.Blocked() {
  63.             return false
  64.         }
  65.         bucket.Insert(command)
  66.     }
  67.  
  68.     return true
  69. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement