巧妙使用channel實現限流器

普通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
        }
    }
}