feat(metrics): add nsq queue stats (#14)

Author: appleboy

go.mod

@@ -3,7 +3,7 @@ module github.com/golang-queue/nsq
 go 1.18
 
 require (
-	github.com/golang-queue/queue v0.0.13-0.20220330060848-d1a0d31ce747
+	github.com/golang-queue/queue v0.0.13-0.20220403053215-9d5c5466dc9a
 	github.com/nsqio/go-nsq v1.1.0
 	github.com/stretchr/testify v1.7.0
 	go.uber.org/goleak v1.1.12

go.sum

@@ -2,6 +2,8 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/golang-queue/queue v0.0.13-0.20220330060848-d1a0d31ce747 h1:uNTbCoWORAcna89KcKgP22WFGv5fsij05e70DCnLrUU=
 github.com/golang-queue/queue v0.0.13-0.20220330060848-d1a0d31ce747/go.mod h1:ku8iyjYffqYY6Duts+xl+QYfN3/KDK4MEvXMZUkHyio=
+github.com/golang-queue/queue v0.0.13-0.20220403053215-9d5c5466dc9a h1:XeBmQF8PCff5DwphR2eZPGKzVnT8/34/JK1IUcSoSpc=
+github.com/golang-queue/queue v0.0.13-0.20220403053215-9d5c5466dc9a/go.mod h1:ku8iyjYffqYY6Duts+xl+QYfN3/KDK4MEvXMZUkHyio=
 github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=

nsq.go

@@ -69,12 +69,6 @@ func (w *Worker) startConsumer(cfg *nsq.Config) error {
 	}
 
 	w.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
-		// re-queue the job if worker has been shutdown.
-		if atomic.LoadInt32(&w.stopFlag) == 1 {
-			msg.Requeue(-1)
-			return nil
-		}
-
 		if len(msg.Body) == 0 {
 			// Returning nil will automatically send a FIN command to NSQ to mark the message as processed.
 			// In this case, a message with an empty body is simply ignored/discarded.
@@ -86,6 +80,7 @@ func (w *Worker) startConsumer(cfg *nsq.Config) error {
 		case <-w.stop:
 			if msg != nil {
 				// re-queue the job if worker has been shutdown.
+				w.opts.logger.Info("re-queue the old job")
 				msg.Requeue(-1)
 			}
 		}
@@ -219,3 +214,12 @@ loop:
 
 	return nil, queue.ErrNoTaskInQueue
 }
+
+// Stats retrieves the current connection and message statistics for a Consumer
+func (w *Worker) Stats() *nsq.ConsumerStats {
+	if w.q == nil {
+		return nil
+	}
+
+	return w.q.Stats()
+}

nsq_test.go

@@ -347,3 +347,52 @@ func TestJobComplete(t *testing.T) {
 	assert.Error(t, err)
 	assert.Equal(t, errors.New("job completed"), err)
 }
+
+func TestNSQStatsinQueue(t *testing.T) {
+	m := mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithAddr(host+":4150"),
+		WithTopic("nsq_stats"),
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			log.Println("get message")
+			return nil
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(1),
+	)
+	assert.NoError(t, err)
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	q.Start()
+	assert.Equal(t, int(1), w.Stats().Connections)
+	time.Sleep(500 * time.Millisecond)
+	assert.Equal(t, uint64(2), w.Stats().MessagesReceived)
+	assert.Equal(t, uint64(2), w.Stats().MessagesFinished)
+	q.Release()
+	assert.Equal(t, int(0), w.Stats().Connections)
+}
+
+func TestNSQStatsInWorker(t *testing.T) {
+	m := mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithAddr(host+":4150"),
+		WithTopic("nsq_stats_queue"),
+	)
+
+	assert.NoError(t, w.Queue(m))
+	assert.NoError(t, w.Queue(m))
+	assert.NoError(t, w.Queue(m))
+	assert.Equal(t, int(1), w.Stats().Connections)
+	time.Sleep(300 * time.Millisecond)
+	assert.Equal(t, uint64(3), w.Stats().MessagesReceived)
+	assert.Equal(t, uint64(1), w.Stats().MessagesFinished)
+	assert.Equal(t, uint64(0), w.Stats().MessagesRequeued)
+
+	_ = w.Shutdown()
+}