ratelimit.go raw

   1  package main
   2  
   3  import "time"
   4  
   5  // RateLimiter is a token bucket that allows a fixed number of actions per minute.
   6  type RateLimiter struct {
   7  	tokens chan struct{}
   8  	ticker *time.Ticker
   9  }
  10  
  11  // NewRateLimiter creates a rate limiter allowing perMinute actions per minute.
  12  func NewRateLimiter(perMinute int) *RateLimiter {
  13  	if perMinute < 1 {
  14  		perMinute = 1
  15  	}
  16  	interval := time.Minute / time.Duration(perMinute)
  17  	rl := &RateLimiter{
  18  		tokens: make(chan struct{}, perMinute),
  19  		ticker: time.NewTicker(interval),
  20  	}
  21  	// Start with one token available.
  22  	rl.tokens <- struct{}{}
  23  
  24  	go func() {
  25  		for range rl.ticker.C {
  26  			select {
  27  			case rl.tokens <- struct{}{}:
  28  			default: // bucket full
  29  			}
  30  		}
  31  	}()
  32  	return rl
  33  }
  34  
  35  // Allow returns true if an action is permitted, consuming one token.
  36  func (rl *RateLimiter) Allow() bool {
  37  	select {
  38  	case <-rl.tokens:
  39  		return true
  40  	default:
  41  		return false
  42  	}
  43  }
  44  
  45  // Stop releases the ticker.
  46  func (rl *RateLimiter) Stop() {
  47  	rl.ticker.Stop()
  48  }
  49