go chan 控制协程并发多数据源采集指定数量数据

场景：最近有个需求，从多个网站端抓取指定数量数据如40条。以下逻辑可满足

• 当某个协程已完成指定数量抓取，则停止其它协程任务抓取
• 当所有任务超过指定时间还未抓取到指定数量数据则停止
• 当所有任务执行完还没抓取到指定数量数据，不等超时立即返回结果

package main

import (
	"fmt"
	"sync/atomic"
	"time"
)

func main() {
	ch := make(chan int)
	taskDone := make(chan int32)
	d := make([]int, 0, 10)
	const limit = 40	// 需要抓取的数据量
	var taskTotal int32 = 0
	var taskCount int32 = 4 // 开启4个协程任务
    startTime := time.Now()
    // 开启多个协程
	go func(chi chan<- int) {
		fmt.Println("start1， 采集任务模拟")
		for i := 0; i < 10; i++ {
			if taskTotal >= limit {
				close(chi)
				fmt.Println("len:", taskTotal)
				break
			}
			atomic.AddInt32(&taskTotal, 1)
            // 将采集的数据推入 chan中
			chi <- i
		}
		// 完成一个任务减一 等于0时表示所有任务完成了
		atomic.AddInt32(&taskCount, -1) // 上报任务进度
        fmt.Println("任务1ok:", taskCount)
		taskDone <- taskCount	// 上报任务进度
	}(ch)
	go func(chi chan<- int) {
		fmt.Println("start2")
		for i := 11; i < 20; i++ {
			if taskTotal >= limit {
				close(chi)
				fmt.Println("len2:", taskTotal)
				break
			}

			atomic.AddInt32(&taskTotal, 1)
			chi <- i
		}

		atomic.AddInt32(&taskCount, -1)
        
        fmt.Println("任务2ok:", taskCount)
		taskDone <- taskCount
	}(ch)

	go func(chi chan<- int) {
		fmt.Println("start4")
		for i := 11; i < 20; i++ {
			if taskTotal >= limit {
				close(chi)
				fmt.Println("len2:", taskTotal)
				break
			}

			atomic.AddInt32(&taskTotal, 1)
			chi <- i
		}

		atomic.AddInt32(&taskCount, -1)
        fmt.Println("任务3ok:", taskCount)
		taskDone <- taskCount
	}(ch)
	go func(chi chan<- int) {
		fmt.Println("start3")
		for i := 21; i < 30; i++ {
			if taskTotal >= limit {
				close(chi)
				fmt.Println("len3:", taskTotal)
				break
			}

			atomic.AddInt32(&taskTotal, 1)
			chi <- i
		}
		atomic.AddInt32(&taskCount, -1)
        fmt.Println("任务4ok:", taskCount)
		taskDone <- taskCount
	}(ch)

	for {
		select {
		case v, ok := <-ch: //如果有数据，下面打印。但是有可能ch一直没数据
			if !ok {
				goto OK
			}
            d = append(d, v)
			fmt.Println("..", v, ok)
		case <-time.After(3 * time.Second): //3秒超时后还没抓取完则停止，注释掉其中一个任务进度报告并且没抓取到指定数量也会触发超时
			fmt.Println("超时")
			close(ch)
			goto OK
		case taskDoneCount := <-taskDone:	// 全部协程执行完，没采集到指定数量, 不等超时直接返回
			fmt.Println("任务数", taskDoneCount)
			if taskDoneCount == 0 {
				goto OK
			}
		}
	}
OK:
	fmt.Println(d, len(d), taskTotal, time.Since(startTime))
}