Skip to Content

Go-Tips-Token-Bucket

Posted on

前几天写了一个docker registry的lib https://github.com/wrfly/reglib

然后在写example的时候发现并发太高, 把registry搞得502了, 然后就想办法怎么限制一下这个并发. 其实这个问题在之前面试的时候被问到过, 当时回答的不好. 因为一般的限流令牌桶的话都是有一个时间往里补充的, 我当时还没接触过别的, 所以就比较尴尬.

但是那天下午忽然就想到可以用channel去实现这个令牌桶, 但是操作恰好相反, 操作的时候不是取出令牌再进行下一步, 而是往里放, 对应的, 操作完成也不是再还回令牌, 而是取出.

可以看一个简单的例子:

package main

import (
	"fmt"
	"sync"
	"time"
)

func main() {
	count := 15
	limit := 5

	var wg sync.WaitGroup
	wg.Add(count)

	tbChan := make(chan struct{}, limit)
	for index := 0; index < count; index++ {
		go func(index int) {
			tbChan <- struct{}{}
			defer func() { <-tbChan }()
			defer wg.Done()

			fmt.Printf("this is %d\n", index)
			time.Sleep(time.Second)
		}(index)
	}

	wg.Wait()
}

https://play.golang.org/p/pdDf7NWhNob

运行的时候可以看到, 每次打印五个, 间隔一秒.

在playground里运行这个程序, 可以看到是打印是有顺序的:

this is 14
this is 0
this is 1
this is 2
this is 3
this is 4
this is 5
this is 6
this is 7
this is 8
this is 9
this is 10
this is 11
this is 12
this is 13

这是因为playground只给了你一个core, 你只有一个CPU, 所以一次只能处理一个goroutine, for循环结束之后, wg.Wait() 等待结束, 此时调度器切换到go出去的goroutine队列里, 那为什么先打印14, 而不是0呢? 是因为第14个goroutine还没有放到goroutine的队列里, 而是在main里捏着, 当main goroutine跑完的时候, 才将CPU的控制权交出, 并开始调度队列里的01234…

但如果你是在本机运行这段代码, 就会发现打印的时候不是按照顺序来的, 而是一种随机打印了, 这也说明, 在golang调度的时候, 每个goroutine的运行都与当前CPU状态有关. 但是, 每次第一行又都是14, 也就是最后一个goroutine, 这与上面的解释相同, 即当前goroutine还没有放入队列中.

  tokenbucket go run main.go
this is 14
this is 0
this is 8
this is 9
this is 7
this is 11
this is 10
this is 12
this is 13
this is 3
this is 1
this is 2
this is 4
this is 5
this is 6
  tokenbucket go run main.go
this is 14
this is 10
this is 0
this is 6
this is 12
this is 1
this is 4
this is 3
this is 5
this is 2
this is 8
this is 7
this is 9
this is 13
this is 11
  tokenbucket

如果想要达到playground的效果, 可以在运行程序之前加上GOMAXPROCS=1, 或者在程序开头加上一行 runtime.GOMAXPROCS(1).

然后再来看这个:

package main

import (
	"fmt"
	"runtime"
	"sync"
	"time"
)

func main() {
	count := 15
	limit := 5

	runtime.GOMAXPROCS(1)
	var wg sync.WaitGroup
	wg.Add(count)

	tbChan := make(chan struct{}, limit)
	for index := 0; index < count; index++ {
		go func(index int) {
			tbChan <- struct{}{}
			defer func() { <-tbChan }()
			defer wg.Done()

			fmt.Printf("this is %d\n", index)
			time.Sleep(time.Second)
		}(index)
	}

	time.Sleep(time.Millisecond)
	wg.Wait()
}

https://play.golang.org/p/67AruhE1hnV

这次可以看到, 程序输出是按顺序的了. 只因为多了一行 time.Sleep(time.Millisecond)

类似的, 还有:

https://play.golang.org/p/5h5YEnnBox2https://play.golang.org/p/5RRBRmfvz1U

也会达到相同的效果. 其实就是把第14号goroutine放到了队列里. 那怎样才会把goroutine放到队列里呢, 有一些关键词: time.Sleep, runtime.Gosched, go.

也可以看下别人的回答: why-is-time-sleep-required-to-run-certain-goroutines

comments powered by Disqus