[pkg/queue] Add StartConsumersWithFactory function

pkg/queue/bounded_queue.go

@@ -25,6 +25,11 @@ import (
 	uatomic "go.uber.org/atomic"
 )
 
+// Consumer consumes data from a bounded queue
+type Consumer interface {
+	Consume(item interface{})
+}
+
 // BoundedQueue implements a producer-consumer exchange similar to a ring buffer queue,
 // where the queue is bounded and if it fills up due to slow consumers, the new items written by
 // the producer force the earliest items to be dropped. The implementation is actually based on
@@ -38,7 +43,7 @@ type BoundedQueue struct {
 	stopped       *uatomic.Uint32
 	items         *chan interface{}
 	onDroppedItem func(item interface{})
-	consumer      func(item interface{})
+	factory       func() Consumer
 	stopCh        chan struct{}
 }
 
@@ -56,25 +61,26 @@ func NewBoundedQueue(capacity int, onDroppedItem func(item interface{})) *Bounde
 	}
 }
 
-// StartConsumers starts a given number of goroutines consuming items from the queue
-// and passing them into the consumer callback.
-func (q *BoundedQueue) StartConsumers(num int, consumer func(item interface{})) {
+// StartConsumersWithFactory creates a given number of consumers consuming items
+// from the queue in separate goroutines.
+func (q *BoundedQueue) StartConsumersWithFactory(num int, factory func() Consumer) {
 	q.workers = num
-	q.consumer = consumer
+	q.factory = factory
 	var startWG sync.WaitGroup
 	for i := 0; i < q.workers; i++ {
 		q.stopWG.Add(1)
 		startWG.Add(1)
 		go func() {
 			startWG.Done()
 			defer q.stopWG.Done()
+			consumer := q.factory()
 			queue := *q.items
 			for {
 				select {
 				case item, ok := <-queue:
 					if ok {
 						q.size.Sub(1)
-						q.consumer(item)
+						consumer.Consume(item)
 					} else {
 						// channel closed, finish worker
 						return
@@ -89,6 +95,25 @@ func (q *BoundedQueue) StartConsumers(num int, consumer func(item interface{}))
 	startWG.Wait()
 }
 
+// statelessConsumer wraps a consume function callback
+type statelessConsumer struct {
+	consumefn func(item interface{})
+}
+
+// Consumer consumes an item from a bounded queue
+func (c *statelessConsumer) Consume(item interface{}) {
+	c.consumefn(item)
+}
+
+// StartConsumers starts a given number of goroutines consuming items from the queue
+// and passing them into the consumer callback.
+func (q *BoundedQueue) StartConsumers(num int, callback func(item interface{})) {
+	consumer := &statelessConsumer{callback}
+	q.StartConsumersWithFactory(num, func() Consumer {
+		return consumer
+	})
+}
+
 // Produce is used by the producer to submit new item to the queue. Returns false in case of queue overflow.
 func (q *BoundedQueue) Produce(item interface{}) bool {
 	if q.stopped.Load() != 0 {
@@ -171,7 +196,7 @@ func (q *BoundedQueue) Resize(capacity int) bool {
 	swapped := atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(&q.items)), unsafe.Pointer(q.items), unsafe.Pointer(&queue))
 	if swapped {
 		// start a new set of consumers, based on the information given previously
-		q.StartConsumers(q.workers, q.consumer)
+		q.StartConsumersWithFactory(q.workers, q.factory)
 
 		// gracefully drain the existing queue
 		close(previous)

pkg/queue/bounded_queue_test.go

@@ -30,33 +30,14 @@ import (
 	uatomic "go.uber.org/atomic"
 )
 
-// In this test we run a queue with capacity 1 and a single consumer.
-// We want to test the overflow behavior, so we block the consumer
-// by holding a startLock before submitting items to the queue.
-func TestBoundedQueue(t *testing.T) {
-	mFact := metricstest.NewFactory(0)
-	counter := mFact.Counter(metrics.Options{Name: "dropped", Tags: nil})
+func checkQueue(
+	t *testing.T,
+	q *BoundedQueue,
+	startLock *sync.Mutex,
+	consumerState *consumerState,
+	mFact *metricstest.Factory,
+) {
 	gauge := mFact.Gauge(metrics.Options{Name: "size", Tags: nil})
-
-	q := NewBoundedQueue(1, func(item interface{}) {
-		counter.Inc(1)
-	})
-	assert.Equal(t, 1, q.Capacity())
-
-	var startLock sync.Mutex
-
-	startLock.Lock() // block consumers
-	consumerState := newConsumerState(t)
-
-	q.StartConsumers(1, func(item interface{}) {
-		consumerState.record(item.(string))
-
-		// block further processing until startLock is released
-		startLock.Lock()
-		//lint:ignore SA2001 empty section is ok
-		startLock.Unlock()
-	})
-
 	assert.True(t, q.Produce("a"))
 
 	// at this point "a" may or may not have been received by the consumer go-routine
@@ -112,6 +93,58 @@ func TestBoundedQueue(t *testing.T) {
 	assert.False(t, q.Produce("x"), "cannot push to closed queue")
 }
 
+// In this test we run a queue with capacity 1 and a single consumer.
+// We want to test the overflow behavior, so we block the consumer
+// by holding a startLock before submitting items to the queue.
+func TestBoundedQueue(t *testing.T) {
+	mFact := metricstest.NewFactory(0)
+	counter := mFact.Counter(metrics.Options{Name: "dropped", Tags: nil})
+
+	q := NewBoundedQueue(1, func(item interface{}) {
+		counter.Inc(1)
+	})
+	assert.Equal(t, 1, q.Capacity())
+
+	var startLock sync.Mutex
+
+	startLock.Lock() // block consumers
+	consumerState := newConsumerState(t)
+
+	q.StartConsumers(1, func(item interface{}) {
+		consumerState.record(item.(string))
+
+		// block further processing until startLock is released
+		startLock.Lock()
+		//lint:ignore SA2001 empty section is ok
+		startLock.Unlock()
+	})
+
+	checkQueue(t, q, &startLock, consumerState, mFact)
+}
+
+// This test is identical to the previous one but we start the
+// queue using a consumer factory instead of a callback.
+func TestBoundedQueueWithFactory(t *testing.T) {
+	mFact := metricstest.NewFactory(0)
+	counter := mFact.Counter(metrics.Options{Name: "dropped", Tags: nil})
+
+	q := NewBoundedQueue(1, func(item interface{}) {
+		counter.Inc(1)
+	})
+	assert.Equal(t, 1, q.Capacity())
+
+	var startLock sync.Mutex
+
+	startLock.Lock() // block consumers
+	consumerState := newConsumerState(t)
+
+	q.StartConsumersWithFactory(1, func() Consumer {
+		return newStatefulConsumer(&startLock, consumerState)
+	})
+
+	checkQueue(t, q, &startLock, consumerState, mFact)
+}
+
 type consumerState struct {
 	sync.Mutex
 	t            *testing.T
@@ -161,6 +194,24 @@ func (s *consumerState) assertConsumed(expected map[string]bool) {
 	assert.Equal(s.t, expected, s.snapshot())
 }
 
+type statefulConsumer struct {
+	*sync.Mutex
+	*consumerState
+}
+
+func (s *statefulConsumer) Consume(item interface{}) {
+	s.record(item.(string))
+
+	// block further processing until the lock is released
+	s.Lock()
+	//lint:ignore SA2001 empty section is ok
+	s.Unlock()
+}
+
+func newStatefulConsumer(startLock *sync.Mutex, cs *consumerState) Consumer {
+	return &statefulConsumer{startLock, cs}
+}
+
 func TestResizeUp(t *testing.T) {
 	q := NewBoundedQueue(2, func(item interface{}) {
 		fmt.Printf("dropped: %v\n", item)
@@ -332,3 +383,21 @@ func BenchmarkBoundedQueue(b *testing.B) {
 		q.Produce(n)
 	}
 }
+
+// nopConsumer is a no-op consumer
+type nopConsumer struct{}
+
+func (*nopConsumer) Consume(item interface{}) {}
+
+func BenchmarkBoundedQueueWithFactory(b *testing.B) {
+	q := NewBoundedQueue(1000, func(item interface{}) {
+	})
+
+	q.StartConsumersWithFactory(10, func() Consumer {
+		return &nopConsumer{}
+	})
+
+	for n := 0; n < b.N; n++ {
+		q.Produce(n)
+	}
+}