Go语言小技巧(1)-控制并发数

需求

虽然goroutine很廉价，但是有时候还是希望控制goroutine并发数。

实现

来自K8S源代码

k8s.io\client-go\util\workqueue\parallelizer.go

type DoWorkPieceFunc func(piece int)

// Parallelize is a very simple framework that allow for parallelizing
// N independent pieces of work.

func Parallelize(workers, pieces int, doWorkPiece DoWorkPieceFunc) {
    toProcess := make(chan int, pieces)
    for i := 0; i < pieces; i++ {
        toProcess <- i
    }
    close(toProcess)

    if pieces < workers {
        workers = pieces
    }

    wg := sync.WaitGroup{}
    wg.Add(workers)
    for i := 0; i < workers; i++ {
        go func() {
            defer utilruntime.HandleCrash()
            defer wg.Done()
            for piece := range toProcess {
                doWorkPiece(piece)
            }
        }()
    }
    wg.Wait()
}

分析

• workers: int, goroutine并发数
• pieces: int, 任务总数
• doWorkPiece： func(piece int)，任务处理函数

1. 将任务索引全部放入channel中；
2. 通过for循环，创建worker个goroutine；
3. 每个goroutine从channel中领取任务