chore(queue): handle panic and refactor the job handle (#18)

Author: appleboy

queue.go

@@ -138,6 +138,12 @@ func (q *Queue) QueueWithTimeout(timeout time.Duration, job QueuedMessage) error
 }
 
 func (q *Queue) work() {
+	select {
+	case <-q.quit:
+		return
+	default:
+	}
+
 	num := atomic.AddInt32(&q.runningWorkers, 1)
 	if err := q.worker.BeforeRun(); err != nil {
 		q.logger.Fatal(err)

simple/simple.go

@@ -36,6 +36,45 @@ func (s *Worker) AfterRun() error {
 	return nil
 }
 
+func (s *Worker) handle(m interface{}) error {
+	// create channel with buffer size 1 to avoid goroutine leak
+	done := make(chan error, 1)
+	panicChan := make(chan interface{}, 1)
+	job, _ := m.(queue.Job)
+	ctx, cancel := context.WithTimeout(context.Background(), job.Timeout)
+	defer cancel()
+
+	// run the job
+	go func() {
+		// handle panic issue
+		defer func() {
+			if p := recover(); p != nil {
+				panicChan <- p
+			}
+		}()
+
+		// run custom process function
+		done <- s.runFunc(ctx, job)
+	}()
+
+	select {
+	case p := <-panicChan:
+		panic(p)
+	case <-ctx.Done(): // timeout reached
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			s.logger.Infof("job timeout: %s", job.Timeout.String())
+		}
+		// wait job
+		return <-done
+	case <-s.stop: // shutdown service
+		cancel()
+		// wait job
+		return <-done
+	case err := <-done: // job finish and continue to worker
+		return err
+	}
+}
+
 // Run start the worker
 func (s *Worker) Run() error {
 	// check queue status
@@ -46,33 +85,9 @@ func (s *Worker) Run() error {
 	}
 
 	for task := range s.taskQueue {
-		done := make(chan struct{})
-		v, _ := task.(queue.Job)
-		ctx, cancel := context.WithTimeout(context.Background(), v.Timeout)
-		// vet doesn't complain if I do this
-		_ = cancel
-
-		// run the job
-		go func() {
-			// run custom process function
-			_ = s.runFunc(ctx, task)
-			close(done)
-		}()
-
-		select {
-		case <-ctx.Done(): // timeout reached
-			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-				s.logger.Infof("job timeout: %s", v.Timeout.String())
-			}
-			// wait job
-			<-done
-		case <-s.stop: // shutdown service
-			cancel()
-			// wait job
-			<-done
-		case <-done: // job finish and continue to work
+		if err := s.handle(task); err != nil {
+			s.logger.Error(err.Error())
 		}
-
 	}
 	return nil
 }

simple/simple_test.go

@@ -64,10 +64,10 @@ func TestCustomFuncAndWait(t *testing.T) {
 	assert.NoError(t, err)
 	q.Start()
 	time.Sleep(100 * time.Millisecond)
-	assert.NoError(t, w.Queue(m))
-	assert.NoError(t, w.Queue(m))
-	assert.NoError(t, w.Queue(m))
-	assert.NoError(t, w.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
 	time.Sleep(600 * time.Millisecond)
 	q.Shutdown()
 	q.Wait()
@@ -226,3 +226,25 @@ func TestGoroutineLeak(t *testing.T) {
 	q.Wait()
 	fmt.Println("number of goroutines:", runtime.NumGoroutine())
 }
+
+func TestGoroutinePanic(t *testing.T) {
+	m := mockMessage{
+		msg: "foo",
+	}
+	w := NewWorker(
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			panic("missing something")
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(50 * time.Millisecond)
+	assert.NoError(t, q.Queue(m))
+	time.Sleep(50 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+}