chore(metric): Add metric interface (#40)

Author: appleboy

consumer.go

@@ -15,25 +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
-	busyWorkers uint64
+	taskQueue chan QueuedMessage
+	runFunc   func(context.Context, QueuedMessage) error
+	stop      chan struct{}
+	logger    Logger
+	stopOnce  sync.Once
+	stopFlag  int32
+	metric    Metric
 }
 
 func (s *Consumer) incBusyWorker() {
-	atomic.AddUint64(&s.busyWorkers, 1)
+	s.metric.IncBusyWorker()
 }
 
 func (s *Consumer) decBusyWorker() {
-	atomic.AddUint64(&s.busyWorkers, ^uint64(0))
+	s.metric.DecBusyWorker()
 }
 
 func (s *Consumer) BusyWorkers() uint64 {
-	return atomic.LoadUint64(&s.busyWorkers)
+	return s.metric.BusyWorkers()
 }
 
 // BeforeRun run script before start worker
@@ -168,6 +168,7 @@ func NewConsumer(opts ...Option) *Consumer {
 		stop:      make(chan struct{}),
 		logger:    o.logger,
 		runFunc:   o.fn,
+		metric:    o.metric,
 	}
 
 	return w

metric.go

@@ -0,0 +1,30 @@
+package queue
+
+import "sync/atomic"
+
+// Metric interface
+type Metric interface {
+	IncBusyWorker()
+	DecBusyWorker()
+	BusyWorkers() uint64
+}
+
+type metric struct {
+	busyWorkers uint64
+}
+
+func newMetric() Metric {
+	return &metric{}
+}
+
+func (m *metric) IncBusyWorker() {
+	atomic.AddUint64(&m.busyWorkers, 1)
+}
+
+func (m *metric) DecBusyWorker() {
+	atomic.AddUint64(&m.busyWorkers, ^uint64(0))
+}
+
+func (m *metric) BusyWorkers() uint64 {
+	return atomic.LoadUint64(&m.busyWorkers)
+}

options.go

@@ -12,6 +12,7 @@ var (
 	defaultTimeout     = 60 * time.Minute
 	defaultNewLogger   = NewLogger()
 	defaultFn          = func(context.Context, QueuedMessage) error { return nil }
+	defaultMetric      = newMetric()
 )
 
 // Option for queue system
@@ -38,6 +39,13 @@ func WithLogger(l Logger) Option {
 	}
 }
 
+// WithMetric set custom Metric
+func WithMetric(m Metric) Option {
+	return func(q *Options) {
+		q.metric = m
+	}
+}
+
 // WithWorker set custom worker
 func WithWorker(w Worker) Option {
 	return func(q *Options) {
@@ -66,6 +74,7 @@ type Options struct {
 	queueSize   int
 	worker      Worker
 	fn          func(context.Context, QueuedMessage) error
+	metric      Metric
 }
 
 func NewOptions(opts ...Option) *Options {
@@ -76,6 +85,7 @@ func NewOptions(opts ...Option) *Options {
 		logger:      defaultNewLogger,
 		worker:      nil,
 		fn:          defaultFn,
+		metric:      defaultMetric,
 	}
 
 	// Loop through each option