chore(queue): cancel the channel.

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

Author: appleboy

rabbitmq.go

@@ -20,7 +20,6 @@ type Worker struct {
 	client    *amqp.Connection
 	channel   *amqp.Channel
 	stop      chan struct{}
-	exit      chan struct{}
 	stopFlag  int32
 	stopOnce  sync.Once
 	startOnce sync.Once
@@ -34,7 +33,6 @@ func NewWorker(opts ...Option) *Worker {
 	w := &Worker{
 		opts:  newOptions(opts...),
 		stop:  make(chan struct{}),
-		exit:  make(chan struct{}),
 		tasks: make(chan amqp.Delivery),
 	}
 
@@ -151,10 +149,6 @@ func (w *Worker) Shutdown() error {
 
 	w.stopOnce.Do(func() {
 		close(w.stop)
-		select {
-		case <-w.exit:
-		case <-time.After(50 * time.Millisecond):
-		}
 		if err := w.channel.Cancel("", true); err != nil {
 			w.opts.logger.Error(err)
 		}

rabbitmq_test.go

@@ -1,9 +1,16 @@
 package rabbitmq
 
 import (
+	"context"
+	"errors"
+	"fmt"
+	"log"
+	"runtime"
 	"testing"
+	"time"
 
 	"github.com/golang-queue/queue"
+	"github.com/golang-queue/queue/core"
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/goleak"
@@ -38,3 +45,186 @@ func TestDefaultWorkFlow(t *testing.T) {
 	q.Start()
 	q.Release()
 }
+
+func TestShutdownWorkFlow(t *testing.T) {
+	w := NewWorker(
+		WithSubj("test"),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(1 * time.Second)
+	q.Shutdown()
+	// check shutdown once
+	q.Shutdown()
+	q.Wait()
+}
+
+func TestCustomFuncAndWait(t *testing.T) {
+	m := &mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithSubj("test"),
+		WithRunFunc(func(ctx context.Context, m core.QueuedMessage) error {
+			log.Println("show message: " + string(m.Bytes()))
+			time.Sleep(500 * time.Millisecond)
+			return nil
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(100 * time.Millisecond)
+	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()
+	// you will see the execute time > 1000ms
+}
+
+func TestEnqueueJobAfterShutdown(t *testing.T) {
+	m := mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker()
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(50 * time.Millisecond)
+	q.Shutdown()
+	// can't queue task after shutdown
+	err = q.Queue(m)
+	assert.Error(t, err)
+	assert.Equal(t, queue.ErrQueueShutdown, err)
+	q.Wait()
+}
+
+func TestJobReachTimeout(t *testing.T) {
+	m := mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithSubj("JobReachTimeout"),
+		WithRunFunc(func(ctx context.Context, m core.QueuedMessage) error {
+			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")
+					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+						log.Println("job deadline exceeded")
+					}
+					return nil
+				default:
+				}
+				time.Sleep(50 * time.Millisecond)
+			}
+		}),
+	)
+	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.QueueWithTimeout(20*time.Millisecond, m))
+	time.Sleep(100 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+}
+
+func TestCancelJobAfterShutdown(t *testing.T) {
+	m := mockMessage{
+		Message: "test",
+	}
+	w := NewWorker(
+		WithSubj("CancelJob"),
+		WithLogger(queue.NewLogger()),
+		WithRunFunc(func(ctx context.Context, m core.QueuedMessage) error {
+			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")
+					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+						log.Println("job deadline exceeded")
+					}
+					return nil
+				default:
+				}
+				time.Sleep(50 * time.Millisecond)
+			}
+		}),
+	)
+	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.QueueWithTimeout(150*time.Millisecond, m))
+	time.Sleep(100 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+}
+
+func TestGoroutineLeak(t *testing.T) {
+	m := mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithSubj("GoroutineLeak"),
+		WithLogger(queue.NewEmptyLogger()),
+		WithRunFunc(func(ctx context.Context, m core.QueuedMessage) error {
+			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")
+					} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+						log.Println("job deadline exceeded")
+					}
+					return nil
+				default:
+					log.Println("get data:", string(m.Bytes()))
+					time.Sleep(50 * time.Millisecond)
+					return nil
+				}
+			}
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithLogger(queue.NewEmptyLogger()),
+		queue.WithWorker(w),
+		queue.WithWorkerCount(10),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(50 * time.Millisecond)
+	for i := 0; i < 500; i++ {
+		m.Message = fmt.Sprintf("foobar: %d", i+1)
+		assert.NoError(t, q.Queue(m))
+	}
+	time.Sleep(200 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+	fmt.Println("number of goroutines:", runtime.NumGoroutine())
+}