Microservice – current limiting: I Implementation of token bucket algorithm by golang

Time:2022-5-4

At first, it was because we had to pull some third-party data, and the API interfaces of the third-party were limited flow measures. For example, 6000 / min, 500 / min. If you want to pull data, you can use multiple collaborative processes. But it is easy to overclock, so I want to write a current limiting Dongdong. There is a token bucket on the Internet, which is similar to the following: (schematic diagram on the Internet)

Token bucket principle

  1. There is a bucket, and the bucket has capacity (Cap: the capacity of the bucket).
  2. Then add a token to the bucket at a constant speed.
  3. If the bucket has reached its capacity, the newly added token will be discarded.
  4. Each consumption is to take a token from the bucket.

It feels very simple, so let’s practice it. First, let’s write the structure of the bucket

Barrel structure

Here, token is the place where tokens are stored. An empty struct {} is used here. As we all know, empty struct {} consumes very little memory. Mu is a mutually exclusive lock. Because it involves multiple coprocess operations to tokens in the bucket, a lock is added here.

package limit

import (
	"sync"
	"time"
)

//Token bucket
type bucket struct {
	Rate int // frequency per minute (how many tokens are added per minute)
	Token Chan struct {} // where tokens are stored
    Cap int // capacity
    mu    *sync. Mutex // lock in barrel
    Pause bool // pause
    Stop bool // stop
}

Instantiate a bucket

It is judged here that the capacity of the bucket must be greater than 0.

//Get new bucket
//Rate: how many times per minute
//Cap: the capacity of the bucket must be greater than or equal to 1
func NewBucket(rate, cap int) *bucket {
	if cap < 1 {
		panic("limit bucket cap error")
	}
	return &bucket{
		token: make(chan struct{}, cap),
		rate:  rate,
		mu:    new(sync.Mutex),
		cap:   cap,
	}
}

Start timing

Here’s a new goroutine for timing

//Start
func (b *bucket) Start() {
	go b.addToken()
}

//Join token
func (b *bucket) addToken() {
	for {
		b.mu.Lock()
		if b.stop {
			close(b.token)
			b.mu.Unlock()
			return
		}
		if b.pause {
			b.mu.Unlock()
			time.Sleep(time.Second)
			continue
		}
		b.token

Consume a token

To get a token here is to get a data from Chan

//Consumption, there will be automatic blocking
func (b *bucket) GetToken() {

Pause, stop, reset

Using the attributes of bucket: pause and stop, you can also add some small functions.

//Pause
func (b *bucket) Pause() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.pause = true
}

//Stop
func (b *bucket) Stop() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.stop = true
}

//Reset
func (b *bucket) Reset() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.token = make(chan struct{}, b.cap)
}

We have basically realized the functions we want. Test it. At this time, my colleague threw me a bag: golang org/x/time/rate。