fix(worker): handle all jobs with graceful shutdown (#54)

Author: appleboy

consumer.go

@@ -76,13 +76,6 @@ func (s *Consumer) handle(job Job) error {
 
 // Run start the worker
 func (s *Consumer) Run(task QueuedMessage) error {
-	// check queue status
-	select {
-	case <-s.stop:
-		return ErrQueueShutdown
-	default:
-	}
-
 	var data Job
 	_ = json.Unmarshal(task.Bytes(), &data)
 	if v, ok := task.(Job); ok {

consumer_test.go

@@ -143,7 +143,7 @@ func TestCancelJobAfterShutdown(t *testing.T) {
 
 func TestGoroutineLeak(t *testing.T) {
 	w := NewConsumer(
-		WithLogger(NewEmptyLogger()),
+		WithLogger(NewLogger()),
 		WithFn(func(ctx context.Context, m QueuedMessage) error {
 			for {
 				select {
@@ -178,7 +178,7 @@ func TestGoroutineLeak(t *testing.T) {
 	}
 
 	q.Start()
-	time.Sleep(2 * time.Second)
+	time.Sleep(1 * time.Second)
 	q.Release()
 	fmt.Println("number of goroutines:", runtime.NumGoroutine())
 }

queue.go

@@ -81,7 +81,9 @@ func NewQueue(opts ...Option) (*Queue, error) {
 
 // Start to enable all worker
 func (q *Queue) Start() {
-	go q.start()
+	q.routineGroup.Run(func() {
+		q.start()
+	})
 }
 
 // Shutdown stops all queues.
@@ -220,6 +222,7 @@ func (q *Queue) work(task QueuedMessage) {
 	}
 }
 
+// UpdateWorkerCount to update worker number dynamically.
 func (q *Queue) UpdateWorkerCount(num int) {
 	q.workerCount = num
 	q.schedule()
@@ -244,13 +247,9 @@ func (q *Queue) start() {
 
 	for {
 		var task QueuedMessage
-		if atomic.LoadInt32(&q.stopFlag) == 1 {
-			return
-		}
 
 		// request task from queue in background
 		q.routineGroup.Run(func() {
-		loop:
 			for {
 				select {
 				case <-q.quit:
@@ -261,15 +260,15 @@ func (q *Queue) start() {
 						if err != nil {
 							select {
 							case <-q.quit:
-								break loop
+								return
 							case <-time.After(time.Second):
 								// sleep 1 second to fetch new task
 							}
 						}
 					}
 					if t != nil {
 						tasks <- t
-						break loop
+						return
 					}
 				}
 			}
@@ -292,24 +291,8 @@ func (q *Queue) start() {
 
 		// check worker number
 		q.schedule()
-
-		// get worker to execute new task
-		select {
-		case <-q.quit:
-			if err := q.worker.Queue(task); err != nil {
-				q.logger.Errorf("can't re-queue task: %v", err)
-			}
-			return
-		case <-q.ready:
-			select {
-			case <-q.quit:
-				if err := q.worker.Queue(task); err != nil {
-					q.logger.Errorf("can't re-queue task: %v", err)
-				}
-				return
-			default:
-			}
-		}
+		// wait worker ready
+		<-q.ready
 
 		// start new task
 		q.metric.IncBusyWorker()