Multi batch support

Author: sfc-gh-sili

exporter/exporterhelper/internal/base_exporter.go

@@ -50,6 +50,7 @@ type BaseExporter struct {
 	queueBatchSettings QueueBatchSettings[request.Request]
 	queueCfg           QueueConfig
 	batcherCfg         exporterbatcher.Config
+	keyFunc            queuebatch.KeyFunc
 }
 
 func NewBaseExporter(set exporter.Settings, signal pipeline.Signal, pusher sender.SendFunc[request.Request], options ...Option) (*BaseExporter, error) {
@@ -244,6 +245,15 @@ func WithBatcher(cfg exporterbatcher.Config) Option {
 	}
 }
 
+// WithMultiBatcher is same WithBatcher except that it allows key-based batching.
+func WithMultiBatcher(cfg exporterbatcher.Config, keyFunc queuebatch.KeyFunc) Option {
+	return func(o *BaseExporter) error {
+		o.batcherCfg = cfg
+		o.keyFunc = keyFunc
+		return nil
+	}
+}
+
 // WithQueueBatchSettings is used to set the QueueBatchSettings for the new request based exporter helper.
 // It must be provided as the first option when creating a new exporter helper.
 func WithQueueBatchSettings(set QueueBatchSettings[request.Request]) Option {

exporter/exporterhelper/internal/queuebatch/batch_manager.go

@@ -0,0 +1,90 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package queuebatch // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch"
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/request"
+)
+
+type batch struct {
+	ctx     context.Context
+	req     request.Request
+	done    multiDone
+	created time.Time
+}
+
+type batchEntry struct {
+	*batch
+	mu sync.Mutex
+}
+
+func newBatchEntry() *batchEntry {
+	return &batchEntry{
+		nil,
+		sync.Mutex{},
+	}
+}
+
+type batchManager interface {
+	getBatch(ctx context.Context, req request.Request) *batchEntry
+	forEachBatch(func(*batchEntry))
+}
+
+func newBatchManager(keyFunc KeyFunc) batchManager {
+	if keyFunc == nil {
+		return &singleBatchManager{
+			batch: newBatchEntry(),
+		}
+	}
+	return &multiBatchManager{
+		batchMap: make(map[string]*batchEntry),
+		keyFunc:  keyFunc,
+	}
+}
+
+type singleBatchManager struct {
+	batch *batchEntry
+}
+
+func (bm *singleBatchManager) getBatch(_ context.Context, _ request.Request) *batchEntry {
+	return bm.batch
+}
+
+func (bm *singleBatchManager) forEachBatch(callback func(*batchEntry)) {
+	callback(bm.batch)
+}
+
+type multiBatchManager struct {
+	batchMap map[string]*batchEntry
+	mu       sync.RWMutex
+	keyFunc  KeyFunc
+}
+
+func (bm *multiBatchManager) forEachBatch(callback func(*batchEntry)) {
+	bm.mu.RLock()
+	for _, batchEntry := range bm.batchMap {
+		callback(batchEntry)
+	}
+	bm.mu.RUnlock()
+}
+
+func (bm *multiBatchManager) getBatch(ctx context.Context, req request.Request) *batchEntry {
+	key := bm.keyFunc(ctx, req)
+
+	bm.mu.RLock()
+	batchEntry, ok := bm.batchMap[key]
+	bm.mu.RUnlock()
+	if ok {
+		return batchEntry
+	}
+	bm.mu.Lock()
+	batchEntry = newBatchEntry()
+	bm.batchMap[key] = batchEntry
+	bm.mu.Unlock()
+	return batchEntry
+}

exporter/exporterhelper/internal/queuebatch/default_batcher.go

@@ -16,26 +16,31 @@ import (
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sender"
 )
 
-type batch struct {
-	ctx     context.Context
-	req     request.Request
-	done    multiDone
-	created time.Time
-}
-
 // defaultBatcher continuously batch incoming requests and flushes asynchronously if minimum size limit is met or on timeout.
 type defaultBatcher struct {
-	batchCfg       BatchConfig
-	workerPool     chan struct{}
-	consumeFunc    sender.SendFunc[request.Request]
-	stopWG         sync.WaitGroup
-	currentBatchMu sync.Mutex
-	currentBatch   *batch
-	ticker         *time.Ticker
-	shutdownCh     chan struct{}
+	batchCfg     BatchConfig
+	workerPool   chan struct{}
+	consumeFunc  sender.SendFunc[request.Request]
+	batchManager batchManager
+	stopWG       sync.WaitGroup
+	ticker       *time.Ticker
+	shutdownCh   chan struct{}
+}
+
+func newDefaultBatcher(
+	batchCfg BatchConfig,
+	consumeFunc sender.SendFunc[request.Request],
+	maxWorkers int,
+) *defaultBatcher {
+	return newDefaultBatcherWithKeyFunc(batchCfg, nil, consumeFunc, maxWorkers)
 }
 
-func newDefaultBatcher(batchCfg BatchConfig, consumeFunc sender.SendFunc[request.Request], maxWorkers int) *defaultBatcher {
+func newDefaultBatcherWithKeyFunc(
+	batchCfg BatchConfig,
+	keyFunc KeyFunc,
+	consumeFunc sender.SendFunc[request.Request],
+	maxWorkers int,
+) *defaultBatcher {
 	// TODO: Determine what is the right behavior for this in combination with async queue.
 	var workerPool chan struct{}
 	if maxWorkers != 0 {
@@ -45,22 +50,24 @@ func newDefaultBatcher(batchCfg BatchConfig, consumeFunc sender.SendFunc[request
 		}
 	}
 	return &defaultBatcher{
-		batchCfg:    batchCfg,
-		workerPool:  workerPool,
-		consumeFunc: consumeFunc,
-		stopWG:      sync.WaitGroup{},
-		shutdownCh:  make(chan struct{}, 1),
+		batchCfg:     batchCfg,
+		workerPool:   workerPool,
+		consumeFunc:  consumeFunc,
+		batchManager: newBatchManager(keyFunc),
+		stopWG:       sync.WaitGroup{},
+		shutdownCh:   make(chan struct{}, 1),
 	}
 }
 
 func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done Done) {
-	qb.currentBatchMu.Lock()
+	batchEntry := qb.batchManager.getBatch(ctx, req)
+	batchEntry.mu.Lock()
 
-	if qb.currentBatch == nil {
+	if batchEntry.batch == nil {
 		reqList, mergeSplitErr := req.MergeSplit(ctx, qb.batchCfg.MaxSize, exporterbatcher.SizerTypeItems, nil)
 		if mergeSplitErr != nil || len(reqList) == 0 {
 			done.OnDone(mergeSplitErr)
-			qb.currentBatchMu.Unlock()
+			batchEntry.mu.Unlock()
 			return
 		}
 
@@ -75,27 +82,27 @@ func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done
 		if lastReq.ItemsCount() < qb.batchCfg.MinSize {
 			// Do not flush the last item and add it to the current batch.
 			reqList = reqList[:len(reqList)-1]
-			qb.currentBatch = &batch{
+			batchEntry.batch = &batch{
 				ctx:     ctx,
 				req:     lastReq,
 				done:    multiDone{done},
 				created: time.Now(),
 			}
 		}
 
-		qb.currentBatchMu.Unlock()
+		batchEntry.mu.Unlock()
 		for i := 0; i < len(reqList); i++ {
 			qb.flush(ctx, reqList[i], done)
 		}
 
 		return
 	}
 
-	reqList, mergeSplitErr := qb.currentBatch.req.MergeSplit(ctx, qb.batchCfg.MaxSize, exporterbatcher.SizerTypeItems, req)
+	reqList, mergeSplitErr := batchEntry.req.MergeSplit(ctx, qb.batchCfg.MaxSize, exporterbatcher.SizerTypeItems, req)
 	// If failed to merge signal all Done callbacks from current batch as well as the current request and reset the current batch.
 	if mergeSplitErr != nil || len(reqList) == 0 {
 		done.OnDone(mergeSplitErr)
-		qb.currentBatchMu.Unlock()
+		batchEntry.mu.Unlock()
 		return
 	}
 
@@ -111,15 +118,15 @@ func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done
 
 	// Logic on how to deal with the current batch:
 	// TODO: Deal with merging Context.
-	qb.currentBatch.req = reqList[0]
-	qb.currentBatch.done = append(qb.currentBatch.done, done)
+	batchEntry.req = reqList[0]
+	batchEntry.done = append(batchEntry.done, done)
 	// Save the "currentBatch" if we need to flush it, because we want to execute flush without holding the lock, and
 	// cannot unlock and re-lock because we are not done processing all the responses.
 	var firstBatch *batch
 	// Need to check the currentBatch if more than 1 result returned or if 1 result return but larger than MinSize.
-	if len(reqList) > 1 || qb.currentBatch.req.ItemsCount() >= qb.batchCfg.MinSize {
-		firstBatch = qb.currentBatch
-		qb.currentBatch = nil
+	if len(reqList) > 1 || batchEntry.req.ItemsCount() >= qb.batchCfg.MinSize {
+		firstBatch = batchEntry.batch
+		batchEntry.batch = nil
 	}
 	// At this moment we dealt with the first result which is iter in the currentBatch or in the `firstBatch` we will flush.
 	reqList = reqList[1:]
@@ -130,7 +137,7 @@ func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done
 		if lastReq.ItemsCount() < qb.batchCfg.MinSize {
 			// Do not flush the last item and add it to the current batch.
 			reqList = reqList[:len(reqList)-1]
-			qb.currentBatch = &batch{
+			batchEntry.batch = &batch{
 				ctx:     ctx,
 				req:     lastReq,
 				done:    multiDone{done},
@@ -139,7 +146,7 @@ func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done
 		}
 	}
 
-	qb.currentBatchMu.Unlock()
+	batchEntry.mu.Unlock()
 	if firstBatch != nil {
 		qb.flush(firstBatch.ctx, firstBatch.req, firstBatch.done)
 	}
@@ -176,21 +183,23 @@ func (qb *defaultBatcher) Start(_ context.Context, _ component.Host) error {
 
 // flushCurrentBatchIfNecessary sends out the current request batch if it is not nil
 func (qb *defaultBatcher) flushCurrentBatchIfNecessary(forceFlush bool) {
-	qb.currentBatchMu.Lock()
-	if qb.currentBatch == nil {
-		qb.currentBatchMu.Unlock()
-		return
-	}
-	if !forceFlush && time.Since(qb.currentBatch.created) < qb.batchCfg.FlushTimeout {
-		qb.currentBatchMu.Unlock()
-		return
-	}
-	batchToFlush := qb.currentBatch
-	qb.currentBatch = nil
-	qb.currentBatchMu.Unlock()
+	qb.batchManager.forEachBatch(func(batchEntry *batchEntry) {
+		batchEntry.mu.Lock()
+		if batchEntry.batch == nil {
+			batchEntry.mu.Unlock()
+			return
+		}
+		if !forceFlush && time.Since(batchEntry.created) < qb.batchCfg.FlushTimeout {
+			batchEntry.mu.Unlock()
+			return
+		}
+		batchToFlush := batchEntry.batch
+		batchEntry.batch = nil
+		batchEntry.mu.Unlock()
 
-	// flush() blocks until successfully started a goroutine for flushing.
-	qb.flush(batchToFlush.ctx, batchToFlush.req, batchToFlush.done)
+		// flush() blocks until successfully started a goroutine for flushing.
+		qb.flush(batchToFlush.ctx, batchToFlush.req, batchToFlush.done)
+	})
 }
 
 // flush starts a goroutine that calls consumeFunc. It blocks until a worker is available if necessary.

exporter/exporterhelper/internal/queuebatch/key_func.go

@@ -0,0 +1,12 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package queuebatch // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch"
+
+import (
+	"context"
+
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/request"
+)
+
+type KeyFunc func(ctx context.Context, req request.Request) string

exporter/exporterhelper/queue_batch.go

@@ -33,6 +33,10 @@ func WithBatcher(cfg exporterbatcher.Config) Option {
 	return internal.WithBatcher(cfg)
 }
 
+func WithMultiBatcher(cfg exporterbatcher.Config, keyFunc queuebatch.KeyFunc) Option {
+	return internal.WithMultiBatcher(cfg, keyFunc)
+}
+
 // QueueBatchConfig defines configuration for queueing and batching for the exporter.
 type QueueBatchConfig = internal.QueueConfig