grpool协程池—控制goroutine数量利器

目的

grpool的目的就是控制执行任务的goroutine数量，避免创建过多的goroutine导致内存占用飙升

简单使用

Simple example

package main

import (
  "fmt"
  "runtime"
  "time"

  "github.com/ivpusic/grpool"
)

func main() {
  // number of workers, and size of job queue
  pool := grpool.NewPool(100, 50)

  // release resources used by pool
  defer pool.Release()

  // submit one or more jobs to pool
  for i := 0; i < 10; i++ {
    count := i

    pool.JobQueue <- func() {
      fmt.Printf("I am worker! Number %d\n", count)
    }
  }

  // dummy wait until jobs are finished
  time.Sleep(1 * time.Second)
}

Example with waiting jobs to finish

package main

import (
  "fmt"
  "runtime"

  "github.com/ivpusic/grpool"
)

func main() {
  // number of workers, and size of job queue
  pool := grpool.NewPool(100, 50)
  defer pool.Release()

  // how many jobs we should wait
  pool.WaitCount(10)

  // submit one or more jobs to pool
  for i := 0; i < 10; i++ {
    count := i

    pool.JobQueue <- func() {
      // say that job is done, so we can know how many jobs are finished
      defer pool.JobDone()

      fmt.Printf("hello %d\n", count)
    }
  }

  // wait until we call JobDone for all jobs
  pool.WaitAll()
}

源码分析

关键API

func NewPool(numWorkers int, jobQueueLen int) *Pool
func (p *Pool) Release()
func (p *Pool) WaitCount(count int)
func (p *Pool) JobDone()
func (p *Pool) WaitAll()

NewPool：创建worker goroutine池，numWorkers代表池中创建的worker数量，jobQueueLen代表可以接收的任务数量如果满了那么会阻塞

Pool.Release：释放池中创建的worker goroutine，注意这里没有释放jobQueue中的job

Pool.WaitCount：设置等待完成的任务数

Pool.JobDone：每次任务完成的时候应该调用这个方法，将等待完成的任务数减1

Pool.WaitAll：阻塞等待所有任务完成，当等待完成的任务数为0时继续往后执行

数据结构

Pool
goroutine池，包含三个属性JobQueue、Dispatcher、wg
JobQueue：任务队列，所有的任务按照先进先出的顺序执行，Dispatcher和Pool共享一个JobQueue
Dispatcher：任务分发器，负责从workerPool中取出空闲的worker、从jobQueue中取出job，将job交给worker执行
wg：用于等待任务完成

Job
要执行的任务，是一个无参无返回值的方法

dispatcher
任务分发器，包含三个属性workerPool、jobQueue、stop
workerPool：worker队列
jobQueue：任务队列
stop：无缓冲channel，用于传输stop信号

worker
执行任务的goroutine，包含三个属性workerPool、jobChannel、stop
workerPool：worker队列，后面讲解源码流程会提到为什么worker要引用worker队列
jobChannel：worker自身的无缓冲channel，用于传输worker要执行的任务
stop：无缓冲channel，用于传输stop信号

关键API源码分析

NewPool

func NewPool(numWorkers int, jobQueueLen int) *Pool {
    jobQueue := make(chan Job, jobQueueLen)
    workerPool := make(chan *worker, numWorkers)

    pool := &Pool{
        JobQueue:   jobQueue,
        dispatcher: newDispatcher(workerPool, jobQueue),
    }

    return pool
}

创建job队列、worker池，然后创建dispatcher任务分发器，最后创建Pool池对象

这里newDispatcher是关键

newDispatcher

func newDispatcher(workerPool chan *worker, jobQueue chan Job) *dispatcher {
    d := &dispatcher{
        workerPool: workerPool,
        jobQueue:   jobQueue,
        stop:       make(chan struct{}),
    }

    for i := 0; i < cap(d.workerPool); i++ {
        worker := newWorker(d.workerPool)
        worker.start()
    }

    go d.dispatch()
    return d
}

创建dispatcher对象，然后创建指定数量的worker并且启动worker，最后开启一个goroutine做任务分发

newWorker

func newWorker(pool chan *worker) *worker {
    return &worker{
        workerPool: pool,
        jobChannel: make(chan Job),
        stop:       make(chan struct{}),
    }
}

创建worker对象，workerPool引用之前创建的worker池对象，jobChannel、stop都创建一个属于自己的无缓冲chan

worker.start

func (w *worker) start() {
    go func() {
        var job Job
        for {
            // worker free, add it to pool
            w.workerPool <- w

            select {
            case job = <-w.jobChannel:
                job()
            case <-w.stop:
                w.stop <- struct{}{}
                return
            }
        }
    }()
}

死循环，将自身放到worker池中，如果自己的job队列中有job那么取出job执行，如果没有job但是stopChan中有stop信号，取出stop信号再往stopChan中发送一个stop信号后结束

这里有几个注意点：
1）如果jobChannel中有任务并且stopChan中也有stop信号，事实上是随机选择一条路执行
2）每次循环都在往workerPool中放入worker自身，那么什么时候取出来呢？
3）worker已经从stopChan中接收了stop信号，为什么还要往stopChan中发送一个stop信号呢？

其中2、3这两个注意点的答案都在接下来的dispatcher.dispatch方法，因此我把worker.start的流程图也放到下面的dispatcher.dispatch方法讲解中

dispatcher.dispatch

func (d *dispatcher) dispatch() {
    for {
        select {
        case job := <-d.jobQueue:
            worker := <-d.workerPool
            worker.jobChannel <- job
        case <-d.stop:
            for i := 0; i < cap(d.workerPool); i++ {
                worker := <-d.workerPool

                worker.stop <- struct{}{}
                <-worker.stop
            }

            d.stop <- struct{}{}
            return
        }
    }
}

死循环，如果job队列中有job，那么从workerPool中获取空闲的worker，如果获取不到则一直阻塞等待，如果获取到那么将job扔到worker自己的job队列；如果job队列中没有job、stopChan中有stop信号，那么给每个worker的stopChan中发送stop信号停止worker，并且在给每个worker发送完stop信号后要再次从worker的stopChan中接收到stop信号才继续停止下一个worker，所有worker停止完成后再往dispatcher自己的stopChan中发送stop信号后结束

这里有几个关键点：

1）如果jobQueue中有任务并且stopChan中也有stop信号，事实上是随机选择一条路执行
2）dispatcher分发任务的时候，是从workerPool中取出一个worker然后将任务分发给这个worker，这里就解答了上面的第二个关注点，也是dispatcher和worker配合处理job的方式
3）停止一个worker后还要再从worker的stopChan中接收stop信号后才继续停止下一个worker，这里和上面的第三个关注点呼应，目的是保证一个worker关闭后再去关闭下一个worker
4）所有worker停止后，为什么dispatcher还要往自身的stopChan中发送一个stop信号后才返回呢？这个问题的答案在下面的Pool.Release方法中会有解答

Pool.Release

func (p *Pool) Release() {
    p.dispatcher.stop <- struct{}{}
    <-p.dispatcher.stop
}

释放workerPool中所有的worker

往dispatcher的stopChan中发送完stop信号后，还要再次从dispatcher的stopChan中接收一个stop信号后才返回，这里和上面的第四个关键点呼应，目的是保证workerPool中所有的worker都释放后，Release方法才返回

Pool.WaitCount

func (p *Pool) WaitCount(count int) {
	p.wg.Add(count)
}

设置等待完成的任务数

整个任务的等待执行就是通过sync.WaitGroup来实现的

Pool.JobDone

func (p *Pool) JobDone() {
	p.wg.Done()
}

每次任务完成的时候应该调用这个方法，将等待完成的任务数减1

Pool.WaitAll

func (p *Pool) WaitAll() {
	p.wg.Wait()
}

阻塞等待所有任务完成，当等待完成的任务数为0时继续往后执行

总结

整个grpool的设计比较简单，关键是搞懂worker.start、dispatcher.dispatch这两个方法是如何配合进行job的分发以及job的执行