chore: Add stop flag in Queue and refactor to handle stop flow (#19)

Author: appleboy

queue.go

@@ -22,6 +22,7 @@ type (
 		stopOnce       sync.Once
 		runningWorkers int32
 		timeout        time.Duration
+		stopFlag       int32
 	}
 
 	// Job with Timeout
@@ -103,6 +104,10 @@ func (q *Queue) Start() {
 
 // Shutdown stops all queues.
 func (q *Queue) Shutdown() {
+	if !atomic.CompareAndSwapInt32(&q.stopFlag, 0, 1) {
+		return
+	}
+
 	q.stopOnce.Do(func() {
 		if err := q.worker.Shutdown(); err != nil {
 			q.logger.Error(err)
@@ -138,10 +143,8 @@ func (q *Queue) QueueWithTimeout(timeout time.Duration, job QueuedMessage) error
 }
 
 func (q *Queue) work() {
-	select {
-	case <-q.quit:
+	if atomic.LoadInt32(&q.stopFlag) == 1 {
 		return
-	default:
 	}
 
 	num := atomic.AddInt32(&q.runningWorkers, 1)

simple/simple.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"sync"
+	"time"
 
 	"github.com/appleboy/queue"
 )
@@ -41,6 +42,7 @@ func (s *Worker) handle(m interface{}) error {
 	done := make(chan error, 1)
 	panicChan := make(chan interface{}, 1)
 	job, _ := m.(queue.Job)
+	startTime := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), job.Timeout)
 	defer cancel()
 
@@ -61,16 +63,22 @@ func (s *Worker) handle(m interface{}) error {
 	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
+		return ctx.Err()
 	case <-s.stop: // shutdown service
+		// cancel job
 		cancel()
+
+		leftTime := job.Timeout - time.Since(startTime)
 		// wait job
-		return <-done
-	case err := <-done: // job finish and continue to worker
+		select {
+		case <-time.After(leftTime):
+			return context.DeadlineExceeded
+		case err := <-done: // job finish
+			return err
+		case p := <-panicChan:
+			panic(p)
+		}
+	case err := <-done: // job finish
 		return err
 	}
 }

simple/simple_test.go

@@ -248,3 +248,82 @@ func TestGoroutinePanic(t *testing.T) {
 	q.Shutdown()
 	q.Wait()
 }
+
+func TestHandleTimeout(t *testing.T) {
+	job := queue.Job{
+		Timeout: 100 * time.Millisecond,
+		Body:    []byte("foo"),
+	}
+	w := NewWorker(
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			time.Sleep(200 * time.Millisecond)
+			return nil
+		}),
+	)
+
+	err := w.handle(job)
+	assert.Error(t, err)
+	assert.Equal(t, context.DeadlineExceeded, err)
+
+	job = queue.Job{
+		Timeout: 150 * time.Millisecond,
+		Body:    []byte("foo"),
+	}
+
+	w = NewWorker(
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			time.Sleep(200 * time.Millisecond)
+			return nil
+		}),
+	)
+
+	done := make(chan error)
+	go func() {
+		done <- w.handle(job)
+	}()
+
+	assert.NoError(t, w.Shutdown())
+
+	err = <-done
+	assert.Error(t, err)
+	assert.Equal(t, context.DeadlineExceeded, err)
+}
+
+func TestJobComplete(t *testing.T) {
+	job := queue.Job{
+		Timeout: 100 * time.Millisecond,
+		Body:    []byte("foo"),
+	}
+	w := NewWorker(
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			return errors.New("job completed")
+		}),
+	)
+
+	err := w.handle(job)
+	assert.Error(t, err)
+	assert.Equal(t, errors.New("job completed"), err)
+
+	job = queue.Job{
+		Timeout: 250 * time.Millisecond,
+		Body:    []byte("foo"),
+	}
+
+	w = NewWorker(
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			time.Sleep(200 * time.Millisecond)
+			return errors.New("job completed")
+		}),
+	)
+
+	done := make(chan error)
+	go func() {
+		done <- w.handle(job)
+	}()
+
+	assert.NoError(t, w.Shutdown())
+
+	err = <-done
+	assert.Error(t, err)
+	assert.Equal(t, errors.New("job completed"), err)
+}