Golang channel实现limiter

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

// 基于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)
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

// 令牌桶算法
// 产生令牌：均为间隔时间内（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
        }
    }
}