chore: add stop channel in worker (#12)

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

nsq/nsq.go

@@ -30,11 +30,13 @@ type Worker struct {
 	q           *nsq.Consumer
 	p           *nsq.Producer
 	startOnce   sync.Once
+	stopOnce    sync.Once
+	stop        chan struct{}
 	maxInFlight int
 	addr        string
 	topic       string
 	channel     string
-	runFunc     func(queue.QueuedMessage) error
+	runFunc     func(queue.QueuedMessage, <-chan struct{}) error
 }
 
 // WithAddr setup the addr of NSQ
@@ -59,7 +61,7 @@ func WithChannel(channel string) Option {
 }
 
 // WithRunFunc setup the run func of queue
-func WithRunFunc(fn func(queue.QueuedMessage) error) Option {
+func WithRunFunc(fn func(queue.QueuedMessage, <-chan struct{}) error) Option {
 	return func(w *Worker) {
 		w.runFunc = fn
 	}
@@ -80,7 +82,8 @@ func NewWorker(opts ...Option) *Worker {
 		topic:       "gorush",
 		channel:     "ch",
 		maxInFlight: runtime.NumCPU(),
-		runFunc: func(queue.QueuedMessage) error {
+		stop:        make(chan struct{}),
+		runFunc: func(queue.QueuedMessage, <-chan struct{}) error {
 			return nil
 		},
 	}
@@ -125,7 +128,7 @@ func (s *Worker) AfterRun() error {
 }
 
 // Run start the worker
-func (s *Worker) Run(quit chan struct{}) error {
+func (s *Worker) Run() error {
 	wg := &sync.WaitGroup{}
 	s.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
 		wg.Add(1)
@@ -141,11 +144,11 @@ func (s *Worker) Run(quit chan struct{}) error {
 		}
 
 		// run custom process function
-		return s.runFunc(job)
+		return s.runFunc(job, s.stop)
 	}))
 
 	// wait close signal
-	<-quit
+	<-s.stop
 
 	// wait job completed
 	wg.Wait()
@@ -155,8 +158,11 @@ func (s *Worker) Run(quit chan struct{}) error {
 
 // Shutdown worker
 func (s *Worker) Shutdown() error {
-	s.q.Stop()
-	s.p.Stop()
+	s.stopOnce.Do(func() {
+		s.q.Stop()
+		s.p.Stop()
+		close(s.stop)
+	})
 	return nil
 }
 

nsq/nsq_test.go

@@ -60,7 +60,7 @@ func TestCustomFuncAndWait(t *testing.T) {
 		WithAddr(host+":4150"),
 		WithTopic("test"),
 		WithMaxInFlight(2),
-		WithRunFunc(func(msg queue.QueuedMessage) error {
+		WithRunFunc(func(msg queue.QueuedMessage, s <-chan struct{}) error {
 			time.Sleep(500 * time.Millisecond)
 			return nil
 		}),

queue.go

@@ -126,7 +126,7 @@ func (q *Queue) work() {
 			}
 		}()
 		q.logger.Infof("start the worker num: %d", num)
-		if err := q.worker.Run(q.quit); err != nil {
+		if err := q.worker.Run(); err != nil {
 			q.logger.Error(err)
 		}
 		q.logger.Infof("stop the worker num: %d", num)

simple/simple.go

@@ -3,6 +3,7 @@ package simple
 import (
 	"errors"
 	"runtime"
+	"sync"
 
 	"github.com/appleboy/queue"
 )
@@ -17,7 +18,9 @@ var errMaxCapacity = errors.New("max capacity reached")
 // Worker for simple queue using channel
 type Worker struct {
 	queueNotification chan queue.QueuedMessage
-	runFunc           func(queue.QueuedMessage) error
+	runFunc           func(queue.QueuedMessage, <-chan struct{}) error
+	stop              chan struct{}
+	stopOnce          sync.Once
 }
 
 // BeforeRun run script before start worker
@@ -31,17 +34,20 @@ func (s *Worker) AfterRun() error {
 }
 
 // Run start the worker
-func (s *Worker) Run(_ chan struct{}) error {
+func (s *Worker) Run() error {
 	for notification := range s.queueNotification {
 		// run custom process function
-		_ = s.runFunc(notification)
+		_ = s.runFunc(notification, s.stop)
 	}
 	return nil
 }
 
 // Shutdown worker
 func (s *Worker) Shutdown() error {
-	close(s.queueNotification)
+	s.stopOnce.Do(func() {
+		close(s.queueNotification)
+		close(s.stop)
+	})
 	return nil
 }
 
@@ -73,7 +79,7 @@ func WithQueueNum(num int) Option {
 }
 
 // WithRunFunc setup the run func of queue
-func WithRunFunc(fn func(queue.QueuedMessage) error) Option {
+func WithRunFunc(fn func(queue.QueuedMessage, <-chan struct{}) error) Option {
 	return func(w *Worker) {
 		w.runFunc = fn
 	}
@@ -83,7 +89,8 @@ func WithRunFunc(fn func(queue.QueuedMessage) error) Option {
 func NewWorker(opts ...Option) *Worker {
 	w := &Worker{
 		queueNotification: make(chan queue.QueuedMessage, runtime.NumCPU()<<1),
-		runFunc: func(msg queue.QueuedMessage) error {
+		stop:              make(chan struct{}),
+		runFunc: func(queue.QueuedMessage, <-chan struct{}) error {
 			return nil
 		},
 	}

simple/simple_test.go

@@ -48,7 +48,7 @@ func TestCustomFuncAndWait(t *testing.T) {
 		msg: "foo",
 	}
 	w := NewWorker(
-		WithRunFunc(func(msg queue.QueuedMessage) error {
+		WithRunFunc(func(msg queue.QueuedMessage, s <-chan struct{}) error {
 			time.Sleep(500 * time.Millisecond)
 			return nil
 		}),

worker.go

@@ -10,7 +10,7 @@ type Worker interface {
 	// BeforeRun is called before starting the worker
 	BeforeRun() error
 	// Run is called to start the worker
-	Run(chan struct{}) error
+	Run() error
 	// BeforeRun is called after starting the worker
 	AfterRun() error
 	// Shutdown is called if stop all worker
@@ -37,7 +37,7 @@ type emptyWorker struct{}
 
 func (w *emptyWorker) BeforeRun() error              { return nil }
 func (w *emptyWorker) AfterRun() error               { return nil }
-func (w *emptyWorker) Run(chan struct{}) error       { return nil }
+func (w *emptyWorker) Run() error                    { return nil }
 func (w *emptyWorker) Shutdown() error               { return nil }
 func (w *emptyWorker) Queue(job QueuedMessage) error { return nil }
 func (w *emptyWorker) Capacity() int                 { return 0 }
@@ -49,7 +49,7 @@ type queueWorker struct {
 
 func (w *queueWorker) BeforeRun() error { return nil }
 func (w *queueWorker) AfterRun() error  { return nil }
-func (w *queueWorker) Run(chan struct{}) error {
+func (w *queueWorker) Run() error {
 	for msg := range w.messages {
 		if string(msg.Bytes()) == "panic" {
 			panic("show panic")