跳到主要内容

巧妙使用channel实现限流器

Keson

cover

普通Limiter

特点:无时间限制,只要不超过数量就可通过

// 基于channel阻塞实现
// 缺点:阻塞无时间限制
type (
ChannelLimiter struct {
bufferChannel chan golang.PlaceholderType
}
)

func NewChannelLimiter(limit int) *ChannelLimiter {
return &ChannelLimiter{bufferChannel: make(chan golang.PlaceholderType, limit)}
}

func (l *ChannelLimiter) Allow() bool {
select {
case l.bufferChannel <- golang.Placeholder:
return true
default:
return false
}
}

func (l *ChannelLimiter) Release() bool {
<-l.bufferChannel
return true
}

func (l *ChannelLimiter) Close() {
close(l.bufferChannel)
}

令牌桶Limiter

特点:令牌桶

// 令牌桶算法
// 产生令牌:均为间隔时间内(1秒)向指定桶中产生指定数量的令牌
// 消费令牌:从桶中获取令牌并消费
// 思路:通过channel阻塞原理来实现

type (
TokenBucketLimiter struct {
t *time.Ticker
bucket chan golang.PlaceholderType
doneC channel.DoneChan
limit int
rate int
stop func()
}
)

// rate:token put rates per second
// limit:max limit
func NewTokenBucketLimiter(rate, limit int) *TokenBucketLimiter {
t := time.NewTicker(time.Second)
doneC := channel.NewDoneChan()
bucket := make(chan golang.PlaceholderType, limit)
tbl := &TokenBucketLimiter{
rate: rate,
t: t,
bucket: bucket,
doneC: doneC,
limit: limit,
// only stop once
stop: routine.DoOnce(func() {
doneC.Stop()
close(bucket)
t.Stop()
}),
}
// 定时放置令牌
tbl.asyncPutTokens()
return tbl
}

// 通过尝试put golang.Placeholder 来达到是否有令牌可消费
func (tbl *TokenBucketLimiter) Allow() bool {
select {
case tbl.bucket <- golang.Placeholder:
return true
case <-tbl.doneC.Done():
return false
default:
return false
}
}

func (tbl *TokenBucketLimiter) Close() {
tbl.stop()
}

// 通过排空channel达到放置token的目的
func (tbl *TokenBucketLimiter) asyncPutTokens() {
safe.GoRun(func() {
for {
select {
case <-tbl.t.C:
tbl.drain()
case <-tbl.doneC.Done():
return
}
}
})
}

func (tbl *TokenBucketLimiter) drain() {
for i := 0; i < tbl.limit; i++ {
select {
case <-tbl.bucket:
default:
return
}
}
}