ants Code Reading

概述

goroutine 使用简单，但并不意味着没有成本。

ants 是一个高性能且低损耗的 goroutine 池，组件内部会创建一个固定容量的 goroutine 池，并且管理和回收池中的 goroutine， 主要应用场景是允许开发者限制并发程序中的 goroutine 数量，更多的特性介绍请参考 Github 主页^[1]。

示例代码

我们首先通过两个示例程序来直观地感受一下使用 goroutine 池前后的差异。

未使用 goroutine 池

package main

import (
 "log"
 "runtime"
 "sync"
 "sync/atomic"
)

var (
 sum int32
)

func main() {
 var wg sync.WaitGroup

 var i int32
 // 启动 100 个 goroutine
 for i = 1; i <= 100; i++ {
  wg.Add(1)
  go func(n int32) {
   defer wg.Done()
   atomic.AddInt32(&sum, n)
  }(i)
 }

 log.Printf("goroutines numbers = %dn", runtime.NumGoroutine())

 wg.Wait()

 log.Printf("sum = %d, want = %d", sum, 5050)
}

运行代码输出结果如下:

goroutines numbers = 101
sum = 5050, want = 5050

程序运行过程中一共启动了 101 个 goroutine, 其中 1 个是 main goroutine, 其他的 100 个是通过 go 关键字启动的 goroutine 。

使用 goroutine 池

package main

import (
 "log"
 "runtime"
 "sync"
 "sync/atomic"

 "github.com/panjf2000/ants/v2"
)

var (
 sum int32
)

func main() {
 defer ants.Release()

 var wg sync.WaitGroup

 p, err := ants.NewPoolWithFunc(10, func(i interface{}) {
  defer wg.Done()
  atomic.AddInt32(&sum, i.(int32))
 })
 if err != nil {
  log.Fatal(err)
 }
 defer p.Release()

 for i := 1; i <= 100; i++ {
  wg.Add(1)
  _ = p.Invoke(int32(i))
 }

 log.Printf("goroutines numbers = %dn", runtime.NumGoroutine())

 wg.Wait()

 log.Printf("sum = %d, want = %d", sum, 5050)
}

运行代码输出结果如下:

goroutines numbers = 15
sum = 5050, want = 5050

程序运行过程中一共启动了 15 个 goroutine, 其中 1 个是 main goroutine, 还有 10 个是通过 ants 组件限制数量的 goroutine, 剩余的 4 个是 ants 组件用于管理调度的内部 goroutine 。

源码解析

通过上面两个示例程序的结果比较，ants 确实能够有效降低和控制 goroutine 的数量，这个功能可以有效解决海量 goroutine 场景下的内存暴涨问题。 性能方面，官方给出的基准测试结果是: ants 的吞吐性能相较于原生 goroutine 可以保持在 2~6 倍的性能压制，而内存消耗则可以达到 10~20 倍的节省优势， 感兴趣的读者可以查看 这里的基准测试^[2]。

数据结构

UML

读者在阅读下面的源代码分析时，可以对照着结构图进行校验。

通用池和专用池

ants 中内置了两种不同模式的 goroutine pool:

1. Pool

ants.Pool 表示通用的 goroutine pool, 也就是每个 goroutine 中可以是不同的任务执行函数。

2. PoolWithFunc

ants.PoolWithFunc 表示专用的 goroutine pool, 也就是每个 goroutine 中是相同的任务执行函数，但是参数可以不同 (上面的示例代码中使用就是这种模式)。

本文以 ants.Pool 模式下的源代码来进行分析。

池对象

Pool 表示组件的核心对象，我们挑选几个关键属性字段来进行说明。

type Pool struct {
 // 池的容量
 // 也就是负责执行 worker 的数量上限
 capacity int32

 // 当前正在执行的 worker 数量
 // 也就是 goroutine 的数量 (每个 worker 启动一个 goroutine)
 running int32

 // 标准库的锁接口
 // 组件内部使用了自旋锁作为具体的实现
 lock sync.Locker

 // 存储闲置的 worker
 workers workerQueue

 // 池的状态标识 (标识池是否已关闭)
 state int32

 // 获取闲置 worker 时使用的信号量
 cond *sync.Cond

 // 池的配置选项对象
 options *Options
}

配置项

Options 对象表示池对象的属性配置，可以在创建池对象时通过 FUNCTIONAL OPTIONS 模式注入自定义的属性值。

type Options struct {
 // 任务的过期时间
 // 空闲 worker 的最后运行时间与当前时间之差
 // 如果大于这个值，表示该 worker 已过期
 // 定时清理任务会清理过期的 worker
 ExpiryDuration time.Duration

 // 初始化池的时候是否预分配 worker 队列
 PreAlloc bool

 // 调用 pool.Submit 方法阻塞时，最多可以阻塞的 worker 数量
 // 默认值为 0，表示不限制
 MaxBlockingTasks int

 // 为 true 时，pool.Submit 方法永远不会阻塞，也就是 MaxBlockingTasks 属性不起作用
 // 如果 pool.Submit 方法因阻塞而无法完成时，返回 ErrPoolOverload 错误
 Nonblocking bool

 // 任务 goroutine 触发 panic 时的错误捕获及处理函数
 PanicHandler func(interface{})

 // 自定义日志对象
 // 默认使用标准库中的 log 对象
 Logger Logger

 // 任务 goroutine 是否常驻内存
 DisablePurge bool
}

func WithExpiryDuration(options Options) Option {}
func WithPreAlloc(options Options) Option {}
...

worker 接口

worker 接口是任务执行者的抽象表示，简单地理解就是一个专门干活的 goroutine, 该接口提供了 5 个操作语义:

ants.Pool 池对应的 worker 接口实现对象是 goWorker。

type goWorker struct {
 // 对应的池对象
 pool *Pool

 // 需要执行的任务函数
 task chan func()

 // 最后的任务执行时间
 lastUsed time.Time
}

下面是 goWorker 对象对于接口的具体实现:

func (w *goWorker) run() {
 // 通知 goroutine 池: 运行 goroutine 数量 + 1
 w.pool.addRunning(1)

 // 启动一个 goroutine
 go func() {
  defer func() {
   // 通知 goroutine 池: 运行 goroutine 数量 - 1
   w.pool.addRunning(-1)
   // 将当前 goWorker 对象放入对象池
   // 可以提升后续获取 goWorker 对象操作的性能
   w.pool.workerCache.Put(w)
   if p := recover(); p != nil {
    // panic 处理
    if ph := w.pool.options.PanicHandler; ph != nil {
     ph(p)
    } else {
     w.pool.options.Logger.Printf("worker exits from panic: %vn%sn", p, debug.Stack())
    }
   }
   // 发送闲置 goWorker 对象信号
   w.pool.cond.Signal()
  }()

  for f := range w.task {
   if f == nil {
    // 如果执行函数为 nil (这里的 nil 充当了 goroutine 的退出信号)
    // 直接退出即可
    return
   }
   // 执行具体的函数
   f()
   // 将当前 goWorker 对象放入 worker 队列
   // 最大程度地复用当前 goWorker 对象，同时避免 goroutine 的上下文切换带来的开销
   if ok := w.pool.revertWorker(w); !ok {
    // 如果放入队列失败，可能是由于:
    // 1. Pool 池对象已关闭
    // 2. Pool 池对象当前运行的 goroutine 数量已到达限制
    // 此时直接退出当前 goroutine 即可
    return
   }
  }
 }()
}

// 向任务 channel 发送一个 nil, 充当了 goroutine 的退出信号
func (w *goWorker) finish() {
 w.task <- nil
}

// 返回 worker 最后的任务执行时间
func (w *goWorker) lastUsedTime() time.Time {
 return w.lastUsed
}

// 向 worker 提交一个新的任务
func (w *goWorker) inputFunc(fn func()) {
 w.task <- fn
}

// 因为 Pool.Submit 提交任务的原型为 func (p *Pool) Submit(task func()) error {}
// 所以 goWorker 对象中的任务函数没有参数值
func (w *goWorker) inputParam(interface{}) {
 panic("unreachable")
}

worker 队列接口

workerQueue 接口表示存放 worker 的队列容器接口，和一般的容器类接口相似, 该接口提供了 6 个操作语义:

组件内部给出了两种队列实现方案，分别是基于 Stack (栈) 和 Ring Queue (环形队列)，两种方案的内部实现和常规数据结构算法实现基本一致，这里不再赘述。

方法

池的初始化

func NewPool(size int, options ...Option) (*Pool, error) {
 // Pool 的配置项加载和默认值处理
 opts := loadOptions(options...)

 ...

 // 初始化 Pool 对象
 p := &Pool{
  capacity: int32(size),
  lock:     syncx.NewSpinLock(),
  options:  opts,
 }

 // 注册 worker 对象池
 p.workerCache.New = func() interface{} {
  return &goWorker{
   pool: p,
   task: make(chan func(), workerChanCap),
  }
 }

 // 根据配置决定是否预分配 worker 队列
 if p.options.PreAlloc {
    // 如果使用预分配
    // 使用 Ring Queue 作为队列的数据结构
  p.workers = newWorkerArray(queueTypeLoopQueue, size)
 } else {
  // 如果不使用预分配
  // 使用 Stack 作为队列的数据结构
  p.workers = newWorkerArray(queueTypeStack, 0)
 }

 ...

 // 启动一个 goroutine, 清除过期 worker
 p.goPurge()

 return p, nil
}

池的初始化操作完成后，接下来就是等待调用方提交具体的任务函数。

提交任务到 Pool

调用方提交任务函数到池中的时候，会尝试获取一个可用的 worker 来执行任务函数。

func (p *Pool) Submit(task func()) error {
 if p.IsClosed() {
  // 如果 Pool 已关闭, 直接返回
  return ErrPoolClosed
 }
 if w := p.retrieveWorker(); w != nil {
  // 获取到了可用的 worker
  // 直接将参数任务提交给 worker 即可
  w.inputFunc(task)
  return nil
 }
 // 没有获取到可用的 worker
 return ErrPoolOverload
}

获取可用的 worker

retrieveWorker 方法会根据配置项，采用不同的策略获取可用的 worker 。

func (p *Pool) retrieveWorker() (w worker) {
 spawnWorker := func() {
  // 从对象池中获取一个 worker
  w = p.workerCache.Get().(*goWorker)
  // 获取到 worker 之后直接运行
  w.run()
 }

 p.lock.Lock()

 // 首先尝试从队列获取 worker
 w = p.workers.detach()
 if w != nil {
  // 从队列中获取到了 worker
  p.lock.Unlock()
 } else if capacity := p.Cap(); capacity == -1 || capacity > p.Running() {
  // 如果队列为空并且满足以下两个条件任意一个:
  //    1. Pool 没有容量限制
  //    2. 当前 worker 数量没有达到 Pool 的容量限制
  // 直接创建一个新的 worker
  p.lock.Unlock()
  spawnWorker()
 } else {
  // 如果上述条件都不满足
  // 只能等待其他执行完的 worker 归还到队列中，然后再获取

  if p.options.Nonblocking {
   // 如果 Pool 的任务提交模式是非阻塞的
   // 直接返回
   p.lock.Unlock()
   return
  }
 retry:
  if p.options.MaxBlockingTasks != 0 && p.Waiting() >= p.options.MaxBlockingTasks {
   // 如果等待 worker 的 goroutine 数量超过限制
   // 直接返回
   p.lock.Unlock()
   return
  }

  // 等待 worker 的 goroutine 数量 + 1
  p.addWaiting(1)
  // 等待可用的 worker 信号
  p.cond.Wait()
  // 等待 worker 的 goroutine 数量 - 1
  p.addWaiting(-1)

  if w = p.workers.detach(); w == nil {
   // 如果接收到了信号，但是依然没有从队列中获取到 worker
   // 说明队列中的 worker 被其他 goroutine 抢先一步获取到了
   if p.Free() > 0 {
    // 此时如果还有没有达到 Pool 的容量限制
    // 直接创建一个新的 worker 并返回
    p.lock.Unlock()
    spawnWorker()
    return
   }
   goto retry
  }
  p.lock.Unlock()
 }
 return
}

定时清除过期 worker

func (p *Pool) purgeStaleWorkers(ctx context.Context) {
 // 根据配置项创建定时器
 ticker := time.NewTicker(p.options.ExpiryDuration)

 defer func() {
  // 终止定时器
  ticker.Stop()
  atomic.StoreInt32(&p.purgeDone, 1)
 }()

 for {
  select {
  case <-ctx.Done():
   // 退出信号
   return
  case <-ticker.C:
   // 执行清理操作
  }

  var isDormant bool
  p.lock.Lock()
  // 检测过期的 worker
  staleWorkers := p.workers.refresh(p.options.ExpiryDuration)
  n := p.Running()
  isDormant = n == 0 || n == len(staleWorkers)
  p.lock.Unlock()

  // 停止过期 worker 的执行
  for i := range staleWorkers {
   staleWorkers[i].finish()
   staleWorkers[i] = nil
  }

  // 如果当前没有运行的 worker, 所有的 worker 都已过期
  // 但是与此同时还有正在提交任务的 goroutine 在等待可用的 worker
  // 直接发送信号通知
  if isDormant && p.Waiting() > 0 {
   p.cond.Broadcast()
  }
 }
}

辅助方法

// 返回阻塞在提交任务的 goroutine
func (p *Pool) Waiting() int {}

// 返回 Pool 的容量
func (p *Pool) Cap() int {}

// 返回 Pool 是否已关闭
func (p *Pool) IsClosed() bool {}

...

小结

本文从源代码的角度，剖析了知名的开源组件 ants, 着重分析了 Pool, worker, worker 队列 三个核心对象的设计与实现， 组件代码逻辑可以参照官方的流程图。

组件本身存在的问题

可以参考 TiDB 大佬的文章^[3] 中的 Section: “跟 goroutine pool 或者 worker pool 有什么区别？”。

Reference

• ants^[4]
• Rubbing control theory on the Go scheduler^[5]

链接