github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/workerpool.go (about)

     1  package util
     2  
     3  import (
     4  	"errors"
     5  	"runtime"
     6  	"sync"
     7  
     8  	"github.com/zhiqiangxu/util/logger"
     9  	"go.uber.org/zap"
    10  )
    11  
    12  // WorkerPool is pool of workers
    13  type WorkerPool struct {
    14  	wg       sync.WaitGroup
    15  	doneChan chan struct{}
    16  	workChan chan func()
    17  	once     sync.Once
    18  }
    19  
    20  // NewWorkerPool is ctor for WorkerPool
    21  func NewWorkerPool() *WorkerPool {
    22  	wp := &WorkerPool{doneChan: make(chan struct{}), workChan: make(chan func())}
    23  	wp.Start()
    24  	return wp
    25  }
    26  
    27  // Start worker pool
    28  func (wp *WorkerPool) Start() {
    29  	n := runtime.NumCPU()
    30  	for i := 0; i < n; i++ {
    31  		GoFunc(&wp.wg, func() {
    32  			defer func() {
    33  				err := recover()
    34  				if err != nil {
    35  					logger.Instance().Error("workerPool", zap.Any("err", err))
    36  				}
    37  			}()
    38  			for {
    39  				select {
    40  				case f := <-wp.workChan:
    41  					f()
    42  				case <-wp.doneChan:
    43  					return
    44  				}
    45  			}
    46  		})
    47  	}
    48  }
    49  
    50  // Close the worker pool
    51  func (wp *WorkerPool) Close() {
    52  	wp.once.Do(func() {
    53  		close(wp.doneChan)
    54  		wp.wg.Wait()
    55  	})
    56  }
    57  
    58  var (
    59  	// ErrWorkerPoolClosed when run on closed pool
    60  	ErrWorkerPoolClosed = errors.New("workerPool closed")
    61  )
    62  
    63  // Run a task
    64  func (wp *WorkerPool) Run(f func()) error {
    65  	select {
    66  	case wp.workChan <- f:
    67  		return nil
    68  	case <-wp.doneChan:
    69  		return ErrWorkerPoolClosed
    70  	}
    71  }