• #">背景#
  • #">协程池的概念#
  • #">基础版#
  • #">解决闭包引用问题#
  • #">添加睡眠功能#
  • #">待补充#

    背景#

    因与工作相关，所以本文中的数据都进行了更改，但逻辑是一样的。
    笔者的服务ServerA会请求服务ServerH获取一些数据，但ServerH的接口有个N秒内只能请求M次的限制，并返回false。而笔者的服务瞬时请求量远超M次，所以采用了协程池在收到103错误时，停止worker的运行N秒，然后再启动。

    协程池的概念#

    协程池的相关概念：
    要有一个一定数量大小的池子（pool），池子里存储需要执行的任务（task），还要有若干个工作协程（worker）。

    协程池要有启动，停止，睡眠的功能。

    下面是从零开始记录一下思想过程和遇到的问题。

    基础版#

    在此版本里，除了睡眠的功能，已经实现了一个基本的协程池。

    1. // workpool.go
    2. package workpool
    4. import (
    5.     "context"
    6.     "sync"
    7. )
    9. type TaskFunc func()
    11. type Task struct {
    12.     f TaskFunc
    13. }
    15. type WorkPool struct {
    16.     pool        chan *Task
    17.     workerCount int
    19.     stopCtx        context.Context
    20.     stopCancelFunc context.CancelFunc
    21.     wg             sync.WaitGroup
    22. }
    24. func (t *Task) Execute() {
    25.     t.f()
    26. }
    28. func New(workerCount, poolLen int) *WorkPool {
    29.     return &WorkPool{
    30.         workerCount: workerCount,
    31.         pool:        make(chan *Task, poolLen),
    32.     }
    33. }
    35. func (w *WorkPool) PushTask(t *Task) {
    36.     w.pool <- t
    37. }
    39. func (w *WorkPool) PushTaskFunc(f TaskFunc) {
    40.     w.pool <- &Task{
    41.         f: f,
    42.     }
    43. }
    45. func (w *WorkPool) work() {
    46.     for {
    47.         select {
    48.         case <-w.stopCtx.Done():
    49.             w.wg.Done()
    50.             return
    51.         case t := <-w.pool:
    52.             t.Execute()
    53.         }
    54.     }
    55. }
    57. func (w *WorkPool) Start() *WorkPool {
    58.     w.wg.Add(w.workerCount)
    59.     w.stopCtx, w.stopCancelFunc = context.WithCancel(context.Background())
    60.     for i := 0; i < w.workerCount; i++ {
    61.         go w.work()
    62.     }
    63.     return w
    64. }
    66. func (w *WorkPool) Stop() {
    67.     w.stopCancelFunc()
    68.     w.wg.Wait()
    69. }


    看起来没什么毛病，还挺简洁。其实不然…
    下面的程序是创建一个容量为50的workpool，并将通过3个worker输出100个数字。

    1. // workpool_test.go
    2. package workpool
    4. import (
    5.     "fmt"
    6.     "sync"
    7.     "testing"
    8. )
    10. func TestWorkPool_Start(t *testing.T) {
    11.     wg := sync.WaitGroup{}
    12.     wp := New(3, 50).Start()
    13.     lenth := 100
    14.     wg.Add(lenth)
    15.     for i := 0; i < lenth; i++ {
    16.         wp.PushTaskFunc(func() {
    17.             defer wg.Done()
    18.             fmt.Print(i, " ")
    19.         })
    20.     }
    21.     wg.Wait()
    22. }

    运行后输出结果如下：

    1. 50 50 50 50 50 50 50 50 50 50 50 50 50 50 50 50 50 50 50 50 51 51 51 51 69 72 78 78 80 81 81 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 83 84 84 84 84 50 84
    2. 100 100 100 100 100 100 100 100 100 100 50 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 84 100 100 100


    这和想象中的输出 0-99 相差甚远。
    其原因在于闭包函数对于外部变量是引用的，所以在函数执行的时候，i 的值早就已经改变了。下面是一个关于闭包的简单例子。

    1. x := 1
    2. f := func() {
    3.     println(x)
    4. }
    5. x = 2
    6. x = 3
    7. f() // 3

    可以将 f() 的调用时机对应为协程池中的 t.Execute()。

    解决闭包引用问题#

    既然是因为闭包引用导致的问题，那就不使用闭包了呗。
    可以把参数传到函数内，但是因为并不知道将要执行的函数需要的参数个数及类型，所以只能是使用不定长的interface{}TaskFunc，在使用的时候进行断言。
    以下仅列出改动部分：

    1. // workpool.go
    2. type TaskFunc func(args ...interface{})
    4. type Task struct {
    5.     f    TaskFunc
    6.     args []interface{}
    7. }
    9. func (t *Task) Execute() {
    10.     t.f(t.args...)
    11. }
    13. func (w *WorkPool) PushTaskFunc(f TaskFunc, args ...interface{}) {
    14.     w.pool <- &Task{
    15.         f:    f,
    16.         args: args,
    17.     }
    18. }

    以下是测试程序：

    1. // workpool_test.go
    2. package workpool
    4. import (
    5.     "fmt"
    6.     "sync"
    7.     "testing"
    8. )
    10. func TestWorkPool_Start(t *testing.T) {
    11.     wg := sync.WaitGroup{}
    12.     wp := New(3, 50).Start()
    13.     lenth := 100
    14.     wg.Add(lenth)
    15.     for i := 0; i < lenth; i++ {
    16.         wp.PushTaskFunc(func(args ...interface{}) {
    17.             defer wg.Done()
    18.             fmt.Print(args[0].(int), " ")
    19.         }, i)
    20.     }
    21.     wg.Wait()
    22. }

    输出内容如下：

    1. 0 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 2 1 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 26 48 49 51 52 53 54 55 56 50 58 59 57 61 62 63 64 65 66 25 68 6
    2. 9 70 71 72 73 67 75 76 77 74 79 78 81 82 83 84 60 86 87 88 89 90 91 92 85 94 95 96 97 98 99 80 93

    虽然顺序是错乱的，但这是正常情况，闭包引用问题已解决。

    添加睡眠功能#

    基于开头的应用场景，在任意一个被worker执行的任务收到ServerH的103错误后，要停止所有worker一段时间，因为再一直请求也没有意义。
    这个版本已经与笔者正在使用的相差无几了

    1. // workpool.go
    2. package workpool
    4. import (
    5.     "context"
    6.     "fmt"
    7.     "sync"
    8.     "sync/atomic"
    9.     "time"
    10. )
    12. type Flag int64
    14. const (
    15.     FLAG_OK    Flag = 1 << iota
    16.     FLAG_RETRY Flag = 1 << iota
    17. )
    19. type TaskFunc func(w *WorkPool, args ...interface{}) Flag
    21. type Task struct {
    22.     f    TaskFunc
    23.     args []interface{}
    24. }
    26. type WorkPool struct {
    27.     pool        chan *Task
    28.     workerCount int
    30.         // stop相关
    31.     stopCtx        context.Context
    32.     stopCancelFunc context.CancelFunc
    33.     wg             sync.WaitGroup
    35.         // sleep相关
    36.     sleepCtx        context.Context
    37.     sleepCancelFunc context.CancelFunc
    38.     sleepSeconds    int64
    39.     sleepNotify     chan bool
    40. }
    42. func (t *Task) Execute(w *WorkPool) Flag {
    43.     return t.f(w, t.args...)
    44. }
    46. func New(workerCount, poolLen int) *WorkPool {
    47.     return &WorkPool{
    48.         workerCount: workerCount,
    49.         pool:        make(chan *Task, poolLen),
    50.         sleepNotify: make(chan bool),
    51.     }
    52. }
    54. func (w *WorkPool) PushTask(t *Task) {
    55.     w.pool <- t
    56. }
    58. func (w *WorkPool) PushTaskFunc(f TaskFunc, args ...interface{}) {
    59.     w.pool <- &Task{
    60.         f:    f,
    61.         args: args,
    62.     }
    63. }
    65. func (w *WorkPool) work(i int) {
    66.     for {
    67.         select {
    68.         case <-w.stopCtx.Done():
    69.             w.wg.Done()
    70.             return
    71.         case <-w.sleepCtx.Done():
    72.             time.Sleep(time.Duration(w.sleepSeconds) * time.Second)
    73.         case t := <-w.pool:
    74.             flag := t.Execute(w)
    75.             if flag&FLAG_RETRY != 0 {
    76.                 w.PushTask(t)
    77.                 fmt.Printf("work %v PushTask,pool length %v\n", i, len(w.pool))
    78.             }
    79.         }
    80.     }
    81. }
    83. func (w *WorkPool) Start() *WorkPool {
    84.     fmt.Printf("workpool run %d worker\n", w.workerCount)
    85.     w.wg.Add(w.workerCount + 1)
    86.     w.stopCtx, w.stopCancelFunc = context.WithCancel(context.Background())
    87.     w.sleepCtx, w.sleepCancelFunc = context.WithCancel(context.Background())
    88.     go w.sleepControl()
    89.     for i := 0; i < w.workerCount; i++ {
    90.         go w.work(i)
    91.     }
    92.     return w
    93. }
    95. func (w *WorkPool) Stop() {
    96.     w.stopCancelFunc()
    97.     w.wg.Wait()
    98. }
    100. func (w *WorkPool) sleepControl() {
    101.     fmt.Println("sleepControl start...")
    102.     for {
    103.         select {
    104.         case <-w.stopCtx.Done():
    105.             w.wg.Done()
    106.             return
    107.         case <-w.sleepNotify:
    108.             fmt.Printf("receive sleep notify start...\n")
    109.             w.sleepCtx, w.sleepCancelFunc = context.WithCancel(context.Background())
    110.             w.sleepCancelFunc()
    111.             fmt.Printf("sleepControl will star sleep %v s\n", w.sleepSeconds)
    112.             time.Sleep(time.Duration(w.sleepSeconds) * time.Second)
    113.             w.sleepSeconds = 0
    114.             fmt.Println("sleepControl was end sleep")
    115.         }
    116.     }
    117. }
    120. func (w *WorkPool) SleepNotify(seconds int64) {
    121.     // 因为需要CAS操作，所以sleepSeconds没有采用time.Duration类型
    122.     // 成功设置后才发出通知
    123.     if atomic.CompareAndSwapInt64(&w.sleepSeconds, 0, seconds) {
    124.         fmt.Printf("sleepSeconds set %v\n", seconds)
    125.         w.sleepNotify <- true
    126.     }
    127. }

    下面的测试程序中，模拟了一下ServerH，其使用场景与笔者工作中大同小异。

    1. // workpool_test.go
    2. package workpool
    4. import (
    5.     "fmt"
    6.     "sync"
    7.     "testing"
    8.     "time"
    9. )
    11. // 这里模拟ServerH服务的限流操作
    12. var serverh = &server{max: 10, interval: 5}
    14. type server struct {
    15.     count    int
    16.     max      int
    17.     lasttime time.Time
    18.     interval int64
    19.     mu       sync.Mutex
    20. }
    22. func (s *server) Access(i int) bool {
    23.     now := time.Now()
    25.     s.mu.Lock()
    26.     defer s.mu.Unlock()
    28.     time.Sleep(100 * time.Millisecond)
    30.     if s.lasttime.Unix() <= 0 || s.count >= s.max {
    31.         if now.After(s.lasttime) {
    32.             s.count = 1
    33.             s.lasttime = time.Unix(now.Unix()+s.interval, 0)
    34.             return true
    35.         }
    36.         fmt.Printf("Access false,i=%d \n", i)
    37.         return false
    38.     } else {
    39.         s.count++
    40.         fmt.Printf("Access true,i=%d s.count %d\n", i, s.count)
    41.         return true
    42.     }
    43. }
    45. // 这里是笔者服务的逻辑
    46. func TestWorkPool_Start(t *testing.T) {
    47.     wp := New(3, 100).Start()
    48.     for i := 0; i < 100; i++ {
    49.         time.Sleep(100 * time.Millisecond)
    50.         wp.PushTaskFunc(func(w *WorkPool, args ...interface{}) Flag {
    51.             if !serverh.Access(args[0].(int)) {
    52.                                 // 发送睡眠5秒的通知
    53.                 w.SleepNotify(5) 
    54.                                 // 此次未执行成功，要将该任务放回协程池
    55.                 return FLAG_RETRY 
    56.             }
    57.             return FLAG_OK
    58.         }, i)
    59.     }
    60.     time.Sleep(100 * time.Second)
    61. }


    输出内容如下：

    1. workpool run 3 worker
    2. sleepControl start...
    3. Access true,i=1 s.count 2
    4. Access true,i=2 s.count 3
    5. Access true,i=3 s.count 4
    6. Access true,i=4 s.count 5
    7. Access true,i=5 s.count 6
    8. Access true,i=6 s.count 7
    9. Access true,i=7 s.count 8
    10. Access true,i=8 s.count 9
    11. Access true,i=9 s.count 10
    12. Access false,i=10 
    13. sleepSeconds set 5
    14. ...........

    待补充#

    重试次数的逻辑