Why Pool

go自從出生就身帶“高併發”的標籤，其併發程式設計就是由groutine實現的，因其消耗資源低，效能高效，開發成本低的特性而被廣泛應用到各種場景，例如服務端開發中使用的HTTP服務,在golang net/http包中,每一個被監聽到的tcp連結都是由一個groutine去完成處理其上下文的，由此使得其擁有極其優秀的併發量吞吐量

for {
    // 監聽tcp
    rw,e := l.Accept()
    if e != nil {
      .......
    }
    tempDelay = 0
    c := srv.newConn(rw)
    c.setState(c.rwc,StateNew) // before Serve can return
    // 啟動協程處理上下文
    go c.serve(ctx)
}

雖然建立一個groutine佔用的記憶體極小(大約2KB左右，執行緒通常2M左右)，但是在實際生產環境無限制的開啟協程顯然是不科學的，比如上圖的邏輯，如果來幾千萬個請求就會開啟幾千萬個groutine,當沒有更多記憶體可用時,go的排程器就會阻塞groutine最終導致記憶體溢位乃至嚴重的崩潰，所以本文將通過實現一個簡單的協程池，以及剖析幾個開源的協程池原始碼來探討一下對groutine的併發控制以及多路複用的設計和實現。

一個簡單的協程池

過年前做過一波小需求，是將主播管理系統中資訊不完整的主播找出來然後再到其相對應的直播平臺爬取完整資訊並補全，當時考慮到每一個主播的資料都要訪問一次直播平臺所以就用應對每一個主播開啟一個groutine去抓取資料，雖然這個業務量還遠遠遠遠達不到能造成groutine效能瓶頸的地步，但是心裡總是不舒服，於是放假回來後將其優化成從協程池中控制groutine數量再開啟爬蟲進行資料抓取。思路其實非常簡單，用一個channel當做任務佇列，初始化groutine池時確定好併發量，然後以設定好的併發量開啟groutine同時讀取channel中的任務並執行,模型如下圖

實現

type SimplePool struct {
  wg  sync.WaitGroup
  work chan func() //任務佇列
}

func NewSimplePoll(workers int) *SimplePool {
  p := &SimplePool{
    wg:  sync.WaitGroup{},work: make(chan func()),}
  p.wg.Add(workers)
  //根據指定的併發量去讀取管道並執行
  for i := 0; i < workers; i++ {
    go func() {
      defer func() {
        // 捕獲異常 防止waitGroup阻塞
        if err := recover(); err != nil {
          fmt.Println(err)
          p.wg.Done()
        }
      }()
      // 從workChannel中取出任務執行
      for fn := range p.work {
        fn()
      }
      p.wg.Done()
    }()
  }
  return p
}
// 新增任務
func (p *SimplePool) Add(fn func()) {
  p.work <- fn
}

// 執行
func (p *SimplePool) Run() {
  close(p.work)
  p.wg.Wait()
}

測試

測試設定為在併發數量為20的協程池中併發抓取一百個人的資訊,因為程式碼包含較多業務邏輯所以sleep 1秒模擬爬蟲過程，理論上執行時間為5秒

func TestSimplePool(t *testing.T) {
  p := NewSimplePoll(20)
  for i := 0; i < 100; i++ {
    p.Add(parseTask(i))
  }
  p.Run()
}

func parseTask(i int) func() {
  return func() {
    // 模擬抓取資料的過程
    time.Sleep(time.Second * 1)
    fmt.Println("finish parse ",i)
  }
}

這樣一來最簡單的一個groutine池就完成了

go-playground/pool

上面的groutine池雖然簡單，但是對於每一個併發任務的狀態，pool的狀態缺少控制，所以又去看了一下go-playground/pool的原始碼實現，先從每一個需要執行的任務入手，該庫中對併發單元做了如下的結構體，可以看到除工作單元的值，錯誤，執行函式等，還用了三個分別表示，取消，取消中，寫 的三個併發安全的原子操作值來標識其執行狀態。

// 需要加入pool 中執行的任務
type WorkFunc func(wu WorkUnit) (interface{},error)

// 工作單元
type workUnit struct {
  value   interface{}  // 任務結果 
  err    error     // 任務的報錯
  done    chan struct{} // 通知任務完成
  fn     WorkFunc  
  cancelled atomic.Value  // 任務是否被取消
  cancelling atomic.Value  // 是否正在取消任務
  writing  atomic.Value  // 任務是否正在執行
}

接下來看Pool的結構

type limitedPool struct {
  workers uint      // 併發量 
  work  chan *workUnit // 任務channel
  cancel chan struct{}  // 用於通知結束的channel
  closed bool      // 是否關閉
  m    sync.RWMutex  // 讀寫鎖，主要用來保證 closed值的併發安全
}

初始化groutine池,以及啟動設定好數量的groutine

// 初始化pool,設定併發量
func NewLimited(workers uint) Pool {
  if workers == 0 {
    panic("invalid workers '0'")
  }
  p := &limitedPool{
    workers: workers,}
  p.initialize()
  return p
}

func (p *limitedPool) initialize() {
  p.work = make(chan *workUnit,p.workers*2)
  p.cancel = make(chan struct{})
  p.closed = false
  for i := 0; i < int(p.workers); i++ {
    // 初始化併發單元
    p.newWorker(p.work,p.cancel)
  }
}

// passing work and cancel channels to newWorker() to avoid any potential race condition
// betweeen p.work read & write
func (p *limitedPool) newWorker(work chan *workUnit,cancel chan struct{}) {
  go func(p *limitedPool) {

    var wu *workUnit

    defer func(p *limitedPool) {
      // 捕獲異常，結束掉異常的工作單元，並將其再次作為新的任務啟動
      if err := recover(); err != nil {

        trace := make([]byte,1<<16)
        n := runtime.Stack(trace,true)

        s := fmt.Sprintf(errRecovery,err,string(trace[:int(math.Min(float64(n),float64(7000)))]))

        iwu := wu
        iwu.err = &ErrRecovery{s: s}
        close(iwu.done)

        // need to fire up new worker to replace this one as this one is exiting
        p.newWorker(p.work,p.cancel)
      }
    }(p)

    var value interface{}
    var err error

    for {
      select {
      // workChannel中讀取任務
      case wu = <-work:

        // 防止channel 被關閉後讀取到零值
        if wu == nil {
          continue
        }

        // 先判斷任務是否被取消
        if wu.cancelled.Load() == nil {
          // 執行任務
          value,err = wu.fn(wu)
          wu.writing.Store(struct{}{})
          
          // 任務執行完在寫入結果時需要再次檢查工作單元是否被取消，防止產生競爭條件
          if wu.cancelled.Load() == nil && wu.cancelling.Load() == nil {
            wu.value,wu.err = value,err
            close(wu.done)
          }
        }
      // pool是否被停止
      case <-cancel:
        return
      }
    }

  }(p)
}

往POOL中新增任務,並檢查pool是否關閉

func (p *limitedPool) Queue(fn WorkFunc) WorkUnit {
  w := &workUnit{
    done: make(chan struct{}),fn:  fn,}

  go func() {
    p.m.RLock()
    if p.closed {
      w.err = &ErrPoolClosed{s: errClosed}
      if w.cancelled.Load() == nil {
        close(w.done)
      }
      p.m.RUnlock()
      return
    }
    // 將工作單元寫入workChannel,pool啟動後將由上面newWorker函式中讀取執行
    p.work <- w
    p.m.RUnlock()
  }()

  return w
}

在go-playground/pool包中， limitedPool的批量併發執行還需要藉助batch.go來完成

// batch contains all information for a batch run of WorkUnits
type batch struct {
  pool  Pool     // 上面的limitedPool實現了Pool interface
  m    sync.Mutex  // 互斥鎖，用來判斷closed
  units  []WorkUnit  // 工作單元的slice， 這個主要用在不設併發限制的場景，這裡忽略
  results chan WorkUnit // 結果集,執行完後的workUnit會更新其value,error,可以從結果集channel中讀取
  done  chan struct{} // 通知batch是否完成
  closed bool
  wg   *sync.WaitGroup
}

// go-playground/pool 中有設定併發量和不設併發量的批量任務，都實現Pool interface，初始化batch批量任務時會將之前建立好的Pool傳入newBatch
func newBatch(p Pool) Batch {
  return &batch{
    pool:  p,units:  make([]WorkUnit,4),// capacity it to 4 so it doesn't grow and allocate too many times.
    results: make(chan WorkUnit),done:  make(chan struct{}),wg:   new(sync.WaitGroup),}
}

// 往批量任務中新增workFunc任務
func (b *batch) Queue(fn WorkFunc) {

  b.m.Lock()
  if b.closed {
    b.m.Unlock()
    return
  }
  //往上述的limitPool中新增workFunc
  wu := b.pool.Queue(fn)

  b.units = append(b.units,wu) // keeping a reference for cancellation purposes
  b.wg.Add(1)
  b.m.Unlock()
  
  // 執行完後將workUnit寫入結果集channel
  go func(b *batch,wu WorkUnit) {
    wu.Wait()
    b.results <- wu
    b.wg.Done()
  }(b,wu)
}

// 通知批量任務不再接受新的workFunc,如果新增完workFunc不執行改方法的話將導致取結果集時done channel一直阻塞
func (b *batch) QueueComplete() {
  b.m.Lock()
  b.closed = true
  close(b.done)
  b.m.Unlock()
}

// 獲取批量任務結果集
func (b *batch) Results() <-chan WorkUnit {
  go func(b *batch) {
    <-b.done
    b.m.Lock()
    b.wg.Wait()
    b.m.Unlock()
    close(b.results)
  }(b)
  return b.results
}

測試

func SendMail(int int) pool.WorkFunc {
  fn := func(wu pool.WorkUnit) (interface{},error) {
    // sleep 1s 模擬發郵件過程
    time.Sleep(time.Second * 1)
    // 模擬異常任務需要取消
    if int == 17 {
      wu.Cancel()
    }
    if wu.IsCancelled() {
      return false,nil
    }
    fmt.Println("send to",int)
    return true,nil
  }
  return fn
}

func TestBatchWork(t *testing.T) {
  // 初始化groutine數量為20的pool
  p := pool.NewLimited(20)
  defer p.Close()
  batch := p.Batch()
  // 設定一個批量任務的過期超時時間
  t := time.After(10 * time.Second)
  go func() {
    for i := 0; i < 100; i++ {
      batch.Queue(SendMail(i))
    }
    batch.QueueComplete()
  }()
  // 因為 batch.Results 中要close results channel 所以不能將其放在LOOP中執行
  r := batch.Results()
LOOP:
  for {
    select {
    case <-t:
    // 登臺超時通知
      fmt.Println("recived timeout")
      break LOOP
   
    case email,ok := <-r:
    // 讀取結果集
      if ok {
        if err := email.Error(); err != nil {
          fmt.Println("err",err.Error())
        }
        fmt.Println(email.Value())
      } else {
        fmt.Println("finish")
        break LOOP
      }
    }
  }
}

  

接近理論值5s,通知模擬被取消的work也正常取消

go-playground/pool在比起之前簡單的協程池的基礎上， 對pool,worker的狀態有了很好的管理。但是，但是問題來了，在第一個實現的簡單groutine池和go-playground/pool中，都是先啟動預定好的groutine來完成任務執行，在併發量遠小於任務量的情況下確實能夠做到groutine的複用，如果任務量不多則會導致任務分配到每個groutine不均勻，甚至可能出現啟動的groutine根本不會執行任務從而導致浪費，而且對於協程池也沒有動態的擴容和縮小。所以我又去看了一下ants的設計和實現。

ants

ants是一個受fasthttp啟發的高效能協程池,fasthttp號稱是比go原生的net/http快10倍，其快速高效能的原因之一就是採用了各種池化技術（這個日後再開新坑去讀原始碼),ants相比之前兩種協程池，其模型更像是之前接觸到的資料庫連線池，需要從空餘的worker中取出一個來執行任務,當無可用空餘worker的時候再去建立，而當pool的容量達到上線之後，剩餘的任務阻塞等待當前進行中的worker執行完畢將worker放回pool,直至pool中有空閒worker。 ants在記憶體的管理上做得很好，除了定期清除過期worker(一定時間內沒有分配到任務的worker)，ants還實現了一種適用於大批量相同任務的pool,這種pool與一個需要大批量重複執行的函式鎖繫結，避免了呼叫方不停的建立，更加節省記憶體。

先看一下ants的pool 結構體 (pool.go)

type Pool struct {
  // 協程池的容量 (groutine數量的上限)
  capacity int32
  // 正在執行中的groutine
  running int32
  // 過期清理間隔時間
  expiryDuration time.Duration
  // 當前可用空閒的groutine
  workers []*Worker
  // 表示pool是否關閉
  release int32
  // lock for synchronous operation.
  lock sync.Mutex
  // 用於控制pool等待獲取可用的groutine
  cond *sync.Cond
  // 確保pool只被關閉一次
  once sync.Once
  // worker臨時物件池，在複用worker時減少新物件的建立並加速worker從pool中的獲取速度
  workerCache sync.Pool
  // pool引發panic時的執行函式
  PanicHandler func(interface{})
}

接下來看pool的工作單元 worker (worker.go)

type Worker struct {
  // worker 所屬的poo;
  pool *Pool
  // 任務佇列
  task chan func()
  // 回收時間，即該worker的最後一次結束執行的時間
  recycleTime time.Time
}

執行worker的程式碼 (worker.go)

func (w *Worker) run() {
  // pool中正在執行的worker數+1
  w.pool.incRunning()
  go func() {
    defer func() {
      if p := recover(); p != nil {
        //若worker因各種問題引發panic,//pool中正在執行的worker數 -1，     
        //如果設定了Pool中的PanicHandler，此時會被呼叫
        w.pool.decRunning()
        if w.pool.PanicHandler != nil {
          w.pool.PanicHandler(p)
        } else {
          log.Printf("worker exits from a panic: %v",p)
        }
      }
    }()
    
    // worker 執行任務佇列
    for f := range w.task {
      //任務佇列中的函式全部被執行完後，
      //pool中正在執行的worker數 -1， 
      //將worker 放回物件池
      if f == nil {
        w.pool.decRunning()
        w.pool.workerCache.Put(w)
        return
      }
      f()
      //worker 執行完任務後放回Pool 
      //使得其餘正在阻塞的任務可以獲取worker
      w.pool.revertWorker(w)
    }
  }()
}

瞭解了工作單元worker如何執行任務以及與pool互動後，回到pool中檢視其實現， pool的核心就是取出可用worker提供給任務執行 (pool.go)

// 向pool提交任務
func (p *Pool) Submit(task func()) error {
  if 1 == atomic.LoadInt32(&p.release) {
    return ErrPoolClosed
  }
  // 獲取pool中的可用worker並向其任務佇列中寫入任務
  p.retrieveWorker().task <- task
  return nil
}


// **核心程式碼** 獲取可用worker
func (p *Pool) retrieveWorker() *Worker {
  var w *Worker

  p.lock.Lock()
  idleWorkers := p.workers
  n := len(idleWorkers) - 1
 // 當前pool中有可用worker,取出(隊尾)worker並執行
  if n >= 0 {
    w = idleWorkers[n]
    idleWorkers[n] = nil
    p.workers = idleWorkers[:n]
    p.lock.Unlock()
  } else if p.Running() < p.Cap() {
    p.lock.Unlock()
    // 當前pool中無空閒worker,且pool數量未達到上線
    // pool會先從臨時物件池中尋找是否有已完成任務的worker,// 若臨時物件池中不存在，則重新建立一個worker並將其啟動
    if cacheWorker := p.workerCache.Get(); cacheWorker != nil {
      w = cacheWorker.(*Worker)
    } else {
      w = &Worker{
        pool: p,task: make(chan func(),workerChanCap),}
    }
    w.run()
  } else {
    // pool中沒有空餘worker且達到併發上限
    // 任務會阻塞等待當前執行的worker完成任務釋放會pool
    for {
      p.cond.Wait() // 等待通知， 暫時阻塞
      l := len(p.workers) - 1
      if l < 0 {
        continue
      }
      // 當有可用worker釋放回pool之後， 取出
      w = p.workers[l]
      p.workers[l] = nil
      p.workers = p.workers[:l]
      break
    }
    p.lock.Unlock()
  }
  return w
}

// 釋放worker回pool
func (p *Pool) revertWorker(worker *Worker) {
  worker.recycleTime = time.Now()
  p.lock.Lock()
  p.workers = append(p.workers,worker)
  // 通知pool中已經獲取鎖的groutine,有一個worker已完成任務
  p.cond.Signal()
  p.lock.Unlock()
}

在批量併發任務的執行過程中， 如果有超過5納秒(ants中預設worker過期時間為5ns)的worker未被分配新的任務，則將其作為過期worker清理掉，從而保證pool中可用的worker都能發揮出最大的作用以及將任務分配得更均勻
(pool.go)

// 該函式會在pool初始化後在協程中啟動
func (p *Pool) periodicallyPurge() {
  // 建立一個5ns定時的心跳
  heartbeat := time.NewTicker(p.expiryDuration)
  defer heartbeat.Stop()

  for range heartbeat.C {
    currentTime := time.Now()
    p.lock.Lock()
    idleWorkers := p.workers
    if len(idleWorkers) == 0 && p.Running() == 0 && atomic.LoadInt32(&p.release) == 1 {
      p.lock.Unlock()
      return
    }
    n := -1
    for i,w := range idleWorkers {
      // 因為pool 的worker佇列是先進後出的，所以正序遍歷可用worker時前面的往往裡當前時間越久
      if currentTime.Sub(w.recycleTime) <= p.expiryDuration {
        break
      }  
      // 如果worker最後一次執行時間距現在超過5納秒，視為過期，worker收到nil,執行上述worker.go中 if n == nil 的操作
      n = i
      w.task <- nil
      idleWorkers[i] = nil
    }
    if n > -1 {
      // 全部過期
      if n >= len(idleWorkers)-1 {
        p.workers = idleWorkers[:0]
      } else {
      // 部分過期
        p.workers = idleWorkers[n+1:]
      }
    }
    p.lock.Unlock()
  }
}

測試

func TestAnts(t *testing.T) {
  wg := sync.WaitGroup{}
  pool,_ := ants.NewPool(20)
  defer pool.Release()
  for i := 0; i < 100; i++ {
    wg.Add(1)
    pool.Submit(sendMail(i,&wg))
  }
  wg.Wait()
}

func sendMail(i int,wg *sync.WaitGroup) func() {
  return func() {
    time.Sleep(time.Second * 1)
    fmt.Println("send mail to ",i)
    wg.Done()
  }
}

這裡雖只簡單的測試批量併發任務的場景， 如果大家有興趣可以去看看ants的壓力測試， ants的吞吐量能夠比原生groutine高出N倍，記憶體節省10到20倍， 可謂是協程池中的神器。

借用ants作者的原話來說：

然而又有多少場景是單臺機器需要扛100w甚至1000w同步任務的？基本沒有啊！結果就是造出了屠龍刀，可是世界上沒有龍啊！也是無情…

Over

一口氣從簡單到複雜總結了三個協程池的實現，受益匪淺， 感謝各開源庫的作者， 雖然世界上沒有龍，但是屠龍技是必須練的，因為它就像存款，不一定要全部都用了，但是一定不能沒有！

以上就是本文的全部內容，希望對大家的學習有所幫助，也希望大家多多支援我們。