feat(queue): all task has been done after closing the channel (#56)

Author: appleboy

consumer.go

@@ -119,10 +119,13 @@ func (s *Consumer) Queue(task QueuedMessage) error {
 
 func (s *Consumer) Request() (QueuedMessage, error) {
 	select {
-	case task := <-s.taskQueue:
+	case task, ok := <-s.taskQueue:
+		if !ok {
+			return nil, ErrQueueHasBeenClosed
+		}
 		return task, nil
 	default:
-		return nil, errors.New("no task in queue")
+		return nil, ErrNoTaskInQueue
 	}
 }
 

consumer_test.go

@@ -148,11 +148,10 @@ func TestGoroutineLeak(t *testing.T) {
 			for {
 				select {
 				case <-ctx.Done():
-					log.Println("get data:", string(m.Bytes()))
 					if errors.Is(ctx.Err(), context.Canceled) {
-						log.Println("queue has been shutdown and cancel the job")
+						log.Println("queue has been shutdown and cancel the job: " + string(m.Bytes()))
 					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-						log.Println("job deadline exceeded")
+						log.Println("job deadline exceeded: " + string(m.Bytes()))
 					}
 					return nil
 				default:

errors.go

@@ -0,0 +1,10 @@
+package queue
+
+import "errors"
+
+var (
+	// ErrNoTaskInQueue there is nothing in the queue
+	ErrNoTaskInQueue = errors.New("no task in queue")
+	// ErrQueueHasBeenClosed the current queue is closed
+	ErrQueueHasBeenClosed = errors.New("queue has been closed")
+)

queue.go

@@ -246,46 +246,41 @@ func (q *Queue) start() {
 	tasks := make(chan QueuedMessage, 1)
 
 	for {
-		var task QueuedMessage
-
 		// request task from queue in background
 		q.routineGroup.Run(func() {
 			for {
-				select {
-				case <-q.quit:
-					return
-				default:
-					t, err := q.worker.Request()
-					if t == nil || err != nil {
-						if err != nil {
-							select {
-							case <-q.quit:
+				t, err := q.worker.Request()
+				if t == nil || err != nil {
+					if err != nil {
+						select {
+						case <-q.quit:
+							if !errors.Is(err, ErrNoTaskInQueue) {
+								close(tasks)
 								return
-							case <-time.After(time.Second):
-								// sleep 1 second to fetch new task
 							}
+						case <-time.After(time.Second):
+							// sleep 1 second to fetch new task
 						}
 					}
-					if t != nil {
-						tasks <- t
+				}
+				if t != nil {
+					tasks <- t
+					return
+				}
+
+				select {
+				case <-q.quit:
+					if !errors.Is(err, ErrNoTaskInQueue) {
+						close(tasks)
 						return
 					}
+				default:
 				}
 			}
 		})
 
-		// read task
-		select {
-		case task = <-tasks:
-		case <-q.quit:
-			select {
-			case task = <-tasks:
-				// queue task before shutdown the service
-				if err := q.worker.Queue(task); err != nil {
-					q.logger.Errorf("can't re-queue task: %v", err)
-				}
-			default:
-			}
+		task, ok := <-tasks
+		if !ok {
 			return
 		}