Golang channel實現limiter

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