让 sync.WaitGroup 支持并发数量限制

如何限制协程数量

业务场景中，需要我们要做完N项工作，假设使用一项任务起一个协程，那么代码如下：

func main() {
   var wg = sync.WaitGroup{}

    taskCount := 5 // 任务数量
    for i := 0; i < taskCount; i++ {
        wg.Add(1)

        go func(i int) {
            defer wg.Done()

            fmt.Println("go func ", i)
            time.Sleep(time.Second * time.Duration(rand.Intn(5)))  // 模拟工作
        }(i)
    }

    wg.Wait()
}

但是 如果taskcount 很大，那么我们就会短时间内启动很多协程，这对系统来说，是有风险的。简单改下代码：

var task_chan = make(chan bool, 5) // 5 为 channel长度

func main() {
    wg := sync.WaitGroup{}

    for i := 0; i < math.MaxInt64; i++ {
        wg.Add(1)

        task_chan <- true
        go func(i int) {
            defer func() { <-task_chan }()
            defer wg.Done()

            fmt.Println("go func ", i)
            time.Sleep(time.Second * time.Duration(rand.Intn(5)))    // 模拟工作
        }(i)
    }

    wg.Wait()
}

同样的思路，github有个扩展了WaitGroup的方案可以实现限制并发数量：https://github.com/remeh/sizedwaitgroup

type SizedWaitGroup struct {
    size    int
    current chan struct{}
    wg      sync.WaitGroup
}

func NewSizedWaitGroup(limit int) *SizedWaitGroup {
    size := math.MaxInt32 // 2^32 - 1
    if limit > 0 {
        size = limit
    }

    return &SizedWaitGroup{
        size:    size,
        current: make(chan struct{}, size),
        wg:      sync.WaitGroup{},
    }
}

func (s *SizedWaitGroup) Add() error {
    return s.AddWithContext(context.Background())
}

func (s *SizedWaitGroup) AddWithContext(ctx context.Context) error {
    select {
    case <-ctx.Done():
        return ctx.Err()
    case s.current <- struct{}{}:
        break
    }
    s.wg.Add(1)
    return nil
}

func (s *SizedWaitGroup) Done() {
    <-s.current
    s.wg.Done()
}

func (s *SizedWaitGroup) Wait() {
    s.wg.Wait()
}

使用：

func main() {
    swg := sizedwaitgroup.New(8)
    for i := 0; i < 50; i++ {
        swg.Add()
        go func(i int) {
            defer swg.Done()
            query(i)
        }(i)
    }

    swg.Wait()
}

func query(i int) {
    fmt.Println(i)
    time.Sleep(10 * time.Second)
}