100 行写一个 go 的协程池 (任务池)

前言

go 的 goroutine 提供了一种较线程而言更廉价的方式处理并发场景, go 使用二级线程的模式, 将 goroutine 以 M:N 的形式复用到系统线程上, 节省了 cpu 调度的开销, 也避免了用户级线程(协程)进行系统调用时阻塞整个系统线程的问题。【1】html

但 goroutine 太多仍会致使调度性能降低、GC 频繁、内存暴涨, 引起一系列问题。在面临这样的场景时, 限制 goroutine 的数量、重用 goroutine 显然颇有价值。git

本文正是针对上述状况而提供一种简单的解决方案, 编写一个协程池(任务池)来实现对 goroutine 的管控。github

思路

要解决这个问题, 要思考两个问题安全

  • goroutine 的数量如何限制, goroutine 如何重用
  • 任务如何执行

goroutine 的数量如何限制, goroutine 如何重用

说到限制和重用, 那么最早想到的就是池化。好比 TCP 链接池, 线程池, 都是有效限制、重用资源的最好实践。因此, 咱们能够建立一个 goroutine 池, 用来管理 goroutine。并发

任务如何执行

在使用原生 goroutine 的场景中, 运行一个任务直接启动一个 goroutine 来运行, 在池化的场景而言, 任务也是要在 goroutine 中执行, 可是任务须要任务池来放入 goroutine。函数

生产者消费者模型

在链接池中, 链接在使用时从池中取出, 用完后放入池中。对于 goroutine 而言, goroutine 经过语言关键字启动, 没法像链接同样操做。那么如何让 goroutine 能够执行任务, 且执行后能够从新用来执行其它任务呢?这里就须要使用生产者消费者模型了:性能

生产者 --(生产任务)--> 队列 --(消费任务)--> 消费者测试

用来执行任务的 goroutine 能够做为消费者, 操做任务池的 goroutine 做为生产者, 而队列则可使用 go 的 buffer channel, 任务池的建模到此结束。ui

实现

Talk is cheap. Show me the code.atom

任务的定义

任务要包含须要执行的函数、以及函数要传的参数, 由于参数类型、个数不肯定, 这里使用可变参数和空接口的形式

type Task struct {
    Handler func(v ...interface{})
    Params  []interface{}
}

任务池的定义

任务池的定义包括了池的容量 capacity、当前运行的 worker(goroutine)数量 runningWorkers、任务队列(channel)taskC、关闭任务池的 channel closeC 以及任务池的状态 state(运行中或已关闭, 用于安全关闭任务池)

type Pool struct {
    capacity       uint64
    runningWorkers uint64
    state          int64
    taskC          chan *Task
    closeC         chan bool
}

任务池的构造函数:

var ErrInvalidPoolCap = errors.New("invalid pool cap")

const (
    RUNNING = 1
    STOPED = 0
)

func NewPool(capacity uint64) (*Pool, error) {
    if capacity <= 0 {
        return nil, ErrInvalidPoolCap
    }
    return &Pool{
        capacity: capacity,
        state:    RUNNING,
        // 初始化任务队列, 队列长度为容量
        taskC:    make(chan *Task, capacity),
        closeC:   make(chan bool),
    }, nil
}

启动 worker

新建 run() 方法做为启动 worker 的方法:

func (p *Pool) run() {
    p.runningWorkers++ // 运行中的任务加一

    go func() {
        defer func() {
            p.runningWorkers-- // worker 结束, 运行中的任务减一
        }()

        for {
            select { // 阻塞等待任务、结束信号到来
            case task, ok := <-p.taskC: // 从 channel 中消费任务
                if !ok { // 若是 channel 被关闭, 结束 worker 运行
                    return
                }
                // 执行任务
                task.Handler(task.Params...)
            case <-p.closeC: // 若是收到关闭信号, 结束 worker 运行
                return
            }
        }
    }()
}

上述代码中, runningWorkers 的加减直接使用了自增运算, 可是考虑到启动多个 worker 时, runningWorkers 就会有数据竞争, 因此咱们使用 sync.atomic 包来保证 runningWorkers 的自增操做是原子的。

对 runningWorkers 的操做进行封装:

func (p *Pool) incRunning() { // runningWorkers + 1
    atomic.AddUint64(&p.runningWorkers, 1)
}

func (p *Pool) decRunning() { // runningWorkers - 1
    atomic.AddUint64(&p.runningWorkers, ^uint64(0))
}

func (p *Pool) GetRunningWorkers() uint64 {
    return atomic.LoadUint64(&p.runningWorkers)
}

打铁乘热, 对于 capacity 的操做也考虑数据竞争, 封装 GetCap() 方法:

func (p *Pool) GetCap() uint64 {
    return atomic.LoadUint64(&p.capacity)
}

run() 方法改造:

func (p *Pool) run() {
    p.incRunning()

    go func() {
        defer func() {
            p.decRunning()
        }()

        for {
            select {
            case task, ok := <-p.taskC:
                if !ok {
                    return
                }
                task.Handler(task.Params...)
            case <-p.closeC:
                return
            }
        }
    }()
}

生产任务

新建 Put() 方法用来将任务放入池中:

func (p *Pool) Put(task *Task) {

    if p.GetRunningWorkers() < p.GetCap() { // 若是任务池满, 则再也不建立 worker
        // 建立启动一个 worker
        p.run()
    }
    // 将任务推入队列, 等待消费
    p.taskC <- task
}

任务池安全关闭

当有关闭任务池来节省 goroutine 资源的场景时, 咱们须要有一个关闭任务池的方法。

直接销毁 worker 关闭 channel 并不合适, 由于此时可能还有任务在队列中没有被消费掉。要确保全部任务被安全消费后再销毁掉 worker。

首先, 在关闭任务池时, 须要先关闭掉生产任务的入口。改造 Put() 方法:

var ErrPoolAlreadyClosed = errors.New("pool already closed")

func (p *Pool) Put(task *Task) error {

    if p.state == STOPED { // 若是任务池处于关闭状态, 再 put 任务会返回 ErrPoolAlreadyClosed 错误
        return ErrPoolAlreadyClosed
    }
    
    if p.GetRunningWorkers() < p.GetCap() { 
        p.run()
    }

    p.taskC <- task
    
    return nil
}

在 run() 方法中已经对 closeC 进行了监听, 销毁 worker 只需等待任务被消费完后向 closeC 发出信号。Close() 方法以下:

func (p *Pool) Close() {
    p.state = STOPED // 设置 state 为已中止

    for len(p.taskC) > 0 { // 阻塞等待全部任务被 worker 消费
    }

    p.closeC <- true // 发送销毁 worker 信号
    close(p.taskC) // 关闭任务队列
}

panic handler

每一个 worker 都是一个 goroutine, 若是 goroutine 中产生了 panic, 会致使整个程序崩溃。为了保证程序的安全进行, 任务池须要对每一个 worker 中的 panic 进行 recover 操做, 并提供可订制的 panic handler。

更新任务池定义:

type Pool struct {
    capacity       uint64
    runningWorkers uint64
    state          int64
    taskC          chan *Task
    closeC         chan bool
    PanicHandler   func(interface{})
}

更新 run() 方法:

func (p *Pool) run() {
    p.incRunning()

    go func() {
        defer func() {
            p.decRunning()
            if r := recover(); r != nil { // 恢复 panic
                if p.PanicHandler != nil { // 若是设置了 PanicHandler, 调用
                    p.PanicHandler(r)
                } else { // 默认处理
                    log.Printf("Worker panic: %s\n", r)
                }
            }
        }()

        for {
            select {
            case task, ok := <-p.taskC:
                if !ok {
                    return
                }
                task.Handler(task.Params...)
            case <-p.closeC:
                return
            }
        }
    }()
}

使用

OK, 咱们的任务池就这么简单的写好了, 试试:

func main() {
    // 建立任务池
    pool, err := NewPool(10)
    if err != nil {
        panic(err)
    }

    for i := 0; i < 20; i++ {
        // 任务放入池中
        pool.Put(&Task{
            Handler: func(v ...interface{}) {
                fmt.Println(v)
            },
            Params: []interface{}{i},
        })
    }

    time.Sleep(1e9) // 等待执行
}

详细例子见 mortar/examples

benchmark

做为协程池, 性能和内存占用的指标测试确定是少不了的, 测试数据才是最有说服力的

测试流程

100w 次执行,原子增量操做

测试任务:

var sum int64

func demoTask(v ...interface{}) {
    for i := 0; i < 100; i++ {
        atomic.AddInt64(&sum, 1)
    }
}

测试方法:

var runTimes = 1000000
// 原生 goroutine
func BenchmarkGoroutineSetTimes(b *testing.B) {

    for i := 0; i < runTimes; i++ {
        go demoTask()
    }
}
// 使用协程池
func BenchmarkPutSetTimes(b *testing.B) {
    pool, err := NewPool(20)
    if err != nil {
        b.Error(err)
    }

    ctx := context.Background()
    task := &Task{
        Handler: demoTask,
    }

    for i := 0; i < runTimes; i++ {
        pool.Put(ctx, task)
    }
}

对比结果

模式 操做时间消耗 ns/op 内存分配大小 B/op 内存分配次数 allocs/op
原生 goroutine (100w goroutine) 1596177880 103815552 240022
任务池开启 20 个 worker 20 goroutine) 1378909099 15312 89

使用任务池和原生 goroutine 性能相近(略好于原生)

使用任务池比直接 goroutine 内存分配节省 7000 倍左右, 内存分配次数减小 2700 倍左右

源码地址

该项目的所有源码详见 mortar

参考文章:

【1】线程的 3 种实现方式

相关文章
相关标签/搜索