Golang 并发协程池实现分析（三）

// Pool accepts the tasks from client, it limits the total of goroutines to a given number by recycling goroutines.
type Pool struct {
	// capacity of the pool, a negative value means that the capacity of pool is limitless, an infinite pool is used to
	// avoid potential issue of endless blocking caused by nested usage of a pool: submitting a task to pool
	// which submits a new task to the same pool.
	capacity int32  //ants 最多能创建的 goroutine 数量（协程池数量）

	// running is the number of the currently running goroutines.
	running int32   // 已经创建的 worker goroutine 的数量

	// lock for protecting the worker queue.
	lock sync.Locker    //ants 自己实现了一个自旋锁。用于同步并发操作（用于保护 workerArray）

	// workers is a slice that store the available workers.
	workers workerArray // 存放一组 worker 对象，即一组 goroutine，workerArray 只是一个 interface，存放 goWorker 对象的容器（见下文分析）

	// state is used to notice the pool to closed itself.
	state int32     // 记录池子当前的状态，是否已关闭（CLOSED）

	// cond for waiting to get a idle worker.
	cond *sync.Cond // 条件变量，处理任务等待和唤醒

	// workerCache speeds up the obtainment of a usable worker in function:retrieveWorker.
	workerCache sync.Pool   // 使用 sync.Pool 对象池管理和创建 worker 对象，提升性能

	// blockingNum is the number of the goroutines already been blocked on pool.Submit, protected by pool.lock
	blockingNum int     // 当前阻塞等待的任务数量（阻塞协程数量，指提交任务的协程）

	options *Options	 // 协程池的配置，包括过期时间、是否支持预分配、最大阻塞数量、panic 处理、日志等
}

// workerArray，一个 workerArray 代表一个 worker 池（协程池）
type workerArray interface {
	len() int
	isEmpty() bool
	insert(worker *goWorker) error
	detach() *goWorker
	retrieveExpiry(duration time.Duration) []*goWorker
	reset()
}

// goWorker is the actual executor who runs the tasks,
// it starts a goroutine that accepts tasks and
// performs function calls.
type goWorker struct {
	// pool who owns this worker.	 // 用于记录当前 worker 属于哪一个协程池
	pool *Pool

	// task is a job should be done.  // 接受任务的 chan，多核环境 chan 大小为 1，单核环境 chan 大小为 0(借鉴 fastHttp 的实现)
	task chan func()

	// recycleTime will be updated when putting a worker back into queue.	  // worker 进入队列的时间
	recycleTime time.Time
}

// NewPool generates an instance of ants pool.
func NewPool(size int, options ...Option) (*Pool, error) {
	opts := loadOptions(options...)

	if size <= 0 {
		size = -1
	}

	// 设置回收协程定时器触发时间
	if expiry := opts.ExpiryDuration; expiry < 0 {
		return nil, ErrInvalidPoolExpiry
	} else if expiry == 0 {
		opts.ExpiryDuration = DefaultCleanIntervalTime
	}

	if opts.Logger == nil {
		opts.Logger = defaultLogger
	}
    // 创建 Pool 对象
	p := &Pool{
		capacity: int32(size),
		lock:     internal.NewSpinLock(),
		options:  opts,
	}
	 // 指定 sync.Pool 创建 worker 的方法
	p.workerCache.New = func() interface{} {
		return &goWorker{
			pool: p,
			task: make(chan func(), workerChanCap),
		}
	}
	// 根据预分配标志，使用不同的 workerArray 的实现方式
	if p.options.PreAlloc {
		if size == -1 {
			return nil, ErrInvalidPreAllocSize
		}
		// 预先分配固定 Size 的池子
		p.workers = newWorkerArray(loopQueueType, size)
	} else {
		// 初始化不创建，运行时再创建
		p.workers = newWorkerArray(stackType, 0)
	}

	p.cond = sync.NewCond(p.lock)

	// Start a goroutine to clean up expired workers periodically.
    // 开启一个协程周期清理过期的 worker
	go p.purgePeriodically()

	return p, nil
}

func (p *Pool) purgePeriodically() {
  heartbeat := time.NewTicker(p.options.ExpiryDuration)
  defer heartbeat.Stop()

  for range heartbeat.C {
    if p.IsClosed() {
		// 如果 Pool 被主动关闭，直接退出 goroutine
		// 如果协程池已经被关闭，就退出清理的定时任务
      	break
    }

	// 从 workers 中获取过期的 worker
    p.lock.Lock()
    expiredWorkers := p.workers.retrieveExpiry(p.options.ExpiryDuration)
    p.lock.Unlock()

	// 执行清理过期的 worker
    for i := range expiredWorkers {
	  // 向 worker 的 taskChan 发送终止信号；当 worker 接收到 nil 的任务时，会进入 workerCache 等待 GC
      expiredWorkers[i].task <- nil
	  // 清空 worker 的引用，方便 GC
      expiredWorkers[i] = nil
    }

	// 唤醒获取 worker 的协程
    if p.Running() == 0 {
		// 唤醒阻塞在 p.cond.Wait() 上的 goroutine
      p.cond.Broadcast()
    }
  }
}

func (wq *loopQueue) retrieveExpiry(duration time.Duration) []*goWorker {
	expiryTime := time.Now().Add(-duration)
	index := wq.binarySearch(expiryTime)
	if index == -1 {
		return nil
	}
	wq.expiry = wq.expiry[:0]

	if wq.head <= index {
		wq.expiry = append(wq.expiry, wq.items[wq.head:index+1]...)
		for i := wq.head; i < index+1; i++ {
			wq.items[i] = nil
		}
	} else {
		wq.expiry = append(wq.expiry, wq.items[0:index+1]...)
		wq.expiry = append(wq.expiry, wq.items[wq.head:]...)
		for i := 0; i < index+1; i++ {
			wq.items[i] = nil
		}
		for i := wq.head; i < wq.size; i++ {
			wq.items[i] = nil
		}
	}
	head := (index + 1) % wq.size
	wq.head = head
	if len(wq.expiry) > 0 {
		wq.isFull = false
	}

	// 返回所有过期的 worker 列表
	return wq.expiry
}

func (p *Pool) Submit(task func()) error {
	// 判断 pool 是否关闭
  if p.IsClosed() {
    return ErrPoolClosed
  }
  var w *goWorker
  //retrieveWorker 方法获取空闲的 worker
  if w = p.retrieveWorker(); w == nil {
    return ErrPoolOverload
  }

  // 将待执行的任务 send 到 goWorker 的 任务 channel 中
  w.task <- task
  return nil
}

func (p *Pool) retrieveWorker() (w *goWorker) {
		//spawnWorker 方法
		spawnWorker := func() {
			// 实例化 worker
			w = p.workerCache.Get().(*goWorker)
			// 启动 worker
			w.run()
		}
        p.lock.Lock()

        //1. 从 workers 中取出一个 goWorker（ 优先从 workerArray 中获取 worker）
        //p.workers 是 loopQueue 或者 workerStack 对象，它们都实现了 detach() 方法
        w = p.workers.detach()
        if w != nil {
			//SUCC，有空闲 goroutine，直接返回
            p.lock.Unlock()
        } else if capacity := p.Cap(); capacity == -1 || capacity> p.Running() {
			// 池容量还没用完（即容量大于正在工作的 goWorker 数量），则调用 spawnWorker() 新建一个 goWorker，执行其 run() 方法，直接返回
			// 2. 可扩容，从 workerCache 中获取
			p.lock.Unlock()
			spawnWorker()
        } else {
			if p.options.Nonblocking {
				    // 3. 若设置了非阻塞选项，直接返回 nil
					p.lock.Unlock()
					return
			}
        RETRY:
                if p.options.MaxBlockingTasks != 0 && p.blockingNum >= p.options.MaxBlockingTasks {
					// 如果设置了最大阻塞队列长度限制，并且当前阻塞等待的任务数量已经达到这个上限，直接返回 nil
                        p.lock.Unlock()
                        return
                }
				// 未超过最大阻塞队列长度限制，阻塞等待数量 +1，调用 p.cond.Wait() 阻塞等待（此操作会被 p.cond.Signal()/p.cond.Broadcast() 这两个方法唤醒）
                p.blockingNum++
				// 调用 p.cond.Wait() 等待
				// 注意！这里会阻塞，通过 p.cond.Signal() 方法会唤醒这里的逻辑
                p.cond.Wait()
				// 异步：被唤醒后， 阻塞的任务数减 1
                p.blockingNum--
                var nw int
                if nw = p.Running(); nw == 0 {
					// 判断当前 goWorker 的数量（goroutine 数量）是否等于 0，若为 0，很有可能 Pool 刚刚执行了 Release() ，及 Pool 被主动关闭了
                        p.lock.Unlock()
                        if !p.IsClosed() {
							//Pool 未被关闭，说明可以创建新的 goWorker，调用 spawnWorker() 创建一个新的 goWorker 并执行其 run() 方法
                                spawnWorker()
                        }
						// 如果 Pool 被关闭了，p.IsClosed 为 true，那么直接返回
                        return
                }

				// 如果当前 goWorker 数量不为 0，则调用 p.workers.detach() 方法尝试取出一个空闲的 goWorker
                if w = p.workers.detach(); w == nil {
					// 若取出失败（有可能发生），因为可能同时有多个 goroutine 在等待，唤醒的时候只有部分 goroutine 能获取到 goWorkerr
                        if nw < capacity {
							// 从 Pool 中取出失败，检查容量是否还有额度，直接创建新的 goWorke
                                p.lock.Unlock()
                                spawnWorker()
                                return
                        }
						// 如果没有取到 goWorker，那么就返回到 p.cond.Wait() 继续阻塞等待好了
						// TODO
                        goto RETRY
                }

                p.lock.Unlock()
        }
    return
}

func (p *Pool) Tune(size int) {
  if capacity := p.Cap(); capacity == -1 || size <= 0 || size == capacity || p.options.PreAlloc {
    return
  }
  atomic.StoreInt32(&p.capacity, int32(size))
}

func (p *Pool) Release() {
	// 设置 pool 的关闭状态
	atomic.StoreInt32(&p.state, CLOSED)
	p.lock.Lock()

	// 关闭 workerArray 中的每个 goroutine
	p.workers.reset()
	p.lock.Unlock()
	// 注意！为了防止有 goroutine 阻塞在 p.cond.Wait() 上，执行一次 p.cond.Broadcast() 唤醒这些阻塞的任务！
	p.cond.Broadcast()
}

// workerStack.reset
func (wq *workerStack) reset() {
	for i := 0; i <wq.len(); i++ {
		// 对每个 worker，发送 nil 到 task 通道从而结束 goroutine
		wq.items[i].task <- nil
		wq.items[i] = nil
  	}
	//free wq.items，好习惯
  	wq.items = wq.items[:0]
}

func (p *Pool) Reboot() {
	// 打开 pool 的状态
	if atomic.CompareAndSwapInt32(&p.state, CLOSED, OPENED) {
		// 由于 p.purgePeriodically() 在 p.Release() 之后检测到池关闭就直接退出了，需要重新开启 purgePeriodically 方法
		go p.purgePeriodically()
	}
}

func (p *Pool) revertWorker(worker *goWorker) bool {
  if capacity := p.Cap(); (capacity> 0 && p.Running() > capacity) || p.IsClosed() {
	  // 不满足放回条件，返回 false
    return false
  }
  worker.recycleTime = time.Now()	// 重置空闲计时器，用于判定过期
  p.lock.Lock()

  if p.IsClosed() {
	  // 如果 Pool 被关闭，那么退出
    p.lock.Unlock()
    return false
  }

	// 调用 workerArray 的 insert 方法，放回 pool
  err := p.workers.insert(worker)
  if err != nil {
    p.lock.Unlock()
    return false
  }

	// 放回成功，通过 p.cond.Signal() 唤醒一个可能阻塞的 goroutine
  p.cond.Signal()
  p.lock.Unlock()
  return true
}

// goWorker is the actual executor who runs the tasks,
// it starts a goroutine that accepts tasks and
// performs function calls.
type goWorker struct {
	// pool who owns this worker.
	pool *Pool	// 指向属主池

	// task is a job should be done.
	task chan func()	// 非常重要！任务通道，通过这个通道将类型为 func () 的函数作为任务发送给 goWorker 执行

	// recycleTime will be updated when putting a worker back into queue.
	recycleTime time.Time	// 此字段记录 goroutine 放回池中的时间（即什么时候开始空闲），注意：此字段只在非 PreAlloc 模式下才起效
}

// run starts a goroutine to repeat the process
// that performs the function calls.
func (w *goWorker) run() {
	// 任务 + 1
	w.pool.incRunning()		// 正在运行的协程数 + 1
	go func() {
		// 异常处理！
		defer func() {
			// 任务执行失败，goroutine 结束，运行数量减 1
			w.pool.decRunning()		// 正在运行的协程数 - 1
			//goWorker 对象可以重复利用，利用 sync.Pool 回收，将 goWorker 对象放回 sync.Pool 池中
			w.pool.workerCache.Put(w)	// 将 worker 放入 workerCache 中等待 GC
			if p := recover(); p != nil {
				if ph := w.pool.options.PanicHandler; ph != nil {
					// 自定义 panic_handler
					ph(p)
				} else {
					w.pool.options.Logger.Printf("worker exits from a panic: %v\n", p)
					var buf [4096]byte
					n := runtime.Stack(buf[:], false)
					w.pool.options.Logger.Printf("worker exits from panic: %s\n", string(buf[:n]))
				}
			}
			// Call Signal() here in case there are goroutines waiting for available workers.
			// 这里有意思，调用 w.pool.cond.Signal() 通知现在有空闲的 goWorker 了
			// 因为实际运行的 goWorker 数量由于 panic 少了一个，而池中可能有其他任务在等待处理
			w.pool.cond.Signal()	 // 唤醒获取 worker 的协程
		}()

		//CORE!
		// 不断消费 taskChan 中的任务
		for f := range w.task {
			//for loop here......
			if f == nil {	 // 接收到终止信号 nil 时，退出循环进入清理环节
				// 外部主动通知本协程关闭
				return
			}
			f()	// 执行任务

			// 调用池的 revertWorker() 方法放回 Pool，返回 false，goroutine 退出；返回 true，说明此 goroutine 被正确的放回了 pool，阻塞在 range 上，等待下一次 task 被塞入任务！
			if ok := w.pool.revertWorker(w); !ok {	// 将 worker 放回到 workerArray 中
				// 如果放回操作失败，则会调用 return，这会让 goroutine 运行结束，防止 goroutine 泄漏
				return
			}
		}
	}()
}

type workerArray interface {
  len() int
  isEmpty() bool
  insert(worker *goWorker) error    // 任务执行结束后，将相应的 worker （goroutine）放回 workerArray 中
  detach() *goWorker                // 从 workerArray 中取出一个 worker（goroutine）
  retrieveExpiry(duration time.Duration) []*goWorker    // 取出所有的过期 worker
  reset()
}

type loopQueue struct {
	items  []*goWorker
	expiry []*goWorker
	head   int
	tail   int
	size   int
	isFull bool
}

type workerStack struct {
  items  []*goWorker       // 当前可用（空闲的） worker
  expiry []*goWorker        // 过期的 worker
  size   int
}

// 非预先创建，采用运行时创建的逻辑，insert 始终是插入到末尾位置
func (wq *workerStack) insert(worker *goWorker) error {
  wq.items = append(wq.items, worker)
  return nil
}

func (wq *workerStack) detach() *goWorker {
  l := wq.len()
  if l == 0 {
    return nil
  }

  w := wq.items[l-1]
  wq.items[l-1] = nil // avoid memory leaks
  wq.items = wq.items[:l-1]

  return w
}

// 二分查找的是最近过期的 worker，即将过期的 worker 的前一个位置，在此位置之前的 worker 已经全部过期了
func (wq *workerStack) binarySearch(l, r int, expiryTime time.Time) int {
  var mid int
  for l <= r {
    mid = (l + r) / 2
    //expiryTime < recycleTime
    if expiryTime.Before(wq.items[mid].recycleTime) {
      r = mid - 1
    } else {
      l = mid + 1
    }
  }
  return r
}

func (wq *workerStack) retrieveExpiry(duration time.Duration) []*goWorker {
  n := wq.len()
  if n == 0 {
    return nil
  }

  // 获取超时的底限刻度
  expiryTime := time.Now().Add(-duration)
  // 采用二分法查找 index
  index := wq.binarySearch(0, n-1, expiryTime)

  //reset wq.expiry
  wq.expiry = wq.expiry[:0]
  if index != -1 {
    wq.expiry = append(wq.expiry, wq.items[:index+1]...)    // 在 0:index+1 区间的 goroutine 都需要被回收

    // 将 wq.items[index+1:] 部分的 goroutine 重新放回 wq.items（这些元素都是未过期的）
    m := copy(wq.items, wq.items[index+1:])
    for i := m; i < n; i++ {
        //reset origin zone
      wq.items[i] = nil
    }
    // 重新生成 wq.items
    wq.items = wq.items[:m]
  }
  return wq.expiry
}

type Options struct {
	ExpiryDuration time.Duration	// 过期时间，表示 goroutine 空闲多长时间之后会被 ants 池回收
	PreAlloc bool	// 是否预分配
	MaxBlockingTasks int	// 最大阻塞任务数量，即池中 goroutine 数量已到池容量，且所有 goroutine 都处于繁忙状态，这时到来的任务会在阻塞列表等待。阻塞的任务数量达到这个值后，后续任务提交直接返回失败
	Nonblocking bool	// 阻塞开关，提交任务时，如果 ants 池中 goroutine 已到上限且全部繁忙，阻塞的池会将任务添加的阻塞列表等待（受限于阻塞列表长度）。非阻塞下直接返回失败
	PanicHandler func(interface{})	//panic 钩子方法，遇到 panic 会调用这里设置的处理函数
	Logger Logger
}

// retrieveWorker returns an available worker to run the tasks.
func (p *Pool) retrieveWorker() (w *goWorker) {
	//...
	p.lock.Lock()

	w = p.workers.detach()
	if w != nil { // first try to fetch the worker from the queue
		p.lock.Unlock()
	} else if capacity := p.Cap(); capacity == -1 || capacity> p.Running() {
		// if the worker queue is empty and we don't run out of the pool capacity,
		// then just spawn a new worker goroutine.
		p.lock.Unlock()
		spawnWorker()
	} else { // otherwise, we'll have to keep them blocked and wait for at least one worker to be put back into pool.
		if p.options.Nonblocking {
			p.lock.Unlock()
			return
		}
	retry:
		if p.options.MaxBlockingTasks != 0 && p.blockingNum >= p.options.MaxBlockingTasks {
			p.lock.Unlock()
			return
		}
		p.blockingNum++
		p.cond.Wait() // block and wait for an available worker
		p.blockingNum--
		var nw int
		if nw = p.Running(); nw == 0 { // awakened by the scavenger
			p.lock.Unlock()
			if !p.IsClosed() {
				spawnWorker()
			}
			return
		}
		if w = p.workers.detach(); w == nil {
			if nw < capacity {
				p.lock.Unlock()
				spawnWorker()
				return
			}
			goto retry
		}

		p.lock.Unlock()
	}
	return
}

p.cond = sync.NewCond(p.lock)

func (c *Cond) Wait() {
	c.checker.check()
	t := runtime_notifyListAdd(&c.notify)
	// 先解锁
	c.L.Unlock()
	// 等待唤醒
	runtime_notifyListWait(&c.notify, t)
	// 加锁
	c.L.Lock()
}

// Signal wakes one goroutine waiting on c, if there is any.
//
// It is allowed but not required for the caller to hold c.L
// during the call.
func (c *Cond) Signal() {
	c.checker.check()
	runtime_notifyListNotifyOne(&c.notify)
}

// Broadcast wakes all goroutines waiting on c.
//
// It is allowed but not required for the caller to hold c.L
// during the call.
func (c *Cond) Broadcast() {
	c.checker.check()
	runtime_notifyListNotifyAll(&c.notify)
}

func (wq *workerStack) detach() *goWorker {
  l := wq.len()
  if l == 0 {
    return nil
  }

  w := wq.items[l-1]
  wq.items[l-1] = nil // avoid memory leaks
  wq.items = wq.items[:l-1]

  return w
}