如何優雅的控制goroutine的數量

1,為什么要控制goroutine的數量？ 

goroutine固然好，但是數量太多了，往往會帶來很多麻煩，比如耗盡系統資源導致程序崩潰，或者CPU使用率過高導致系統忙不過來。比如：

for i:=0; i < 10000; i++ {
    go work()
}

2,用什么方法控制goroutine的數量？ 

要在每一次執行go之前判斷goroutine的數量，如果數量超了，就要阻塞go的執行。第一時間想到的就是使用通道。每次執行的go之前向通道寫入值，直到通道滿的時候就阻塞了，如下：

var ch chan int

func work() {
    //do something
    <-ch
}

func main() {
    ch = make(chan int, 10)
    for i:=0; i < 10000; i++ {
       ch <- 1
       go work()
    }
}

這樣每次同時運行的goroutine就被限制為10個了。但是新的問題出現了，因為並不是所有的goroutine都執行完了，在main函數退出之后，還有一些goroutine沒有執行完就被強制結束了。這個時候我們就需要用到sync.WaitGroup。使用WaitGroup等待所有的goroutine退出。如下：

var wg *sync.WaitGroup

func work() {
    defer wg.Done()
    //do something
}

func main() {
    wg = &sync.WaitGroup{}
    for i:=0; i < 10000; i++ {
       wg.Add(1)
       go work()
    }
    wg.Wait()//等待所有goroutine退出
}

3,優雅的使用並控制goroutine的數量 

綜上所述，我們封裝一下，代碼如下：

package gpool

import (
    "sync"
)

type pool struct {
    queue chan int
    wg    *sync.WaitGroup
}

func New(size int) *pool {
    if size <= 0 {
        size = 1
    }
    return &pool{
        queue: make(chan int, size),
        wg:    &sync.WaitGroup{},
    }
}

func (p *pool) Add(delta int) {
    for i := 0; i < delta; i++ {
        p.queue <- 1
    }
    for i := 0; i > delta; i-- {
        <-p.queue
    }
    p.wg.Add(delta)
}

func (p *pool) Done() {
    <-p.queue
    p.wg.Done()
}

func (p *pool) Wait() {
    p.wg.Wait()
}

來段測試代碼：

package gpool_test

import (
    "runtime"
    "testing"
    "time"
    "gpool"
)

func Test_Example(t *testing.T) {
    pool := gpool.New(100)
    println(runtime.NumGoroutine())
    for i := 0; i < 1000; i++ {
        pool.Add(1)
        go func() {
            time.Sleep(time.Second)
            println(runtime.NumGoroutine())
            pool.Done()
        }()
    }
    pool.Wait()
    println(runtime.NumGoroutine())
}

good job，Over～