chore(metric): support busy workers count. (#3)

Author: appleboy

go.mod

@@ -3,7 +3,7 @@ module github.com/golang-queue/nsq
 go 1.16
 
 require (
-	github.com/golang-queue/queue v0.0.7
+	github.com/golang-queue/queue v0.0.8-0.20210905085819-3cd1dfe014e2
 	github.com/nsqio/go-nsq v1.0.8
 	github.com/stretchr/testify v1.7.0
 )

go.sum

@@ -1,7 +1,7 @@
 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.7 h1:WENCPyPBcIWYgBFSAZ8USGtwmxeCeMkhjwuyM0MAi84=
-github.com/golang-queue/queue v0.0.7/go.mod h1:JS5tYJacahCjafcplU5idNLX2vkYioqh6wEDX5o9Nms=
+github.com/golang-queue/queue v0.0.8-0.20210905085819-3cd1dfe014e2 h1:C2nn3ZHokfvJ787A/IiUDgweK5QZa5RGm/tfvhNfOtg=
+github.com/golang-queue/queue v0.0.8-0.20210905085819-3cd1dfe014e2/go.mod h1:JS5tYJacahCjafcplU5idNLX2vkYioqh6wEDX5o9Nms=
 github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/nsqio/go-nsq v1.0.8 h1:3L2F8tNLlwXXlp2slDUrUWSBn2O3nMh8R1/KEDFTHPk=

nsq.go

@@ -33,6 +33,19 @@ type Worker struct {
 	logger      queue.Logger
 	stopFlag    int32
 	startFlag   int32
+	busyWorkers uint64
+}
+
+func (w *Worker) incBusyWorker() {
+	atomic.AddUint64(&w.busyWorkers, 1)
+}
+
+func (w *Worker) decBusyWorker() {
+	atomic.AddUint64(&w.busyWorkers, ^uint64(0))
+}
+
+func (w *Worker) BusyWorkers() uint64 {
+	return atomic.LoadUint64(&w.busyWorkers)
 }
 
 // WithAddr setup the addr of NSQ
@@ -114,32 +127,36 @@ func NewWorker(opts ...Option) *Worker {
 }
 
 // BeforeRun run script before start worker
-func (s *Worker) BeforeRun() error {
+func (w *Worker) BeforeRun() error {
 	return nil
 }
 
 // AfterRun run script after start worker
-func (s *Worker) AfterRun() error {
-	s.startOnce.Do(func() {
+func (w *Worker) AfterRun() error {
+	w.startOnce.Do(func() {
 		time.Sleep(100 * time.Millisecond)
-		err := s.q.ConnectToNSQD(s.addr)
+		err := w.q.ConnectToNSQD(w.addr)
 		if err != nil {
 			panic("Could not connect nsq server: " + err.Error())
 		}
 
-		atomic.CompareAndSwapInt32(&s.startFlag, 0, 1)
+		atomic.CompareAndSwapInt32(&w.startFlag, 0, 1)
 	})
 
 	return nil
 }
 
-func (s *Worker) handle(job queue.Job) error {
+func (w *Worker) handle(job queue.Job) error {
 	// create channel with buffer size 1 to avoid goroutine leak
 	done := make(chan error, 1)
 	panicChan := make(chan interface{}, 1)
 	startTime := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), job.Timeout)
-	defer cancel()
+	w.incBusyWorker()
+	defer func() {
+		cancel()
+		w.decBusyWorker()
+	}()
 
 	// run the job
 	go func() {
@@ -151,15 +168,15 @@ func (s *Worker) handle(job queue.Job) error {
 		}()
 
 		// run custom process function
-		done <- s.runFunc(ctx, job)
+		done <- w.runFunc(ctx, job)
 	}()
 
 	select {
 	case p := <-panicChan:
 		panic(p)
 	case <-ctx.Done(): // timeout reached
 		return ctx.Err()
-	case <-s.stop: // shutdown service
+	case <-w.stop: // shutdown service
 		// cancel job
 		cancel()
 
@@ -179,10 +196,10 @@ func (s *Worker) handle(job queue.Job) error {
 }
 
 // Run start the worker
-func (s *Worker) Run() error {
+func (w *Worker) Run() error {
 	wg := &sync.WaitGroup{}
 	panicChan := make(chan interface{}, 1)
-	s.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
+	w.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
 		wg.Add(1)
 		defer func() {
 			wg.Done()
@@ -198,14 +215,14 @@ func (s *Worker) Run() error {
 
 		var data queue.Job
 		_ = json.Unmarshal(msg.Body, &data)
-		return s.handle(data)
+		return w.handle(data)
 	}))
 
 	// wait close signal
 	select {
-	case <-s.stop:
+	case <-w.stop:
 	case err := <-panicChan:
-		s.logger.Error(err)
+		w.logger.Error(err)
 	}
 
 	// wait job completed
@@ -215,39 +232,39 @@ func (s *Worker) Run() error {
 }
 
 // Shutdown worker
-func (s *Worker) Shutdown() error {
-	if !atomic.CompareAndSwapInt32(&s.stopFlag, 0, 1) {
+func (w *Worker) Shutdown() error {
+	if !atomic.CompareAndSwapInt32(&w.stopFlag, 0, 1) {
 		return queue.ErrQueueShutdown
 	}
 
-	s.stopOnce.Do(func() {
-		if atomic.LoadInt32(&s.startFlag) == 1 {
-			s.q.Stop()
-			s.p.Stop()
+	w.stopOnce.Do(func() {
+		if atomic.LoadInt32(&w.startFlag) == 1 {
+			w.q.Stop()
+			w.p.Stop()
 		}
 
-		close(s.stop)
+		close(w.stop)
 	})
 	return nil
 }
 
 // Capacity for channel
-func (s *Worker) Capacity() int {
+func (w *Worker) Capacity() int {
 	return 0
 }
 
 // Usage for count of channel usage
-func (s *Worker) Usage() int {
+func (w *Worker) Usage() int {
 	return 0
 }
 
 // Queue send notification to queue
-func (s *Worker) Queue(job queue.QueuedMessage) error {
-	if atomic.LoadInt32(&s.stopFlag) == 1 {
+func (w *Worker) Queue(job queue.QueuedMessage) error {
+	if atomic.LoadInt32(&w.stopFlag) == 1 {
 		return queue.ErrQueueShutdown
 	}
 
-	err := s.p.Publish(s.topic, job.Bytes())
+	err := w.p.Publish(w.topic, job.Bytes())
 	if err != nil {
 		return err
 	}

nsq_test.go

@@ -365,3 +365,32 @@ func TestJobComplete(t *testing.T) {
 	assert.Error(t, err)
 	assert.Equal(t, errors.New("job completed"), err)
 }
+
+func TestBusyWorkerCount(t *testing.T) {
+	job := queue.Job{
+		Timeout: 500 * time.Millisecond,
+		Body:    []byte("foo"),
+	}
+
+	w := NewWorker(
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			time.Sleep(200 * time.Millisecond)
+			return nil
+		}),
+	)
+
+	assert.Equal(t, uint64(0), w.BusyWorkers())
+	go func() {
+		assert.NoError(t, w.handle(job))
+	}()
+	go func() {
+		assert.NoError(t, w.handle(job))
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+	assert.Equal(t, uint64(2), w.BusyWorkers())
+	time.Sleep(200 * time.Millisecond)
+	assert.Equal(t, uint64(0), w.BusyWorkers())
+
+	assert.NoError(t, w.Shutdown())
+}