go中控制goroutine數(shù)量的方法

var (
 // channel長度
 poolCount      = 5
 // 復(fù)用的goroutine數(shù)量
 goroutineCount = 10
)

func pool() {
 jobsChan := make(chan int, poolCount)

 // workers
 var wg sync.WaitGroup
 for i := 0; i  goroutineCount; i++ {
  wg.Add(1)
  go func() {
   defer wg.Done()
   for item := range jobsChan {
    // ...
    fmt.Println(item)
   }
  }()
 }

 // senders
 for i := 0; i  1000; i++ {
  jobsChan - i
 }

 // 關(guān)閉channel，上游的goroutine在讀完channel的內(nèi)容，就會通過wg的done退出
 close(jobsChan)
 wg.Wait()
}

package main

import (
 "context"
 "fmt"
 "sync"
 "time"

 "golang.org/x/sync/semaphore"
)

const (
 // 同時運(yùn)行的goroutine上限
 Limit = 3
 // 信號量的權(quán)重
 Weight = 1
)

func main() {
 names := []string{
  "小白",
  "小紅",
  "小明",
  "小李",
  "小花",
 }

 sem := semaphore.NewWeighted(Limit)
 var w sync.WaitGroup
 for _, name := range names {
  w.Add(1)
  go func(name string) {
   sem.Acquire(context.Background(), Weight)
   // ... 具體的業(yè)務(wù)邏輯
   fmt.Println(name, "-吃飯了")
   time.Sleep(2 * time.Second)
   sem.Release(Weight)
   w.Done()
  }(name)
 }
 w.Wait()

 fmt.Println("ending--------")
}

// workerPool通過一組工作池服務(wù)傳入的連接
// 按照FILO（先進(jìn)后出）的順序，即最近停止的工作人員將為下一個工作傳入的連接。
//
// 這種方案能夠保持cpu的緩存保持高效（理論上）
type workerPool struct {
 // 這個函數(shù)用于server的連接
 // It must leave c unclosed.
 WorkerFunc ServeHandler

 // 最大的Workers數(shù)量
 MaxWorkersCount int

 LogAllErrors bool

 MaxIdleWorkerDuration time.Duration

 Logger Logger

 lock         sync.Mutex
 // 當(dāng)前worker的數(shù)量
 workersCount int
 // worker停止的標(biāo)識
 mustStop     bool

 // 等待使用的workerChan
 // 可能會被清理
 ready []*workerChan

 // 用來標(biāo)識start和stop
 stopCh chan struct{}

 // workerChan的緩存池，通過sync.Pool實(shí)現(xiàn)
 workerChanPool sync.Pool

 connState func(net.Conn, ConnState)
}

// workerChan的結(jié)構(gòu)
type workerChan struct {
 lastUseTime time.Time
 ch          chan net.Conn
}

func (wp *workerPool) Start() {
 // 判斷是否已經(jīng)Start過了
 if wp.stopCh != nil {
  panic("BUG: workerPool already started")
 }
 // stopCh塞入值
 wp.stopCh = make(chan struct{})
 stopCh := wp.stopCh
 wp.workerChanPool.New = func() interface{} {
  // 如果單核cpu則讓workerChan阻塞
  // 否則，使用非阻塞，workerChan的長度為1
  return workerChan{
   ch: make(chan net.Conn, workerChanCap),
  }
 }
 go func() {
  var scratch []*workerChan
  for {
   wp.clean(scratch)
   select {
   // 接收到退出信號，退出
   case -stopCh:
    return
   default:
    time.Sleep(wp.getMaxIdleWorkerDuration())
   }
  }
 }()
}

// 如果單核cpu則讓workerChan阻塞
// 否則，使用非阻塞，workerChan的長度為1
var workerChanCap = func() int {
 // 如果GOMAXPROCS=1，workerChan的長度為0，變成一個阻塞的channel
 if runtime.GOMAXPROCS(0) == 1 {
  return 0
 }

 // 如果GOMAXPROCS>1則使用非阻塞的workerChan
 return 1
}()

func (wp *workerPool) Stop() {
 // 同start，stop也只能觸發(fā)一次
 if wp.stopCh == nil {
  panic("BUG: workerPool wasn't started")
 }
 // 關(guān)閉stopCh
 close(wp.stopCh)
 // 將stopCh置為nil
 wp.stopCh = nil

 // 停止所有的等待獲取連接的workers
 // 正在運(yùn)行的workers，不需要等待他們退出，他們會在完成connection或mustStop被設(shè)置成true退出
 wp.lock.Lock()
 ready := wp.ready
 // 循環(huán)將ready的workerChan置為nil
 for i := range ready {
  ready[i].ch - nil
  ready[i] = nil
 }
 wp.ready = ready[:0]
 // 設(shè)置mustStop為true
 wp.mustStop = true
 wp.lock.Unlock()
}

func (wp *workerPool) clean(scratch *[]*workerChan) {
 maxIdleWorkerDuration := wp.getMaxIdleWorkerDuration()

 // 清理掉最近最少使用的workers如果他們過了maxIdleWorkerDuration時間沒有提供服務(wù)
 criticalTime := time.Now().Add(-maxIdleWorkerDuration)

 wp.lock.Lock()
 ready := wp.ready
 n := len(ready)

 // 使用二分搜索算法找出最近可以被清除的worker
 // 最后使用的workerChan 一定是放回隊列尾部的。
 l, r, mid := 0, n-1, 0
 for l = r {
  mid = (l + r) / 2
  if criticalTime.After(wp.ready[mid].lastUseTime) {
   l = mid + 1
  } else {
   r = mid - 1
  }
 }
 i := r
 if i == -1 {
  wp.lock.Unlock()
  return
 }

 // 將ready中i之前的的全部清除
 *scratch = append((*scratch)[:0], ready[:i+1]...)
 m := copy(ready, ready[i+1:])
 for i = m; i  n; i++ {
  ready[i] = nil
 }
 wp.ready = ready[:m]
 wp.lock.Unlock()

 // 通知淘汰的workers停止
 // 此通知必須位于wp.lock之外，因?yàn)閏h.ch
 // 如果有很多workers，可能會阻塞并且可能會花費(fèi)大量時間
 // 位于非本地CPU上。
 tmp := *scratch
 for i := range tmp {
  tmp[i].ch - nil
  tmp[i] = nil
 }
}

func (wp *workerPool) getCh() *workerChan {
 var ch *workerChan
 createWorker := false

 wp.lock.Lock()
 ready := wp.ready
 n := len(ready) - 1
 // 如果ready為空
 if n  0 {
  if wp.workersCount  wp.MaxWorkersCount {
   createWorker = true
   wp.workersCount++
  }
 } else {
  // 不為空從ready中取一個
  ch = ready[n]
  ready[n] = nil
  wp.ready = ready[:n]
 }
 wp.lock.Unlock()

 // 如果沒拿到ch
 if ch == nil {
  if !createWorker {
   return nil
  }
  // 從緩存中獲取一個ch
  vch := wp.workerChanPool.Get()
  ch = vch.(*workerChan)
  go func() {
   // 具體的執(zhí)行函數(shù)
   wp.workerFunc(ch)
   // 再放入到pool中
   wp.workerChanPool.Put(vch)
  }()
 }
 return ch
}

func (wp *workerPool) workerFunc(ch *workerChan) {
 var c net.Conn

 var err error
 // 監(jiān)聽workerChan
 for c = range ch.ch {
  if c == nil {
   break
  }

  // 具體的業(yè)務(wù)邏輯
  ...
  c = nil

  // 釋放workerChan
  // 在mustStop的時候?qū)鲅h(huán)
  if !wp.release(ch) {
   break
  }
 }

 wp.lock.Lock()
 wp.workersCount--
 wp.lock.Unlock()
}

// 把Conn放入到channel中
func (wp *workerPool) Serve(c net.Conn) bool {
 ch := wp.getCh()
 if ch == nil {
  return false
 }
 ch.ch - c
 return true
}

func (wp *workerPool) release(ch *workerChan) bool {
 // 修改 ch.lastUseTime
 ch.lastUseTime = time.Now()
 wp.lock.Lock()
 // 如果需要停止，直接返回
 if wp.mustStop {
  wp.lock.Unlock()
  return false
 }
 // 將ch放到ready中
 wp.ready = append(wp.ready, ch)
 wp.lock.Unlock()
 return true
}

package main

import (
 "fmt"
 "sync"
 "sync/atomic"
 "time"

 "github.com/panjf2000/ants"
)

func demoFunc() {
 time.Sleep(10 * time.Millisecond)
 fmt.Println("Hello World!")
}

func main() {
 defer ants.Release()

 runTimes := 1000

 var wg sync.WaitGroup
 syncCalculateSum := func() {
  demoFunc()
  wg.Done()
 }
 for i := 0; i  runTimes; i++ {
  wg.Add(1)
  _ = ants.Submit(syncCalculateSum)
 }
 wg.Wait()
 fmt.Printf("running goroutines: %d\n", ants.Running())
 fmt.Printf("finish all tasks.\n")
}

package main

import (
 "fmt"
 "sync"
 "sync/atomic"
 "time"

 "github.com/panjf2000/ants"
)

var sum int32

func myFunc(i interface{}) {
 n := i.(int32)
 atomic.AddInt32(sum, n)
 fmt.Printf("run with %d\n", n)
}

func main() {
 var wg sync.WaitGroup
 runTimes := 1000

 // Use the pool with a method,
 // set 10 to the capacity of goroutine pool and 1 second for expired duration.
 p, _ := ants.NewPoolWithFunc(10, func(i interface{}) {
  myFunc(i)
  wg.Done()
 })
 defer p.Release()
 // Submit tasks one by one.
 for i := 0; i  runTimes; i++ {
  wg.Add(1)
  _ = p.Invoke(int32(i))
 }
 wg.Wait()
 fmt.Printf("running goroutines: %d\n", p.Running())
 fmt.Printf("finish all tasks, result is %d\n", sum)
 if sum != 499500 {
  panic("the final result is wrong!!!")
 }
}

package main

import (
 "fmt"
 "time"

 "gopkg.in/go-playground/pool.v3"
)

func main() {

 p := pool.NewLimited(10)
 defer p.Close()

 user := p.Queue(getUser(13))
 other := p.Queue(getOtherInfo(13))

 user.Wait()
 if err := user.Error(); err != nil {
  // handle error
 }

 // do stuff with user
 username := user.Value().(string)
 fmt.Println(username)

 other.Wait()
 if err := other.Error(); err != nil {
  // handle error
 }

 // do stuff with other
 otherInfo := other.Value().(string)
 fmt.Println(otherInfo)
}

func getUser(id int) pool.WorkFunc {

 return func(wu pool.WorkUnit) (interface{}, error) {

  // simulate waiting for something, like TCP connection to be established
  // or connection from pool grabbed
  time.Sleep(time.Second * 1)

  if wu.IsCancelled() {
   // return values not used
   return nil, nil
  }

  // ready for processing...

  return "Joeybloggs", nil
 }
}

func getOtherInfo(id int) pool.WorkFunc {

 return func(wu pool.WorkUnit) (interface{}, error) {

  // simulate waiting for something, like TCP connection to be established
  // or connection from pool grabbed
  time.Sleep(time.Second * 1)

  if wu.IsCancelled() {
   // return values not used
   return nil, nil
  }

  // ready for processing...

  return "Other Info", nil
 }
}

package main

import (
 "fmt"
 "time"

 "gopkg.in/go-playground/pool.v3"
)

func main() {

 p := pool.NewLimited(10)
 defer p.Close()

 batch := p.Batch()

 // for max speed Queue in another goroutine
 // but it is not required, just can't start reading results
 // until all items are Queued.

 go func() {
  for i := 0; i  10; i++ {
   batch.Queue(sendEmail("email content"))
  }

  // DO NOT FORGET THIS OR GOROUTINES WILL DEADLOCK
  // if calling Cancel() it calles QueueComplete() internally
  batch.QueueComplete()
 }()

 for email := range batch.Results() {

  if err := email.Error(); err != nil {
   // handle error
   // maybe call batch.Cancel()
  }

  // use return value
  fmt.Println(email.Value().(bool))
 }
}

func sendEmail(email string) pool.WorkFunc {
 return func(wu pool.WorkUnit) (interface{}, error) {

  // simulate waiting for something, like TCP connection to be established
  // or connection from pool grabbed
  time.Sleep(time.Second * 1)

  if wu.IsCancelled() {
   // return values not used
   return nil, nil
  }

  // ready for processing...

  return true, nil // everything ok, send nil, error if not
 }
}

// WorkUnit contains a single uint of works values
type WorkUnit interface {

 // 阻塞直到當(dāng)前任務(wù)被完成或被取消
 Wait()

 // 執(zhí)行函數(shù)返回的結(jié)果
 Value() interface{}

 // Error returns the Work Unit's error
 Error() error

 // 取消當(dāng)前的可執(zhí)行任務(wù)
 Cancel()

 // 判斷當(dāng)前的可執(zhí)行單元是否被取消了
 IsCancelled() bool
}

var _ WorkUnit = new(workUnit)

// workUnit contains a single unit of works values
type workUnit struct {
 // 任務(wù)執(zhí)行的結(jié)果
 value      interface{}
 // 錯誤信息
 err        error
 // 通知任務(wù)完成
 done       chan struct{}
 // 需要執(zhí)行的任務(wù)函數(shù)
 fn         WorkFunc
 // 任務(wù)是會否被取消
 cancelled  atomic.Value
 // 是否正在取消任務(wù)
 cancelling atomic.Value
 // 任務(wù)是否正在執(zhí)行
 writing    atomic.Value
}

var _ Pool = new(limitedPool)

// limitedPool contains all information for a limited pool instance.
type limitedPool struct {
 // 并發(fā)量
 workers uint
 // work的channel
 work    chan *workUnit
 // 通知結(jié)束的channel
 cancel  chan struct{}
 // 是否關(guān)閉的標(biāo)識
 closed  bool
 // 讀寫鎖
 m       sync.RWMutex
}

// 初始化一個pool
func NewLimited(workers uint) Pool {

 if workers == 0 {
  panic("invalid workers '0'")
 }
 // 初始化pool的work數(shù)量
 p := limitedPool{
  workers: workers,
 }
 // 初始化pool的操作
 p.initialize()

 return p
}

func (p *limitedPool) initialize() {
 // channel的長度為work數(shù)量的兩倍
 p.work = make(chan *workUnit, p.workers*2)
 p.cancel = make(chan struct{})
 p.closed = false

 // fire up workers here
 for i := 0; i  int(p.workers); i++ {
  p.newWorker(p.work, p.cancel)
 }
}

// 將工作傳遞并取消頻道到newWorker（）以避免任何潛在的競爭狀況
// 在p.work讀寫之間
func (p *limitedPool) newWorker(work chan *workUnit, cancel chan struct{}) {
 go func(p *limitedPool) {

  var wu *workUnit

  defer func(p *limitedPool) {
   // 捕獲異常，結(jié)束掉異常的工作單元，并將其再次作為新的任務(wù)啟動
   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)

    // 重新啟動
    p.newWorker(p.work, p.cancel)
   }
  }(p)

  var value interface{}
  var err error
  // 監(jiān)聽channel，讀取內(nèi)容
  for {
   select {
   // channel中取出數(shù)據(jù)
   case wu = -work:

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

    // 單個和批量的cancellation這個都支持
    if wu.cancelled.Load() == nil {
     // 執(zhí)行我們的業(yè)務(wù)函數(shù)
     value, err = wu.fn(wu)

     wu.writing.Store(struct{}{})

     // 如果WorkFunc取消了此工作單元，則需要再次檢查
     // 防止產(chǎn)生競爭條件
     if wu.cancelled.Load() == nil  wu.cancelling.Load() == nil {
      wu.value, wu.err = value, err

      // 執(zhí)行完成，關(guān)閉當(dāng)前channel
      close(wu.done)
     }
    }
    // 如果取消了，就退出
   case -cancel:
    return
   }
  }

 }(p)
}

// 放置一個執(zhí)行的task到channel，并返回channel
func (p *limitedPool) Queue(fn WorkFunc) WorkUnit {
 // 初始化一個workUnit類型的channel
 w := workUnit{
  done: make(chan struct{}),
  // 具體的執(zhí)行函數(shù)
  fn:   fn,
 }

 go func() {
  p.m.RLock()
  // 如果pool關(guān)閉的時候通知channel關(guān)閉
  if p.closed {
   w.err = ErrPoolClosed{s: errClosed}
   if w.cancelled.Load() == nil {
    close(w.done)
   }
   p.m.RUnlock()
   return
  }
  // 將channel傳遞給pool的work
  p.work - w

  p.m.RUnlock()
 }()

 return w
}

// batch contains all information for a batch run of WorkUnits
type batch struct {
 pool    Pool
 m       sync.Mutex
 // WorkUnit的切片
 units   []WorkUnit
 // 結(jié)果集,執(zhí)行完后的workUnit會更新其value,error,可以從結(jié)果集channel中讀取
 results chan WorkUnit
 // 通知batch是否完成
 done    chan struct{}
 closed  bool
 wg      *sync.WaitGroup
}

// 初始化Batch
func newBatch(p Pool) Batch {
 return batch{
  pool:    p,
  units:   make([]WorkUnit, 0, 4),
  results: make(chan WorkUnit),
  done:    make(chan struct{}),
  wg:      new(sync.WaitGroup),
 }
}


// 將WorkFunc放入到WorkUnit中并保留取消和輸出結(jié)果的參考。
func (b *batch) Queue(fn WorkFunc) {

 b.m.Lock()

 if b.closed {
  b.m.Unlock()
  return
 }
 // 返回一個WorkUnit
 wu := b.pool.Queue(fn)

 // 放到WorkUnit的切片中
 b.units = append(b.units, wu)
 // 通過waitgroup進(jìn)行g(shù)oroutine的執(zhí)行控制
 b.wg.Add(1)
 b.m.Unlock()

 // 執(zhí)行任務(wù)
 go func(b *batch, wu WorkUnit) {
  wu.Wait()
  // 將執(zhí)行的結(jié)果寫入到results中
  b.results - wu
  b.wg.Done()
 }(b, wu)
}


// QueueComplete讓批處理知道不再有排隊的工作單元
// 以便在所有工作完成后可以關(guān)閉結(jié)果渠道。
// 警告：如果未調(diào)用此函數(shù)，則結(jié)果通道將永遠(yuǎn)不會耗盡，
// 但會永遠(yuǎn)阻止以獲取更多結(jié)果。
func (b *batch) QueueComplete() {
 b.m.Lock()
 b.closed = true
 close(b.done)
 b.m.Unlock()
}

// 取消批次的任務(wù)
func (b *batch) Cancel() {

 b.QueueComplete()

 b.m.Lock()

 // 一個個取消units，倒敘的取消
 for i := len(b.units) - 1; i >= 0; i-- {
  b.units[i].Cancel()
 }

 b.m.Unlock()
}

// 輸出執(zhí)行完成的結(jié)果集
func (b *batch) Results() -chan WorkUnit {
 // 啟動一個協(xié)程監(jiān)聽完成的通知
 // waitgroup阻塞直到所有的worker都完成退出
 // 最后關(guān)閉channel
 go func(b *batch) {
  -b.done
  b.m.Lock()
  // 阻塞直到上面waitgroup中的goroutine一個個執(zhí)行完成退出
  b.wg.Wait()
  b.m.Unlock()
  // 關(guān)閉channel
  close(b.results)
 }(b)

 return b.results
}