golang協(xié)程池設(shè)計(jì)詳解

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

type SimplePool struct {
  wg  sync.WaitGroup
  work chan func() //任務(wù)隊(duì)列
}

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

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

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() {
    // 模擬抓取數(shù)據(jù)的過(guò)程
    time.Sleep(time.Second * 1)
    fmt.Println("finish parse ", i)
  }
}

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

// 工作單元
type workUnit struct {
  value   interface{}  // 任務(wù)結(jié)果 
  err    error     // 任務(wù)的報(bào)錯(cuò)
  done    chan struct{} // 通知任務(wù)完成
  fn     WorkFunc  
  cancelled atomic.Value  // 任務(wù)是否被取消
  cancelling atomic.Value  // 是否正在取消任務(wù)
  writing  atomic.Value  // 任務(wù)是否正在執(zhí)行
}

type limitedPool struct {
  workers uint      // 并發(fā)量 
  work  chan *workUnit // 任務(wù)channel
  cancel chan struct{}  // 用于通知結(jié)束的channel
  closed bool      // 是否關(guān)閉
  m    sync.RWMutex  // 讀寫(xiě)鎖，主要用來(lái)保證 closed值的并發(fā)安全
}

// 初始化pool,設(shè)定并發(fā)量
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++ {
    // 初始化并發(fā)單元
    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) {
      // 捕獲異常，結(jié)束掉異常的工作單元，并將其再次作為新的任務(wù)啟動(dòng)
      if err := recover(); err != nil {

        trace := make([]byte, 116)
        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中讀取任務(wù)
      case wu = -work:

        // 防止channel 被關(guān)閉后讀取到零值
        if wu == nil {
          continue
        }

        // 先判斷任務(wù)是否被取消
        if wu.cancelled.Load() == nil {
          // 執(zhí)行任務(wù)
          value, err = wu.fn(wu)
          wu.writing.Store(struct{}{})
          
          // 任務(wù)執(zhí)行完在寫(xiě)入結(jié)果時(shí)需要再次檢查工作單元是否被取消，防止產(chǎn)生競(jìng)爭(zhēng)條件
          if wu.cancelled.Load() == nil  wu.cancelling.Load() == nil {
            wu.value, wu.err = value, err
            close(wu.done)
          }
        }
      // pool是否被停止
      case -cancel:
        return
      }
    }

  }(p)
}

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
    }
    // 將工作單元寫(xiě)入workChannel, pool啟動(dòng)后將由上面newWorker函數(shù)中讀取執(zhí)行
    p.work - w
    p.m.RUnlock()
  }()

  return w
}

// batch contains all information for a batch run of WorkUnits
type batch struct {
  pool  Pool     // 上面的limitedPool實(shí)現(xiàn)了Pool interface
  m    sync.Mutex  // 互斥鎖，用來(lái)判斷closed
  units  []WorkUnit  // 工作單元的slice， 這個(gè)主要用在不設(shè)并發(fā)限制的場(chǎng)景，這里忽略
  results chan WorkUnit // 結(jié)果集,執(zhí)行完后的workUnit會(huì)更新其value,error,可以從結(jié)果集channel中讀取
  done  chan struct{} // 通知batch是否完成
  closed bool
  wg   *sync.WaitGroup
}

// go-playground/pool 中有設(shè)置并發(fā)量和不設(shè)并發(fā)量的批量任務(wù)，都實(shí)現(xiàn)Pool interface，初始化batch批量任務(wù)時(shí)會(huì)將之前創(chuàng)建好的Pool傳入newBatch
func newBatch(p Pool) Batch {
  return batch{
    pool:  p,
    units:  make([]WorkUnit, 0, 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),
  }
}

// 往批量任務(wù)中添加workFunc任務(wù)
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()
  
  // 執(zhí)行完后將workUnit寫(xiě)入結(jié)果集channel
  go func(b *batch, wu WorkUnit) {
    wu.Wait()
    b.results - wu
    b.wg.Done()
  }(b, wu)
}

// 通知批量任務(wù)不再接受新的workFunc, 如果添加完workFunc不執(zhí)行改方法的話(huà)將導(dǎo)致取結(jié)果集時(shí)done channel一直阻塞
func (b *batch) QueueComplete() {
  b.m.Lock()
  b.closed = true
  close(b.done)
  b.m.Unlock()
}

// 獲取批量任務(wù)結(jié)果集
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 模擬發(fā)郵件過(guò)程
    time.Sleep(time.Second * 1)
    // 模擬異常任務(wù)需要取消
    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數(shù)量為20的pool
  p := pool.NewLimited(20)
  defer p.Close()
  batch := p.Batch()
  // 設(shè)置一個(gè)批量任務(wù)的過(guò)期超時(shí)時(shí)間
  t := time.After(10 * time.Second)
  go func() {
    for i := 0; i  100; i++ {
      batch.Queue(SendMail(i))
    }
    batch.QueueComplete()
  }()
  // 因?yàn)?batch.Results 中要close results channel 所以不能將其放在LOOP中執(zhí)行
  r := batch.Results()
LOOP:
  for {
    select {
    case -t:
    // 登臺(tái)超時(shí)通知
      fmt.Println("recived timeout")
      break LOOP
   
    case email, ok := -r:
    // 讀取結(jié)果集
      if ok {
        if err := email.Error(); err != nil {
          fmt.Println("err", err.Error())
        }
        fmt.Println(email.Value())
      } else {
        fmt.Println("finish")
        break LOOP
      }
    }
  }
}

  

type Pool struct {
  // 協(xié)程池的容量 (groutine數(shù)量的上限)
  capacity int32
  // 正在執(zhí)行中的groutine
  running int32
  // 過(guò)期清理間隔時(shí)間
  expiryDuration time.Duration
  // 當(dāng)前可用空閑的groutine
  workers []*Worker
  // 表示pool是否關(guān)閉
  release int32
  // lock for synchronous operation.
  lock sync.Mutex
  // 用于控制pool等待獲取可用的groutine
  cond *sync.Cond
  // 確保pool只被關(guān)閉一次
  once sync.Once
  // worker臨時(shí)對(duì)象池，在復(fù)用worker時(shí)減少新對(duì)象的創(chuàng)建并加速worker從pool中的獲取速度
  workerCache sync.Pool
  // pool引發(fā)panic時(shí)的執(zhí)行函數(shù)
  PanicHandler func(interface{})
}

type Worker struct {
  // worker 所屬的poo;
  pool *Pool
  // 任務(wù)隊(duì)列
  task chan func()
  // 回收時(shí)間，即該worker的最后一次結(jié)束運(yùn)行的時(shí)間
  recycleTime time.Time
}

func (w *Worker) run() {
  // pool中正在執(zhí)行的worker數(shù)+1
  w.pool.incRunning()
  go func() {
    defer func() {
      if p := recover(); p != nil {
        //若worker因各種問(wèn)題引發(fā)panic, 
        //pool中正在執(zhí)行的worker數(shù) -1，     
        //如果設(shè)置了Pool中的PanicHandler，此時(shí)會(huì)被調(diào)用
        w.pool.decRunning()
        if w.pool.PanicHandler != nil {
          w.pool.PanicHandler(p)
        } else {
          log.Printf("worker exits from a panic: %v", p)
        }
      }
    }()
    
    // worker 執(zhí)行任務(wù)隊(duì)列
    for f := range w.task {
      //任務(wù)隊(duì)列中的函數(shù)全部被執(zhí)行完后，
      //pool中正在執(zhí)行的worker數(shù) -1， 
      //將worker 放回對(duì)象池
      if f == nil {
        w.pool.decRunning()
        w.pool.workerCache.Put(w)
        return
      }
      f()
      //worker 執(zhí)行完任務(wù)后放回Pool 
      //使得其余正在阻塞的任務(wù)可以獲取worker
      w.pool.revertWorker(w)
    }
  }()
}

// 向pool提交任務(wù)
func (p *Pool) Submit(task func()) error {
  if 1 == atomic.LoadInt32(p.release) {
    return ErrPoolClosed
  }
  // 獲取pool中的可用worker并向其任務(wù)隊(duì)列中寫(xiě)入任務(wù)
  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
 // 當(dāng)前pool中有可用worker, 取出(隊(duì)尾)worker并執(zhí)行
  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()
    // 當(dāng)前pool中無(wú)空閑worker,且pool數(shù)量未達(dá)到上線(xiàn)
    // pool會(huì)先從臨時(shí)對(duì)象池中尋找是否有已完成任務(wù)的worker,
    // 若臨時(shí)對(duì)象池中不存在，則重新創(chuàng)建一個(gè)worker并將其啟動(dòng)
    if cacheWorker := p.workerCache.Get(); cacheWorker != nil {
      w = cacheWorker.(*Worker)
    } else {
      w = Worker{
        pool: p,
        task: make(chan func(), workerChanCap),
      }
    }
    w.run()
  } else {
    // pool中沒(méi)有空余worker且達(dá)到并發(fā)上限
    // 任務(wù)會(huì)阻塞等待當(dāng)前運(yùn)行的worker完成任務(wù)釋放會(huì)pool
    for {
      p.cond.Wait() // 等待通知， 暫時(shí)阻塞
      l := len(p.workers) - 1
      if l  0 {
        continue
      }
      // 當(dāng)有可用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中已經(jīng)獲取鎖的groutine, 有一個(gè)worker已完成任務(wù)
  p.cond.Signal()
  p.lock.Unlock()
}

// 該函數(shù)會(huì)在pool初始化后在協(xié)程中啟動(dòng)
func (p *Pool) periodicallyPurge() {
  // 創(chuàng)建一個(gè)5ns定時(shí)的心跳
  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 {
      // 因?yàn)閜ool 的worker隊(duì)列是先進(jìn)后出的，所以正序遍歷可用worker時(shí)前面的往往里當(dāng)前時(shí)間越久
      if currentTime.Sub(w.recycleTime) = p.expiryDuration {
        break
      }  
      // 如果worker最后一次運(yùn)行時(shí)間距現(xiàn)在超過(guò)5納秒，視為過(guò)期，worker收到nil, 執(zhí)行上述worker.go中 if n == nil 的操作
      n = i
      w.task - nil
      idleWorkers[i] = nil
    }
    if n > -1 {
      // 全部過(guò)期
      if n >= len(idleWorkers)-1 {
        p.workers = idleWorkers[:0]
      } else {
      // 部分過(guò)期
        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()
  }
}