1
2
3
4
5
6
7
8

type Clock
type Limiter
    func New(rate int, opts ...Option) Limiter
    func NewUnlimited() Limiter
type Option
    func Per(per time.Duration) Option
    func WithClock(clock Clock) Option
    func WithSlack(slack int) Option

1
2
3
4

type Clock interface {
	Now() time.Time
	Sleep(time.Duration)
}

1
2
3
4

type Limiter interface {
	// Take should block to make sure that the RPS is met.
	Take() time.Time
}

1
2
3
4
5
6
7
8
9
10
11
12
13

func NewLimiter(rps int) gin.HandlerFunc {
	limiters := &sync.Map{}
	return func(c *gin.Context) {
		// 获取限速器
		// key 除了 ip 之外也可以是其他的，例如 header，user name 等
		key := c.ClientIP()
		l, _ := limiters.LoadOrStore(key, ratelimit.New(rps))
		now := l.(ratelimit.Limiter).Take()
		fmt.Printf("now: %s\n", now)
		c.Next()
	}
}

1
2
3
4
5
6
7
8
9

func main() {
	e := gin.Default()
	// 新建一个限速器，允许突发 3 个并发
	e.Use(NewLimiter(3))
	e.GET("ping", func(c *gin.Context) {
		c.String(http.StatusOK, "pong")
	})
	e.Run(":8080")
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

go-stress-testing-linux -c 100 -u http://localhost:8080/ping
─────┬───────┬───────┬───────┬────────┬────────┬────────┬────────┬────────┬────────┬────────
 耗时 │ 并发数│  成功数│ 失败数 │   qps  │最长耗时│最短耗时 │平均耗时 │下载字节 │字节每秒 │ 错误码
─────┼───────┼───────┼───────┼────────┼────────┼────────┼────────┼────────┼────────┼────────
   1s│     13│     13│      0│  233.55│  676.10│    5.82│   85.64│      52│      51│200:13
   2s│     16│     16│      0│   62.25│ 1675.17│    5.82│  321.30│      64│      31│200:16
   3s│     19│     19│      0│   31.24│ 2673.94│    5.82│  640.20│      76│      25│200:19
   3s│     20│     20│      0│   26.37│ 3006.49│    5.82│  758.51│      80│      26│200:20
*************************  结果 stat  ****************************
处理协程数量: 20
请求总数（并发数*请求数 -c * -n）: 20 总请求时间: 3.011 秒 successNum: 20 failureNum: 0
*************************  结果 end   ****************************

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

type state struct {
	last     time.Time
	sleepFor time.Duration
}
type atomicLimiter struct {
	state unsafe.Pointer
	//lint:ignore U1000 Padding is unused but it is crucial to maintain performance
	// of this rate limiter in case of collocation with other frequently accessed memory.
	padding [56]byte // cache line size - state pointer size = 64 - 8; created to avoid false sharing.
	perRequest time.Duration
	maxSlack   time.Duration
	clock      Clock
}

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

func (t *atomicLimiter) Take() time.Time {
	var (
        // 状态
		newState state
        // 用于表示原子操作是否成功
		taken    bool
        // 需要 sleep 的时间
		interval time.Duration
	)
    // 如果 CAS 操作不成功就一直尝试
	for !taken {
        // 获取当前的时间
		now := t.clock.Now()
        // load 出上一次调用的时间
		previousStatePointer := atomic.LoadPointer(&t.state)
		oldState := (*state)(previousStatePointer)
		newState = state{
			last:     now,
			sleepFor: oldState.sleepFor,
		}
		// 如果 last 是零值的话，表示之前就没用过，直接保存返回即可
		if oldState.last.IsZero() {
			taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
			continue
		}
		// sleepFor 是需要睡眠的时间，由于引入了松弛时间，所以 sleepFor 可能是一个
        // maxSlack ~ 0 之间的一个值，所以这里需要将现在的需要 sleep 的时间和上一次
        // sleepFor 的值相加
		newState.sleepFor += t.perRequest - now.Sub(oldState.last)
        // 如果距离上一次调用已经很久了，sleepFor 可能会是一个很小的值
        // 最小值只能是 maxSlack 的大小
		if newState.sleepFor < t.maxSlack {
			newState.sleepFor = t.maxSlack
		}
        // 如果 sleepFor 大于 0  的话，计算出需要 sleep 的时间
        // 然后将 state.sleepFor 置零
		if newState.sleepFor > 0 {
			newState.last = newState.last.Add(newState.sleepFor)
			interval, newState.sleepFor = newState.sleepFor, 0
		}
        // 保存状态
		taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
	}
    // sleep interval
	t.clock.Sleep(interval)
	return newState.last
}