[chore] Extract batcher worker pool, cleanup unit tests (#13164)

Signed-off-by: Bogdan Drutu <[email protected]>

Author: bogdandrutu

exporter/exporterhelper/internal/queuebatch/multi_batcher.go

@@ -21,7 +21,7 @@ type batcherSettings[T any] struct {
 
 type multiBatcher struct {
 	cfg         BatchConfig
-	wp          chan struct{}
+	wp          *workerPool
 	sizerType   request.SizerType
 	sizer       request.Sizer[request.Request]
 	partitioner Partitioner[request.Request]
@@ -34,16 +34,9 @@ type multiBatcher struct {
 var _ Batcher[request.Request] = (*multiBatcher)(nil)
 
 func newMultiBatcher(bCfg BatchConfig, bSet batcherSettings[request.Request]) *multiBatcher {
-	var workerPool chan struct{}
-	if bSet.maxWorkers != 0 {
-		workerPool = make(chan struct{}, bSet.maxWorkers)
-		for i := 0; i < bSet.maxWorkers; i++ {
-			workerPool <- struct{}{}
-		}
-	}
 	mb := &multiBatcher{
 		cfg:         bCfg,
-		wp:          workerPool,
+		wp:          newWorkerPool(bSet.maxWorkers),
 		sizerType:   bSet.sizerType,
 		sizer:       bSet.sizer,
 		partitioner: bSet.partitioner,
@@ -68,18 +61,18 @@ func (mb *multiBatcher) getShard(ctx context.Context, req request.Request) *shar
 		return s.(*shardBatcher)
 	}
 	newS := newShard(mb.cfg, mb.sizerType, mb.sizer, mb.wp, mb.consumeFunc)
-	newS.start(ctx, nil)
+	_ = newS.Start(ctx, nil)
 	s, loaded := mb.shards.LoadOrStore(key, newS)
 	// If not loaded, there was a race condition in adding the new shard. Shutdown the newly created shard.
 	if loaded {
-		newS.shutdown(ctx)
+		_ = newS.Shutdown(ctx)
 	}
 	return s.(*shardBatcher)
 }
 
 func (mb *multiBatcher) Start(ctx context.Context, host component.Host) error {
 	if mb.singleShard != nil {
-		mb.singleShard.start(ctx, host)
+		return mb.singleShard.Start(ctx, host)
 	}
 	return nil
 }
@@ -91,16 +84,15 @@ func (mb *multiBatcher) Consume(ctx context.Context, req request.Request, done D
 
 func (mb *multiBatcher) Shutdown(ctx context.Context) error {
 	if mb.singleShard != nil {
-		mb.singleShard.shutdown(ctx)
-		return nil
+		return mb.singleShard.Shutdown(ctx)
 	}
 
 	var wg sync.WaitGroup
 	mb.shards.Range(func(_ any, shard any) bool {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			shard.(*shardBatcher).shutdown(ctx)
+			_ = shard.(*shardBatcher).Shutdown(ctx)
 		}()
 		return true
 	})

exporter/exporterhelper/internal/queuebatch/shard_batcher.go

@@ -24,7 +24,7 @@ type batch struct {
 // shardBatcher continuously batch incoming requests and flushes asynchronously if minimum size limit is met or on timeout.
 type shardBatcher struct {
 	cfg            BatchConfig
-	workerPool     chan struct{}
+	wp             *workerPool
 	sizerType      request.SizerType
 	sizer          request.Sizer[request.Request]
 	consumeFunc    sender.SendFunc[request.Request]
@@ -35,10 +35,10 @@ type shardBatcher struct {
 	shutdownCh     chan struct{}
 }
 
-func newShard(cfg BatchConfig, sizerType request.SizerType, sizer request.Sizer[request.Request], workerPool chan struct{}, next sender.SendFunc[request.Request]) *shardBatcher {
+func newShard(cfg BatchConfig, sizerType request.SizerType, sizer request.Sizer[request.Request], wp *workerPool, next sender.SendFunc[request.Request]) *shardBatcher {
 	return &shardBatcher{
 		cfg:         cfg,
-		workerPool:  workerPool,
+		wp:          wp,
 		sizerType:   sizerType,
 		sizer:       sizer,
 		consumeFunc: next,
@@ -149,22 +149,33 @@ func (qb *shardBatcher) Consume(ctx context.Context, req request.Request, done D
 }
 
 // Start starts the goroutine that reads from the queue and flushes asynchronously.
-func (qb *shardBatcher) start(_ context.Context, _ component.Host) {
-	if qb.cfg.FlushTimeout > 0 {
-		qb.timer = time.NewTimer(qb.cfg.FlushTimeout)
-		qb.stopWG.Add(1)
-		go func() {
-			defer qb.stopWG.Done()
-			for {
-				select {
-				case <-qb.shutdownCh:
-					return
-				case <-qb.timer.C:
-					qb.flushCurrentBatchIfNecessary()
-				}
-			}
-		}()
+func (qb *shardBatcher) Start(context.Context, component.Host) error {
+	if qb.cfg.FlushTimeout <= 0 {
+		return nil
 	}
+	qb.timer = time.NewTimer(qb.cfg.FlushTimeout)
+	qb.stopWG.Add(1)
+	go func() {
+		defer qb.stopWG.Done()
+		for {
+			select {
+			case <-qb.shutdownCh:
+				return
+			case <-qb.timer.C:
+				qb.flushCurrentBatchIfNecessary()
+			}
+		}
+	}()
+	return nil
+}
+
+// Shutdown ensures that queue and all Batcher are stopped.
+func (qb *shardBatcher) Shutdown(context.Context) error {
+	close(qb.shutdownCh)
+	// Make sure execute one last flush if necessary.
+	qb.flushCurrentBatchIfNecessary()
+	qb.stopWG.Wait()
+	return nil
 }
 
 // flushCurrentBatchIfNecessary sends out the current request batch if it is not nil
@@ -186,24 +197,28 @@ func (qb *shardBatcher) flushCurrentBatchIfNecessary() {
 // flush starts a goroutine that calls consumeFunc. It blocks until a worker is available if necessary.
 func (qb *shardBatcher) flush(ctx context.Context, req request.Request, done Done) {
 	qb.stopWG.Add(1)
-	if qb.workerPool != nil {
-		<-qb.workerPool
-	}
-	go func() {
+	qb.wp.execute(func() {
 		defer qb.stopWG.Done()
 		done.OnDone(qb.consumeFunc(ctx, req))
-		if qb.workerPool != nil {
-			qb.workerPool <- struct{}{}
-		}
-	}()
+	})
 }
 
-// Shutdown ensures that queue and all Batcher are stopped.
-func (qb *shardBatcher) shutdown(_ context.Context) {
-	close(qb.shutdownCh)
-	// Make sure execute one last flush if necessary.
-	qb.flushCurrentBatchIfNecessary()
-	qb.stopWG.Wait()
+type workerPool struct {
+	workers chan struct{}
+}
+
+func newWorkerPool(maxWorkers int) *workerPool {
+	workers := make(chan struct{}, maxWorkers)
+	for i := 0; i < maxWorkers; i++ {
+		workers <- struct{}{}
+	}
+	return &workerPool{workers: workers}
+}
+
+func (wp *workerPool) execute(f func()) {
+	<-wp.workers
+	go f()
+	wp.workers <- struct{}{}
 }
 
 type multiDone []Done

exporter/exporterhelper/internal/queuebatch/shard_batcher_test.go

@@ -59,13 +59,7 @@ func TestShardBatcher_NoSplit_MinThresholdZero_TimeoutDisabled(t *testing.T) {
 			}
 
 			sink := requesttest.NewSink()
-			ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
-				sizerType:   tt.sizerType,
-				sizer:       tt.sizer,
-				partitioner: nil,
-				next:        sink.Export,
-				maxWorkers:  tt.maxWorkers,
-			})
+			ba := newShard(cfg, tt.sizerType, tt.sizer, newWorkerPool(tt.maxWorkers), sink.Export)
 			require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
 			t.Cleanup(func() {
 				require.NoError(t, ba.Shutdown(context.Background()))
@@ -83,7 +77,7 @@ func TestShardBatcher_NoSplit_MinThresholdZero_TimeoutDisabled(t *testing.T) {
 			assert.Eventually(t, func() bool {
 				return sink.RequestsCount() == 5 && (sink.ItemsCount() == 75 || sink.BytesCount() == 75)
 			}, 1*time.Second, 10*time.Millisecond)
-			// Check that done callback is called for the right amount of times.
+			// Check that done callback is called for the right number of times.
 			assert.EqualValues(t, 1, done.errors.Load())
 			assert.EqualValues(t, 5, done.success.Load())
 		})
@@ -130,13 +124,7 @@ func TestShardBatcher_NoSplit_TimeoutDisabled(t *testing.T) {
 			}
 
 			sink := requesttest.NewSink()
-			ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
-				sizerType:   tt.sizerType,
-				sizer:       tt.sizer,
-				partitioner: nil,
-				next:        sink.Export,
-				maxWorkers:  tt.maxWorkers,
-			})
+			ba := newShard(cfg, tt.sizerType, tt.sizer, newWorkerPool(tt.maxWorkers), sink.Export)
 			require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
 
 			done := newFakeDone()
@@ -165,7 +153,7 @@ func TestShardBatcher_NoSplit_TimeoutDisabled(t *testing.T) {
 			assert.Equal(t, 3, sink.RequestsCount())
 			assert.True(t, sink.ItemsCount() == 57 || sink.BytesCount() == 57)
 
-			// Check that done callback is called for the right amount of times.
+			// Check that done callback is called for the right number of times.
 			assert.EqualValues(t, 3, done.errors.Load())
 			assert.EqualValues(t, 4, done.success.Load())
 		})
@@ -216,13 +204,7 @@ func TestShardBatcher_NoSplit_WithTimeout(t *testing.T) {
 			}
 
 			sink := requesttest.NewSink()
-			ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
-				sizerType:   tt.sizerType,
-				sizer:       tt.sizer,
-				partitioner: nil,
-				next:        sink.Export,
-				maxWorkers:  tt.maxWorkers,
-			})
+			ba := newShard(cfg, tt.sizerType, tt.sizer, newWorkerPool(tt.maxWorkers), sink.Export)
 			require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
 			t.Cleanup(func() {
 				require.NoError(t, ba.Shutdown(context.Background()))
@@ -241,7 +223,7 @@ func TestShardBatcher_NoSplit_WithTimeout(t *testing.T) {
 				return sink.RequestsCount() == 1 && (sink.ItemsCount() == 75 || sink.BytesCount() == 75)
 			}, 1*time.Second, 10*time.Millisecond)
 
-			// Check that done callback is called for the right amount of times.
+			// Check that done callback is called for the right number of times.
 			assert.EqualValues(t, 1, done.errors.Load())
 			assert.EqualValues(t, 5, done.success.Load())
 		})
@@ -293,13 +275,7 @@ func TestShardBatcher_Split_TimeoutDisabled(t *testing.T) {
 			}
 
 			sink := requesttest.NewSink()
-			ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
-				sizerType:   tt.sizerType,
-				sizer:       tt.sizer,
-				partitioner: nil,
-				next:        sink.Export,
-				maxWorkers:  tt.maxWorkers,
-			})
+			ba := newShard(cfg, tt.sizerType, tt.sizer, newWorkerPool(tt.maxWorkers), sink.Export)
 			require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
 
 			done := newFakeDone()
@@ -332,7 +308,7 @@ func TestShardBatcher_Split_TimeoutDisabled(t *testing.T) {
 			assert.Equal(t, 11, sink.RequestsCount())
 			assert.True(t, sink.ItemsCount() == 1005 || sink.BytesCount() == 1005)
 
-			// Check that done callback is called for the right amount of times.
+			// Check that done callback is called for the right number of times.
 			assert.EqualValues(t, 2, done.errors.Load())
 			assert.EqualValues(t, 7, done.success.Load())
 		})
@@ -346,13 +322,7 @@ func TestShardBatcher_Shutdown(t *testing.T) {
 	}
 
 	sink := requesttest.NewSink()
-	ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
-		sizerType:   request.SizerTypeItems,
-		sizer:       request.NewItemsSizer(),
-		partitioner: nil,
-		next:        sink.Export,
-		maxWorkers:  2,
-	})
+	ba := newShard(cfg, request.SizerTypeItems, request.NewItemsSizer(), newWorkerPool(2), sink.Export)
 	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
 
 	done := newFakeDone()
@@ -367,7 +337,7 @@ func TestShardBatcher_Shutdown(t *testing.T) {
 	assert.Equal(t, 1, sink.RequestsCount())
 	assert.Equal(t, 3, sink.ItemsCount())
 
-	// Check that done callback is called for the right amount of times.
+	// Check that done callback is called for the right number of times.
 	assert.EqualValues(t, 0, done.errors.Load())
 	assert.EqualValues(t, 2, done.success.Load())
 }
@@ -380,14 +350,7 @@ func TestShardBatcher_MergeError(t *testing.T) {
 	}
 
 	sink := requesttest.NewSink()
-	ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
-		sizerType:   request.SizerTypeItems,
-		sizer:       request.NewItemsSizer(),
-		partitioner: nil,
-		next:        sink.Export,
-		maxWorkers:  2,
-	})
-
+	ba := newShard(cfg, request.SizerTypeItems, request.NewItemsSizer(), newWorkerPool(2), sink.Export)
 	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
 	t.Cleanup(func() {
 		require.NoError(t, ba.Shutdown(context.Background()))
@@ -405,7 +368,7 @@ func TestShardBatcher_MergeError(t *testing.T) {
 		return done.errors.Load() == 2
 	}, 1*time.Second, 10*time.Millisecond)
 
-	// Check that done callback is called for the right amount of times.
+	// Check that done callback is called for the right number of times.
 	assert.EqualValues(t, 2, done.errors.Load())
 	assert.EqualValues(t, 0, done.success.Load())
 }