vitess.io/vitess@v0.16.2/go/ratelimiter/ratelimiter.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package ratelimiter implements rate limiting functionality.
    18  package ratelimiter
    19  
    20  import (
    21  	"sync"
    22  	"time"
    23  )
    24  
    25  // RateLimiter was inspired by https://github.com/golang/go/wiki/RateLimiting.
    26  // However, the go example is not good for setting high qps limits because
    27  // it will cause the ticker to fire too often. Also, the ticker will continue
    28  // to fire when the system is idle. This new Ratelimiter achieves the same thing,
    29  // but by using just counters with no tickers or channels.
    30  type RateLimiter struct {
    31  	maxCount int
    32  	interval time.Duration
    33  
    34  	mu       sync.Mutex
    35  	curCount int
    36  	lastTime time.Time
    37  }
    38  
    39  // NewRateLimiter creates a new RateLimiter. maxCount is the max burst allowed
    40  // while interval specifies the duration for a burst. The effective rate limit is
    41  // equal to maxCount/interval. For example, if you want to a max QPS of 5000,
    42  // and want to limit bursts to no more than 500, you'd specify a maxCount of 500
    43  // and an interval of 100*time.Millilsecond.
    44  func NewRateLimiter(maxCount int, interval time.Duration) *RateLimiter {
    45  	return &RateLimiter{
    46  		maxCount: maxCount,
    47  		interval: interval,
    48  	}
    49  }
    50  
    51  // Allow returns true if a request is within the rate limit norms.
    52  // Otherwise, it returns false.
    53  func (rl *RateLimiter) Allow() bool {
    54  	rl.mu.Lock()
    55  	defer rl.mu.Unlock()
    56  	if time.Since(rl.lastTime) < rl.interval {
    57  		if rl.curCount > 0 {
    58  			rl.curCount--
    59  			return true
    60  		}
    61  		return false
    62  	}
    63  	rl.curCount = rl.maxCount - 1
    64  	rl.lastTime = time.Now()
    65  	return true
    66  }