Token bucket

Posted 2020-12-15 littlepage

Token bucket

令牌桶限流机制：

在长连接时，当一个http请求来临，我们需要进行GetConn，放入一个令牌进桶中，当我们的连接释放，于是可以释放连接。
我们控制桶的大小即可控制http的流量。

令牌桶算法，最近在学Golang，而对于HTTP的限流，Golang的令牌桶实现更为简洁，Golang的channel可以控制协程通信，所以令牌桶会简单很多。

shared channel instead of shared memory.

我们可以知道，channel可以作为共享内存。

Go语言代码案例：

// 令牌桶的数据结构
type ConnLimiter struct {
	concurrentConn int
	bucket chan int
}

// 初始化令牌桶 cc为capacity，令牌桶容量
func NewConnLimiter(cc int) *ConnLimiter {
	return &ConnLimiter {
		concurrentConn: cc,
		bucket: make(chan int, cc),
	}
}

// 获取一个连接
func (cl *ConnLimiter) GetConn() bool {
	if len(cl.bucket) >= cl.concurrentConn {
		log.Printf("Reached the rate limitation.")
		return false
	}

	cl.bucket <- 1
	return true
}

// 释放一个连接
func (cl *ConnLimiter) ReleaseConn() {
	c :=<- cl.bucket
	log.Printf("New connction coming: %d", c)
}

最后我们在实际场景中，把令牌桶添加到一个http中间件中即可

func (m middleWareHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !m.l.GetConn() {
		sendErrorResponse(w, http.StatusTooManyRequests, "Too many requests")
		return
	}
	m.r.ServeHTTP(w, r)
	defer m.l.ReleaseConn()
}