fix(NSQ): close service after waiting all job completed. (#601)

queue/nsq/nsq.go

@@ -98,7 +98,10 @@ func (s *Worker) AfterRun() error {
 
 // Run start the worker
 func (s *Worker) Run(quit chan struct{}) error {
+	wg := &sync.WaitGroup{}
 	s.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
+		wg.Add(1)
+		defer wg.Done()
 		var notification gorush.PushNotification
 		if err := json.Unmarshal(msg.Body, &notification); err != nil {
 			return err
@@ -110,6 +113,8 @@ func (s *Worker) Run(quit chan struct{}) error {
 	select {
 	case <-quit:
 	}
+
+	wg.Wait()
 	return nil
 }
 

queue/nsq/nsq_test.go

@@ -0,0 +1,34 @@
+package nsq
+
+import (
+	"log"
+	"testing"
+	"time"
+
+	"github.com/appleboy/gorush/logx"
+	"github.com/appleboy/gorush/queue"
+)
+
+func TestMain(m *testing.M) {
+	if err := logx.InitLog(
+		"debug",
+		"stdout",
+		"debug",
+		"stdout",
+	); err != nil {
+		log.Fatalf("Can't load log module, error: %v", err)
+	}
+
+	m.Run()
+}
+
+func TestShutdown(t *testing.T) {
+	w := NewWorker(
+		WithAddr("nsq:4150"),
+		WithTopic("test"),
+	)
+	q := queue.NewQueue(w, 2)
+	q.Start()
+	time.Sleep(1 * time.Second)
+	q.Shutdown()
+}

queue/queue.go

@@ -64,7 +64,9 @@ func (q *Queue) Queue(job interface{}) error {
 }
 
 func (q *Queue) work(num int) {
-	q.worker.BeforeRun()
+	if err := q.worker.BeforeRun(); err != nil {
+		logx.LogError.Fatal(err)
+	}
 	q.routineGroup.Run(func() {
 		// to handle panic cases from inside the worker
 		// in such case, we start a new goroutine
@@ -79,7 +81,9 @@ func (q *Queue) work(num int) {
 		q.worker.Run(q.quit)
 		logx.LogAccess.Info("closed the worker num ", num)
 	})
-	q.worker.AfterRun()
+	if err := q.worker.AfterRun(); err != nil {
+		logx.LogError.Fatal(err)
+	}
 }
 
 func (q *Queue) startWorker() {