前言
本篇文章将分析Kitex源码中关于限流的部分，代码位于github.com/cloudwego/kitex/pkg/limit和github.com/cloudwego/kitex/pkg/limiter文件夹下，Kitex版本为v0.11.3。如有问题欢迎指正。
limit
// kitex/pkg/limit/limit.go
// 定义接口Updater，能够动态修改limit
type Updater interface {
UpdateLimit(opt *Option) (updated bool)
}
// 用于配置一个limiter的最大连接数和最大QPS信息
type Option struct {
MaxConnections int
MaxQPS int
// 用于接受上述的Updater以动态修改limit
UpdateControl func(u Updater)
}
// 正确性校验
func (lo *Option) Valid() bool {
return lo.MaxConnections > 0 || lo.MaxQPS > 0
}
参照server.go中的相关代码
// kitex/server/server.go
if limits != nil && limits.UpdateControl != nil {
updater := limiter.NewLimiterWrapper(connLimit, qpsLimit)
limits.UpdateControl(updater)
}
可以看到这段代码中传入的Updater是通过limiter.NewLimiterWrapper实现的，而该方法就是将参数中的ConcurrencyLimiter和RateLimiter封装成了一个Updater
// kitex/pkg/limiter/limiter.go
func NewLimiterWrapper(conLimit ConcurrencyLimiter, qpsLimit RateLimiter) limit.Updater {
return &limitWrapper{
conLimit: conLimit,
qpsLimit: qpsLimit,
}
}
type limitWrapper struct {
conLimit ConcurrencyLimiter
qpsLimit RateLimiter
}
connection_limiter
connection_limiter实现了ConcurrencyLimiter接口
// kitex/pkg/limiter/limiter.go
// ConcurrencyLimiter limits the number of concurrent access towards the protected resource.
// The implementation of ConcurrencyLimiter should be concurrent safe.
// 注意到此处要求ConcurrencyLimiter的实现应该满足并发安全
type ConcurrencyLimiter interface {
// Acquire reports if next access to the protected resource is allowed.
Acquire(ctx context.Context) bool
// Release claims a previous taken access has released the resource.
Release(ctx context.Context)
// Status returns the total quota and occupied.
Status(ctx context.Context) (limit, occupied int)
}
在connectionLimiter中定义了两个属性，lim表示该limiter的Maxconnections（小于等于0表示无限制），curr表示当前的连接总数。
// kitex/pkg/limiter/connection_limiter.go
type connectionLimiter struct {
lim int32
curr int32
}
Acquire方法获取一个连接，将curr减一；Release方法释放一个连接，将curr加一。在代码中可以看到，为了保证并发安全，对于limiter中属性的操作均使用atomic库下的方法进行。此外，Acquire方法没有采用先判断后分配连接的方式，而是先直接将curr减一，若无法超过limit数量的限制则不分配，直接执行Release恢复curr的数据。
// kitex/pkg/limiter/limiter.go
// Acquire tries to increase the connection counter.
// The return value indicates whether the operation is allowed under the connection limitation.
// Acquired is executed in `OnActive` which is called when a new connection is accepted, so even if the limit is reached
// the count is still need increase, but return false will lead the connection is closed then Release also be executed.
func (ml *connectionLimiter) Acquire(ctx context.Context) bool {
limit := atomic.LoadInt32(&ml.lim)
x := atomic.AddInt32(&ml.curr, 1)
return x <= limit || limit <= 0
}
// Release decrease the connection counter.
func (ml *connectionLimiter) Release(ctx context.Context) {
atomic.AddInt32(&ml.curr, -1)
}
// UpdateLimit updates the limit.
func (ml *connectionLimiter) UpdateLimit(lim int) {
atomic.StoreInt32(&ml.lim, int32(lim))
}
// Status returns the current status.
func (ml *connectionLimiter) Status(ctx context.Context) (limit, occupied int) {
limit = int(atomic.LoadInt32(&ml.lim))
occupied = int(atomic.LoadInt32(&ml.curr))
return
}
qps_limiter
qps_limiter实现了RateLimiter接口，采用了令牌桶算法。有关令牌桶算法的原理可参考这篇文章：[服务或接口限流算法1/2]-漏桶算法及令牌桶算法解析
// kitex/pkg/limiter/limiter.go
// RateLimiter limits the access rate towards the protected resource.
type RateLimiter interface {
// Acquire reports if next access to the protected resource is allowed.
Acquire(ctx context.Context) bool
// Status returns the rate limit.
Status(ctx context.Context) (max, current int, interval time.Duration)
}
在qpsLimiter中定义了以下属性：
limit表示MaxQPS，tokens为目前桶内剩余的令牌数量，interval是令牌刷新的时间间隔，once为每次刷新释放的令牌数，ticker利用time.Ticker执行定时任务
// kitex/pkg/limiter/qps_limiter.go
type qpsLimiter struct {
limit int32
tokens int32
interval time.Duration
once int32
ticker *time.Ticker
tickerDone chan bool
}
分析一下calcOnce方法
// kitex/pkg/limiter/qps_limiter.go
func calcOnce(interval time.Duration, limit int) int32 {
if interval > time.Second {
interval = time.Second
}
once := int32(float64(limit) / (fixedWindowTime.Seconds() / interval.Seconds()))
if once < 0 {
once = 0
}
return once
}
可以看到当interval小于1秒时，once计算出了每间隔interval的时间所需要释放的tokens数量。而当interval大于1秒时，once一次性就释放limit数量的tokens。（和connection_limiter类似，limit小于0表示无限制）
然后来看一下startTicker方法
// kitex/pkg/limiter/qps_limiter.go
func (l *qpsLimiter) startTicker(interval time.Duration) {
l.ticker = time.NewTicker(interval)
defer l.ticker.Stop()
l.tickerDone = make(chan bool, 1)
tc := l.ticker.C
td := l.tickerDone
// ticker and tickerDone can be reset, cannot use l.ticker or l.tickerDone directly
for {
select {
case <-tc:
l.updateToken()
case <-td:
return
}
}
}
这里启动了一个时间间隔为interval的定时任务，当定时任务正常触发时，执行updateToken方法；当使用stopTicker将tickerDone置为true时，停止当前的定时任务。
// kitex/pkg/limiter/qps_limiter.go
func (l *qpsLimiter) updateToken() {
if atomic.LoadInt32(&l.limit) < atomic.LoadInt32(&l.tokens) {
return
}
once := atomic.LoadInt32(&l.once)
delta := atomic.LoadInt32(&l.limit) - atomic.LoadInt32(&l.tokens)
if delta > once || delta < 0 {
delta = once
}
newTokens := atomic.AddInt32(&l.tokens, delta)
if newTokens < once {
atomic.StoreInt32(&l.tokens, once)
}
}
updateToken实现的就是间隔interval时间后释放tokens的操作，即将tokens数量设置为once。
