[exporter][batcher] Multi-batch support - Version 2

exporter/exporterhelper/internal/base_exporter.go

@@ -90,11 +90,12 @@ func NewBaseExporter(set exporter.Settings, signal pipeline.Signal, pusher sende
 
 	if be.queueCfg.Enabled || be.batcherCfg.Enabled {
 		qSet := queuebatch.Settings[request.Request]{
-			Signal:    signal,
-			ID:        set.ID,
-			Telemetry: set.TelemetrySettings,
-			Encoding:  be.queueBatchSettings.Encoding,
-			Sizers:    be.queueBatchSettings.Sizers,
+			Signal:      signal,
+			ID:          set.ID,
+			Telemetry:   set.TelemetrySettings,
+			Encoding:    be.queueBatchSettings.Encoding,
+			Sizers:      be.queueBatchSettings.Sizers,
+			Partitioner: be.queueBatchSettings.Partitioner,
 		}
 		be.QueueSender, err = NewQueueSender(qSet, be.queueCfg, be.batcherCfg, be.ExportFailureMessage, be.firstSender)
 		if err != nil {

exporter/exporterhelper/internal/queuebatch/default_batcher.go

@@ -22,16 +22,17 @@ type batch struct {
 }
 
 type batcherSettings[T any] struct {
-	sizerType  request.SizerType
-	sizer      request.Sizer[T]
-	next       sender.SendFunc[T]
-	maxWorkers int
+	sizerType   request.SizerType
+	sizer       request.Sizer[T]
+	partitioner Partitioner[T]
+	next        sender.SendFunc[T]
+	maxWorkers  int
 }
 
 // defaultBatcher continuously batch incoming requests and flushes asynchronously if minimum size limit is met or on timeout.
 type defaultBatcher struct {
 	cfg            BatchConfig
-	workerPool     chan struct{}
+	workerPool     *chan struct{}
 	sizerType      request.SizerType
 	sizer          request.Sizer[request.Request]
 	consumeFunc    sender.SendFunc[request.Request]
@@ -53,7 +54,7 @@ func newDefaultBatcher(bCfg BatchConfig, bSet batcherSettings[request.Request])
 	}
 	return &defaultBatcher{
 		cfg:         bCfg,
-		workerPool:  workerPool,
+		workerPool:  &workerPool,
 		sizerType:   bSet.sizerType,
 		sizer:       bSet.sizer,
 		consumeFunc: bSet.next,
@@ -210,13 +211,13 @@ func (qb *defaultBatcher) flushCurrentBatchIfNecessary() {
 func (qb *defaultBatcher) flush(ctx context.Context, req request.Request, done Done) {
 	qb.stopWG.Add(1)
 	if qb.workerPool != nil {
-		<-qb.workerPool
+		<-*qb.workerPool
 	}
 	go func() {
 		defer qb.stopWG.Done()
 		done.OnDone(qb.consumeFunc(ctx, req))
 		if qb.workerPool != nil {
-			qb.workerPool <- struct{}{}
+			*qb.workerPool <- struct{}{}
 		}
 	}()
 }

exporter/exporterhelper/internal/queuebatch/multi_batcher.go

@@ -0,0 +1,92 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package queuebatch // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch"
+import (
+	"context"
+	"sync"
+
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/request"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sender"
+)
+
+type multiBatcher struct {
+	cfg         BatchConfig
+	workerPool  *chan struct{}
+	sizerType   request.SizerType
+	sizer       request.Sizer[request.Request]
+	partitioner Partitioner[request.Request]
+	consumeFunc sender.SendFunc[request.Request]
+
+	shardMapMu sync.Mutex
+	shards     map[string]*defaultBatcher
+}
+
+func newMultiBatcher(bCfg BatchConfig, bSet batcherSettings[request.Request]) *multiBatcher {
+	// TODO: Determine what is the right behavior for this in combination with async queue.
+	var workerPool chan struct{}
+	if bSet.maxWorkers != 0 {
+		workerPool = make(chan struct{}, bSet.maxWorkers)
+		for i := 0; i < bSet.maxWorkers; i++ {
+			workerPool <- struct{}{}
+		}
+	}
+	return &multiBatcher{
+		cfg:         bCfg,
+		workerPool:  &workerPool,
+		sizerType:   bSet.sizerType,
+		sizer:       bSet.sizer,
+		partitioner: bSet.partitioner,
+		consumeFunc: bSet.next,
+		shardMapMu:  sync.Mutex{},
+		shards:      make(map[string]*defaultBatcher),
+	}
+}
+
+func (qb *multiBatcher) getShard(ctx context.Context, req request.Request) *defaultBatcher {
+	key := qb.partitioner.GetKey(ctx, req)
+
+	qb.shardMapMu.Lock()
+	defer qb.shardMapMu.Unlock()
+
+	s, ok := qb.shards[key]
+	if !ok {
+		s = &defaultBatcher{
+			cfg:         qb.cfg,
+			workerPool:  qb.workerPool,
+			sizerType:   qb.sizerType,
+			sizer:       qb.sizer,
+			consumeFunc: qb.consumeFunc,
+			stopWG:      sync.WaitGroup{},
+			shutdownCh:  make(chan struct{}, 1),
+		}
+		qb.shards[key] = s
+		_ = s.Start(ctx, nil)
+	}
+	return s
+}
+
+func (qb *multiBatcher) Start(_ context.Context, _ component.Host) error {
+	return nil
+}
+
+func (qb *multiBatcher) Consume(ctx context.Context, req request.Request, done Done) {
+	shard := qb.getShard(ctx, req)
+	shard.Consume(ctx, req, done)
+}
+
+func (qb *multiBatcher) Shutdown(ctx context.Context) error {
+	qb.shardMapMu.Lock()
+	defer qb.shardMapMu.Unlock()
+	stopWG := sync.WaitGroup{}
+	for _, shard := range qb.shards {
+		stopWG.Add(1)
+		go func() {
+			_ = shard.Shutdown(ctx)
+			stopWG.Done()
+		}()
+	}
+	stopWG.Wait()
+	return nil
+}

exporter/exporterhelper/internal/queuebatch/multi_batcher_test.go

@@ -0,0 +1,113 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package queuebatch
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/request"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/requesttest"
+)
+
+func TestMultiBatcher_NoTimeout(t *testing.T) {
+	cfg := BatchConfig{
+		FlushTimeout: 0,
+		MinSize:      10,
+	}
+	sink := requesttest.NewSink()
+
+	type partitionKey struct{}
+
+	ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
+		sizerType: request.SizerTypeItems,
+		sizer:     request.NewItemsSizer(),
+		partitioner: NewPartitioner(func(ctx context.Context, _ request.Request) string {
+			return ctx.Value(partitionKey{}).(string)
+		}),
+		next:       sink.Export,
+		maxWorkers: 1,
+	})
+
+	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
+	t.Cleanup(func() {
+		require.NoError(t, ba.Shutdown(context.Background()))
+	})
+
+	done := newFakeDone()
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}, done)
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}, done)
+
+	// Neither batch should be flushed since they haven't reached min threshold.
+	assert.Equal(t, 0, sink.RequestsCount())
+	assert.Equal(t, 0, sink.ItemsCount())
+
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}, done)
+
+	assert.Eventually(t, func() bool {
+		return sink.RequestsCount() == 1 && sink.ItemsCount() == 16
+	}, 500*time.Millisecond, 10*time.Millisecond)
+
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}, done)
+
+	assert.Eventually(t, func() bool {
+		return sink.RequestsCount() == 2 && sink.ItemsCount() == 28
+	}, 500*time.Millisecond, 10*time.Millisecond)
+
+	// Check that done callback is called for the right amount of times.
+	assert.EqualValues(t, 0, done.errors.Load())
+	assert.EqualValues(t, 4, done.success.Load())
+
+	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
+}
+
+func TestMultiBatcher_Timeout(t *testing.T) {
+	cfg := BatchConfig{
+		FlushTimeout: 100 * time.Millisecond,
+		MinSize:      100,
+	}
+	sink := requesttest.NewSink()
+
+	type partitionKey struct{}
+
+	ba := newMultiBatcher(cfg, batcherSettings[request.Request]{
+		sizerType: request.SizerTypeItems,
+		sizer:     request.NewItemsSizer(),
+		partitioner: NewPartitioner(func(ctx context.Context, _ request.Request) string {
+			return ctx.Value(partitionKey{}).(string)
+		}),
+		next:       sink.Export,
+		maxWorkers: 1,
+	})
+
+	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
+	t.Cleanup(func() {
+		require.NoError(t, ba.Shutdown(context.Background()))
+	})
+
+	done := newFakeDone()
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}, done)
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}, done)
+
+	// Neither batch should be flushed since they haven't reached min threshold.
+	assert.Equal(t, 0, sink.RequestsCount())
+	assert.Equal(t, 0, sink.ItemsCount())
+
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}, done)
+	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}, done)
+
+	assert.Eventually(t, func() bool {
+		return sink.RequestsCount() == 2 && sink.ItemsCount() == 28
+	}, 1*time.Second, 10*time.Millisecond)
+	// Check that done callback is called for the right amount of times.
+	assert.EqualValues(t, 0, done.errors.Load())
+	assert.EqualValues(t, 4, done.success.Load())
+
+	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
+}

exporter/exporterhelper/internal/queuebatch/queue_batch.go

@@ -16,11 +16,12 @@ import (
 
 // Settings defines settings for creating a QueueBatch.
 type Settings[T any] struct {
-	Signal    pipeline.Signal
-	ID        component.ID
-	Telemetry component.TelemetrySettings
-	Encoding  Encoding[T]
-	Sizers    map[request.SizerType]request.Sizer[T]
+	Signal      pipeline.Signal
+	ID          component.ID
+	Telemetry   component.TelemetrySettings
+	Encoding    Encoding[T]
+	Sizers      map[request.SizerType]request.Sizer[T]
+	Partitioner Partitioner[T]
 }
 
 type QueueBatch struct {
@@ -74,12 +75,23 @@ func newQueueBatch(
 				maxWorkers: cfg.NumConsumers,
 			})
 		} else {
-			b = newDefaultBatcher(*cfg.Batch, batcherSettings[request.Request]{
-				sizerType:  cfg.Sizer,
-				sizer:      sizer,
-				next:       next,
-				maxWorkers: cfg.NumConsumers,
-			})
+			// If partitioning is not enabled or if paritition is done at queue level, we can use the default batcher.
+			if set.Partitioner == nil {
+				b = newDefaultBatcher(*cfg.Batch, batcherSettings[request.Request]{
+					sizerType:  cfg.Sizer,
+					sizer:      sizer,
+					next:       next,
+					maxWorkers: cfg.NumConsumers,
+				})
+			} else {
+				b = newMultiBatcher(*cfg.Batch, batcherSettings[request.Request]{
+					sizerType:   cfg.Sizer,
+					sizer:       sizer,
+					partitioner: set.Partitioner,
+					next:        next,
+					maxWorkers:  cfg.NumConsumers,
+				})
+			}
 		}
 	} else {
 		b = newDisabledBatcher[request.Request](next)

exporter/exporterhelper/internal/queuebatch/queue_batch_test.go

@@ -404,6 +404,48 @@ func TestQueueBatch_MergeOrSplit(t *testing.T) {
 	require.NoError(t, qb.Shutdown(context.Background()))
 }
 
+func TestQueueBatch_MergeOrSplit_Multibatch(t *testing.T) {
+	sink := requesttest.NewSink()
+	cfg := newTestConfig()
+	cfg.Batch = &BatchConfig{
+		FlushTimeout: 100 * time.Millisecond,
+		MinSize:      10,
+	}
+
+	type partitionKey struct{}
+	set := newFakeRequestSettings()
+	set.Partitioner = NewPartitioner(func(ctx context.Context, _ request.Request) string {
+		key := ctx.Value(partitionKey{}).(string)
+		return key
+	})
+
+	qb, err := NewQueueBatch(set, cfg, sink.Export)
+	require.NoError(t, err)
+	require.NoError(t, qb.Start(context.Background(), componenttest.NewNopHost()))
+
+	// should be sent right away by reaching the minimum items size.
+	require.NoError(t, qb.Send(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}))
+	require.NoError(t, qb.Send(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}))
+
+	// Neither batch should be flushed since they haven't reached min threshold.
+	assert.Equal(t, 0, sink.RequestsCount())
+	assert.Equal(t, 0, sink.ItemsCount())
+
+	require.NoError(t, qb.Send(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}))
+
+	assert.Eventually(t, func() bool {
+		return sink.RequestsCount() == 1 && sink.ItemsCount() == 16
+	}, 500*time.Millisecond, 10*time.Millisecond)
+
+	require.NoError(t, qb.Send(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}))
+
+	assert.Eventually(t, func() bool {
+		return sink.RequestsCount() == 2 && sink.ItemsCount() == 28
+	}, 500*time.Millisecond, 10*time.Millisecond)
+
+	require.NoError(t, qb.Shutdown(context.Background()))
+}
+
 func TestQueueBatch_Shutdown(t *testing.T) {
 	sink := requesttest.NewSink()
 	qb, err := NewQueueBatch(newFakeRequestSettings(), newTestConfig(), sink.Export)