基于chan和WaitGroup实现Goroutine池

最近闲来无事，于是封装了goroutine常用到的一个场景：chan + sync.WaitGroup，实现控制goroutine并发和等待任务执行完毕。

当然，也有其他现成的库实现该功能，比如ants。

gpool.go

// gpool -> goroutine pool with wait
package main

import (
    "sync"
)

type GPool struct {
    sync.WaitGroup  // 等待任务执行完
    buf chan struct{} // 限制并发
}

// n, goroutine并发数
// c, wait group需要等待完成的groutine总数
func NewGPool(n int, c int) *GPool {
    p := &GPool{buf: make(chan struct{}, n)}
    p.Add(c)
    return p
}

func (p *GPool) Exec(f func(interface{}), args interface{}) {
    p.buf <- struct{}{}
    go func(args interface{}) {
        defer p.Done()
        f(args)
        <-p.buf
    }(args)
}

 

main函数调用, main.go

// wg-chan
package main

import (
    "fmt"
    "sync"
    "time"
)

var (
    lock sync.Mutex
    sum  int
)

func task(args interface{}) {
    time.Sleep(time.Second * 1)
    fmt.Printf("[%s]%v\n", time.Now().Format(time.RFC3339Nano), args)
    lock.Lock()
    defer lock.Unlock()
    if v, ok := args.(int); ok {
        sum += v
    }
}

func t1() {
    loops := 10
    var pool = NewGPool(10, loops)
    for i := 0; i < loops; i++ {
        go pool.Exec(task, i)
    }
    pool.Wait()
    fmt.Println("[total]", sum)
}

func main() {
    t1()
}