1. package main
  3. import (
  4.     "errors"
  5.     "fmt"
  6.     "log"
  7.     "sync"
  8.     "sync/atomic"
  9. )
  11. type Task struct {
  12.     Handle func(...interface{})
  13.     Params []interface{}
  14.     wg *sync.WaitGroup
  15. }
  17. type Pool struct {
  18.     Capacity uint64
  19.     WorksNums uint64
  20.     State int32
  21.     TaskChan chan *Task
  22.     FinishChan chan  bool
  23.     PanicHandler func(interface{})
  24. }
  26. func NewPool(c uint64)*Pool{
  27.     return &Pool{
  28.         Capacity:   c,
  29.         WorksNums:  0,
  30.         State:      1,
  31.         TaskChan:   make(chan *Task),
  32.         FinishChan: make(chan bool),
  33.     }
  34. }
  36. func(p *Pool)run(){
  37.     p.incWorkNum()
  38.     go func() {
  39.         defer func() {
  40.             p.desWorkNum()
  41.             if r :=recover();r!=nil{
  42.                 if p.PanicHandler!=nil{
  43.                     p.PanicHandler(r)
  44.                 }else {
  45.                     log.Printf("work panic%s\n",r)
  46.                 }
  47.             }
  48.         }()
  49.         for {
  50.             select {
  51.             case task,ok:=<-p.TaskChan:
  52.                 if !ok{
  53.                     return
  54.                 }
  55.                 task.Handle(task.Params)
  56.                 task.wg.Done()
  57.             case <-p.FinishChan:
  58.                 return
  59.             }
  61.         }
  62.     }()
  63. }
  65. func(p *Pool)Add(task *Task)error{
  66.     if p.State==0{
  67.         return errors.New("pool stack is not start")
  68.     }
  69.     if p.getWorkNum()<p.getCap(){
  70.         p.run()
  71.     }
  72.     p.TaskChan<-task
  73.     return nil
  74. }
  76. func (p *Pool)Close() {
  77.     if atomic.LoadInt32(&p.State)==0 {
  78.         return
  79.     }
  80.     atomic.StoreInt32(&p.State, 0)
  81.     for len(p.TaskChan) > 0 {
  83.     }
  84.     p.FinishChan <- true
  85.     close(p.TaskChan)
  86. }
  88. func(p *Pool)incWorkNum(){
  89.     atomic.AddUint64(&p.WorksNums,1)
  90. }
  91. func (p *Pool)desWorkNum(){
  92.     atomic.AddUint64(&p.WorksNums,^uint64(0))
  93. }
  94. func (p *Pool)getWorkNum()uint64{
  95.     return atomic.LoadUint64(&p.WorksNums)
  96. }
  97. func (p *Pool)getCap()uint64{
  98.     return atomic.LoadUint64(&p.Capacity)
  99. }
  101. func main() {
  102.     pool := NewPool(10)
  103.     var wg sync.WaitGroup
  104.     for i:=0;i<5;i++{
  105.         wg.Add(1)
  106.         pool.Add(&Task{
  107.             Handle: func(i ...interface{}) {
  108.                 fmt.Println(i[0])
  109.             },
  110.             wg:&wg,
  111.             Params: []interface{}{i},
  112.         })
  113.     }
  114.     wg.Wait()
  115. }

参考

https://segmentfault.com/a/1190000018193161
https://blog.csdn.net/KingEasternSun/article/details/78964267