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 }