chore: close queue after shutdown. (#22)

Author: appleboy

queue.go

@@ -128,6 +128,10 @@ func (q *Queue) Wait() {
 
 // Queue to queue all job
 func (q *Queue) Queue(job QueuedMessage) error {
+	if atomic.LoadInt32(&q.stopFlag) == 1 {
+		return ErrQueueShutdown
+	}
+
 	return q.worker.Queue(Job{
 		Timeout: q.timeout,
 		Body:    job.Bytes(),
@@ -136,6 +140,10 @@ func (q *Queue) Queue(job QueuedMessage) error {
 
 // Queue to queue all job
 func (q *Queue) QueueWithTimeout(timeout time.Duration, job QueuedMessage) error {
+	if atomic.LoadInt32(&q.stopFlag) == 1 {
+		return ErrQueueShutdown
+	}
+
 	return q.worker.Queue(Job{
 		Timeout: timeout,
 		Body:    job.Bytes(),

queue_test.go

@@ -146,3 +146,31 @@ func TestCapacityReached(t *testing.T) {
 		message: "foobar",
 	}))
 }
+
+func TestCloseQueueAfterShutdown(t *testing.T) {
+	w := &queueWorker{
+		messages: make(chan QueuedMessage, 10),
+	}
+	q, err := NewQueue(
+		WithWorker(w),
+		WithWorkerCount(5),
+		WithLogger(NewEmptyLogger()),
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, q)
+
+	assert.NoError(t, q.Queue(mockMessage{
+		message: "foobar",
+	}))
+	q.Shutdown()
+	err = q.Queue(mockMessage{
+		message: "foobar",
+	})
+	assert.Error(t, err)
+	assert.Equal(t, ErrQueueShutdown, err)
+	err = q.QueueWithTimeout(10*time.Millisecond, mockMessage{
+		message: "foobar",
+	})
+	assert.Error(t, err)
+	assert.Equal(t, ErrQueueShutdown, err)
+}