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

* chore(metric): support busy workers count.

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

* chore: update testing

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

Author: appleboy

consumer.go

@@ -15,12 +15,25 @@ var errMaxCapacity = errors.New("max capacity reached")
 
 // Worker for simple queue using channel
 type Consumer struct {
-	taskQueue chan QueuedMessage
-	runFunc   func(context.Context, QueuedMessage) error
-	stop      chan struct{}
-	logger    Logger
-	stopOnce  sync.Once
-	stopFlag  int32
+	taskQueue   chan QueuedMessage
+	runFunc     func(context.Context, QueuedMessage) error
+	stop        chan struct{}
+	logger      Logger
+	stopOnce    sync.Once
+	stopFlag    int32
+	busyWorkers uint64
+}
+
+func (s *Consumer) incBusyWorker() {
+	atomic.AddUint64(&s.busyWorkers, 1)
+}
+
+func (s *Consumer) decBusyWorker() {
+	atomic.AddUint64(&s.busyWorkers, ^uint64(0))
+}
+
+func (s *Consumer) BusyWorkers() uint64 {
+	return atomic.LoadUint64(&s.busyWorkers)
 }
 
 // BeforeRun run script before start worker
@@ -39,7 +52,11 @@ func (s *Consumer) handle(job Job) error {
 	panicChan := make(chan interface{}, 1)
 	startTime := time.Now()
 	ctx, cancel := context.WithTimeout(context.Background(), job.Timeout)
-	defer cancel()
+	s.incBusyWorker()
+	defer func() {
+		cancel()
+		s.decBusyWorker()
+	}()
 
 	// run the job
 	go func() {

consumer_test.go

@@ -357,3 +357,30 @@ func TestTaskJobComplete(t *testing.T) {
 	}
 	assert.Equal(t, context.DeadlineExceeded, w.handle(job))
 }
+
+func TestBusyWorkerCount(t *testing.T) {
+	job := Job{
+		Timeout: 200 * time.Millisecond,
+		Task: func(ctx context.Context) error {
+			time.Sleep(100 * time.Millisecond)
+			return nil
+		},
+	}
+
+	w := NewConsumer()
+
+	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(100 * time.Millisecond)
+	assert.Equal(t, uint64(0), w.BusyWorkers())
+
+	assert.NoError(t, w.Shutdown())
+}

queue_test.go

@@ -29,6 +29,7 @@ func TestNewQueue(t *testing.T) {
 	assert.NotNil(t, q)
 
 	q.Start()
+	assert.Equal(t, uint64(0), w.BusyWorkers())
 	q.Shutdown()
 	q.Wait()
 }
@@ -48,6 +49,7 @@ func TestWorkerNum(t *testing.T) {
 	q.Start()
 	time.Sleep(20 * time.Millisecond)
 	assert.Equal(t, 4, q.Workers())
+	assert.Equal(t, uint64(0), w.BusyWorkers())
 	q.Shutdown()
 	q.Wait()
 }
@@ -203,6 +205,7 @@ func TestQueueTaskJob(t *testing.T) {
 		return nil
 	}))
 	time.Sleep(50 * time.Millisecond)
+	assert.Equal(t, uint64(0), w.BusyWorkers())
 	q.Shutdown()
 	assert.Equal(t, ErrQueueShutdown, q.QueueTask(func(ctx context.Context) error {
 		return nil

worker.go

@@ -16,6 +16,8 @@ type Worker interface {
 	Capacity() int
 	// Usage is how many message in queue
 	Usage() int
+	// BusyWorkers return count of busy worker currently
+	BusyWorkers() uint64
 }
 
 // QueuedMessage ...

worker_empty.go

@@ -12,3 +12,4 @@ func (w *emptyWorker) Shutdown() error               { return nil }
 func (w *emptyWorker) Queue(job QueuedMessage) error { return nil }
 func (w *emptyWorker) Capacity() int                 { return 0 }
 func (w *emptyWorker) Usage() int                    { return 0 }
+func (w *emptyWorker) BusyWorkers() uint64           { return uint64(0) }

worker_message.go

@@ -37,5 +37,6 @@ func (w *messageWorker) Queue(job QueuedMessage) error {
 		return errors.New("max capacity reached")
 	}
 }
-func (w *messageWorker) Capacity() int { return cap(w.messages) }
-func (w *messageWorker) Usage() int    { return len(w.messages) }
+func (w *messageWorker) Capacity() int       { return cap(w.messages) }
+func (w *messageWorker) Usage() int          { return len(w.messages) }
+func (w *messageWorker) BusyWorkers() uint64 { return uint64(0) }

worker_task.go

@@ -38,5 +38,6 @@ func (w *taskWorker) Queue(job QueuedMessage) error {
 		return errors.New("max capacity reached")
 	}
 }
-func (w *taskWorker) Capacity() int { return cap(w.messages) }
-func (w *taskWorker) Usage() int    { return len(w.messages) }
+func (w *taskWorker) Capacity() int       { return cap(w.messages) }
+func (w *taskWorker) Usage() int          { return len(w.messages) }
+func (w *taskWorker) BusyWorkers() uint64 { return uint64(0) }