github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/mq/worker.go (about)

     1  package mq
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  func NewWorker(process func() error, numWorkers int) *Worker {
     8  	return &Worker{
     9  		numWorkers: numWorkers,
    10  		process:    process,
    11  	}
    12  }
    13  
    14  type Worker struct {
    15  	numWorkers  int
    16  	process     func() error
    17  	wg          sync.WaitGroup
    18  	stopChannel chan struct{}
    19  }
    20  
    21  func (mq *Worker) Stop() {
    22  	for i := 0; i < mq.numWorkers; i++ {
    23  		mq.stopChannel <- struct{}{}
    24  	}
    25  	mq.wg.Wait()
    26  }
    27  
    28  func (mq *Worker) Start() {
    29  	mq.stopChannel = make(chan struct{}, 1)
    30  	mq.wg.Add(mq.numWorkers)
    31  
    32  	for i := 0; i < mq.numWorkers; i++ {
    33  		go func(workerID int) {
    34  			defer mq.wg.Done()
    35  
    36  			for {
    37  				select {
    38  				case <-mq.stopChannel:
    39  					return
    40  				default:
    41  					if err := mq.process(); err != nil {
    42  						continue
    43  					}
    44  				}
    45  			}
    46  		}(i)
    47  	}
    48  }