This article picks up where the go-Zero left off (1).

In this article, we introduce another tokenLimit, the token bucket limit.

use

const (
	burst   = 100
	rate    = 100
	seconds = 5
)

store := redis.NewRedis("localhost:6379"."node"."")
fmt.Println(store.Ping())
// New tokenLimiter
limiter := limit.NewTokenLimiter(rate, burst, store, "rate-test")
timer := time.NewTimer(time.Second * seconds)
quit := make(chan struct{})
defer timer.Stop()
go func(a) {
  <-timer.C
  close(quit)
}()

var allowed, denied int32
var wait sync.WaitGroup
for i := 0; i < runtime.NumCPU(); i++ {
  wait.Add(1)
  go func(a) {
    for {
      select {
        case <-quit:
          wait.Done()
          return
        default:
          if limiter.Allow() {
            atomic.AddInt32(&allowed, 1)}else {
            atomic.AddInt32(&denied, 1)
          }
      }
    }
  }()
}

wait.Wait()
fmt.Printf("allowed: %d, denied: %d, qps: %d\n", allowed, denied, (allowed+denied)/seconds)
Copy the code

tokenlimit

The overall token bucket production logic is as follows:

  • If the average sending rate is set to R, a token is added to the bucket every 1/r second.
  • Assume that a bucket can hold up to B tokens. If the bucket is full when the token arrives, the token is discarded;
  • When the traffic enters at rate V, the bucket obtains a token at rate V, and the traffic with the token passes. If the traffic with no token fails to pass, the circuit breaker logic is executed.

Go-zero adopts the mode of Lua Script under both types of current limiters. Distributed current limiting can be achieved by relying on Redis, and Lua Script can also achieve atomization of token production and reading operations.

Let’s take a look at a few key properties controlled by Lua Script:

argument mean
ARGV[1] Rate “Generate several tokens per second”
ARGV[2] Burst “Token bucket Max”
ARGV[3] Now_time “Current timestamp”
ARGV[4] Get Token nums Number of tokens developers need to obtain
KEYS[1] Represents the tokenkey of the resource
KEYS[2] Key that represents the refresh time 
Returns whether the expected token can be obtained alive

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

-- fill_time: how long it takes to fill token_bucket
local fill_time = capacity/rate
-- Round down the fill time
local ttl = math.floor(fill_time*2)

-- Gets the number of tokens remaining in the token_bucket
-- Set token_bucket to the maximum number of token buckets if it is the first time
local last_tokens = tonumber(redis.call("get", KEYS[1]))
if last_tokens == nil then
    last_tokens = capacity
end

-- The last time token_bucket was updated
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
if last_refreshed == nil then
    last_refreshed = 0
end

local delta = math.max(0, now-last_refreshed)
-- Calculate the number of new tokens based on the span between the current time and the last update time and the rate of token production
If max_burst is exceeded, excess tokens are discarded
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
if allowed then
    new_tokens = filled_tokens - requested
end

-- Update the new token number and update time
redis.call("setex", KEYS[1], ttl, new_tokens)
redis.call("setex", KEYS[2], ttl, now)

return allowed
Copy the code

As can be seen from the above, Lua Script: only involves token operations to ensure reasonable token production and reading.

Function analysis

It can be seen from the above process:

  1. There are multiple safeguards in place to ensure that limiting traffic will be completed.
  2. ifredis limiterInvalid, at least in processrate limiterOut.
  3. retryredis limiterThe mechanism ensures that it runs as smoothly as possible.

conclusion

The TokenLimit traffic limiting scheme in Go-Zero applies to instantaneous traffic impact, and the actual request scenario does not operate at a constant rate. The token bucket is quite pre-requested and will not be overwhelmed instantly when a real request arrives. When the traffic impact reaches a certain level, consumption will be carried out at a predetermined rate.

However, in the production of tokens, dynamic adjustment cannot be made according to the current traffic situation, which is not flexible enough and can be further optimized. In addition, you can refer to the Token bucket WIKI to layer Token buckets into different queues according to different traffic bandwidths.

reference

  • go-zero tokenlimit
  • Distributed stream limiting library provided by Go-Redis

If you like the article, please click on github star 🤝. And welcome to Go-Zero, github.com/tal-tech/go…