Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bogdandrutu merged 1 commit into from
Jun 6, 2025
+66 −96
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 7 additions & 15 deletions exporter/exporterhelper/internal/queuebatch/multi_batcher.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

type multiBatcher struct {
cfg BatchConfig
wp chan struct{}
wp *workerPool
sizerType request.SizerType
sizer request.Sizer[request.Request]
partitioner Partitioner[request.Request]
Expand All @@ -34,16 +34,9 @@
var _ Batcher[request.Request] = (*multiBatcher)(nil)

func newMultiBatcher(bCfg BatchConfig, bSet batcherSettings[request.Request]) *multiBatcher {
var workerPool chan struct{}
if bSet.maxWorkers != 0 {
Copy link
Copy Markdown
Member Author

@bogdandrutu bogdandrutu Jun 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maxWorkers cannot be 0, because https://github.com/open-telemetry/opentelemetry-collector/blob/main/exporter/exporterhelper/internal/queuebatch/config.go#L77

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,
Expand All @@ -68,18 +61,18 @@
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)

Check warning on line 68 in exporter/exporterhelper/internal/queuebatch/multi_batcher.go

View check run for this annotation

Codecov / codecov/patch

exporter/exporterhelper/internal/queuebatch/multi_batcher.go#L68 

Added line #L68 was not covered by tests
}
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
}
Expand All @@ -91,16 +84,15 @@

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
})
Expand Down
79 changes: 47 additions & 32 deletions exporter/exporterhelper/internal/queuebatch/shard_batcher.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
61 changes: 12 additions & 49 deletions exporter/exporterhelper/internal/queuebatch/shard_batcher_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()))
Expand All @@ -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())
})
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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())
})
Expand Down Expand Up @@ -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()))
Expand All @@ -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())
})
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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())
})
Expand All @@ -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()
Expand All @@ -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())
}
Expand All @@ -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()))
Expand All @@ -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())
}
Expand Down
Loading