Multi batch support

Author: sfc-gh-sili

exporter/exporterhelper/internal/queue_sender.go

@@ -22,8 +22,9 @@ import (
 
 // QueueBatchSettings is a subset of the queuebatch.Settings that are needed when used within an Exporter.
 type QueueBatchSettings[K any] struct {
-	Encoding queuebatch.Encoding[K]
-	Sizers   map[request.SizerType]request.Sizer[K]
+	Encoding       queuebatch.Encoding[K]
+	Sizers         map[request.SizerType]request.Sizer[K]
+	BatchKeyGetter queuebatch.BatchKeyGetter[K]
 }
 
 // NewDefaultQueueConfig returns the default config for QueueConfig.

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(keyGetter BatchKeyGetter[request.Request]) batchManager {
+	if keyGetter == nil {
+		return &singleBatchManager{
+			batch: newBatchEntry(),
+		}
+	}
+	return &multiBatchManager{
+		batchMap:  make(map[string]*batchEntry),
+		keyGetter: keyGetter,
+	}
+}
+
+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
+	keyGetter BatchKeyGetter[request.Request]
+	mu        sync.RWMutex
+}
+
+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.keyGetter.GetKey(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

@@ -15,31 +15,26 @@ import (
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sender"
 )
 
-type batch struct {
-	ctx  context.Context
-	req  request.Request
-	done multiDone
-}
-
 type batcherSettings[K any] struct {
 	sizerType  request.SizerType
 	sizer      request.Sizer[K]
+	keyGetter  BatchKeyGetter[K]
 	next       sender.SendFunc[K]
 	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{}
-	sizerType      request.SizerType
-	sizer          request.Sizer[request.Request]
-	consumeFunc    sender.SendFunc[request.Request]
-	stopWG         sync.WaitGroup
-	currentBatchMu sync.Mutex
-	currentBatch   *batch
-	timer          *time.Timer
-	shutdownCh     chan struct{}
+	cfg         BatchConfig
+	workerPool  chan struct{}
+	sizerType   request.SizerType
+	sizer       request.Sizer[request.Request]
+	consumeFunc sender.SendFunc[request.Request]
+	stopWG      sync.WaitGroup
+	ticker      *time.Ticker
+	shutdownCh  chan struct{}
+
+	batchManager batchManager
 }
 
 func newDefaultBatcher(bCfg BatchConfig, bSet batcherSettings[request.Request]) *defaultBatcher {
@@ -52,30 +47,26 @@ func newDefaultBatcher(bCfg BatchConfig, bSet batcherSettings[request.Request])
 		}
 	}
 	return &defaultBatcher{
-		cfg:         bCfg,
-		workerPool:  workerPool,
-		sizerType:   bSet.sizerType,
-		sizer:       bSet.sizer,
-		consumeFunc: bSet.next,
-		stopWG:      sync.WaitGroup{},
-		shutdownCh:  make(chan struct{}, 1),
-	}
-}
-
-func (qb *defaultBatcher) resetTimer() {
-	if qb.cfg.FlushTimeout > 0 {
-		qb.timer.Reset(qb.cfg.FlushTimeout)
+		cfg:          bCfg,
+		workerPool:   workerPool,
+		sizerType:    bSet.sizerType,
+		sizer:        bSet.sizer,
+		consumeFunc:  bSet.next,
+		stopWG:       sync.WaitGroup{},
+		shutdownCh:   make(chan struct{}, 1),
+		batchManager: newBatchManager(bSet.keyGetter),
 	}
 }
 
 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, int(qb.cfg.MaxSize), qb.sizerType, nil)
 		if mergeSplitErr != nil || len(reqList) == 0 {
 			done.OnDone(mergeSplitErr)
-			qb.currentBatchMu.Unlock()
+			batchEntry.mu.Unlock()
 			return
 		}
 
@@ -90,27 +81,27 @@ func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done
 		if qb.sizer.Sizeof(lastReq) < qb.cfg.MinSize {
 			// Do not flush the last item and add it to the current batch.
 			reqList = reqList[:len(reqList)-1]
-			qb.currentBatch = &batch{
-				ctx:  ctx,
-				req:  lastReq,
-				done: multiDone{done},
+			batchEntry.batch = &batch{
+				ctx:     ctx,
+				req:     lastReq,
+				done:    multiDone{done},
+				created: time.Now(),
 			}
-			qb.resetTimer()
 		}
 
-		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, int(qb.cfg.MaxSize), qb.sizerType, req)
+	reqList, mergeSplitErr := batchEntry.batch.req.MergeSplit(ctx, int(qb.cfg.MaxSize), qb.sizerType, 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
 	}
 
@@ -126,15 +117,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.sizer.Sizeof(qb.currentBatch.req) >= qb.cfg.MinSize {
-		firstBatch = qb.currentBatch
-		qb.currentBatch = nil
+	if len(reqList) > 1 || qb.sizer.Sizeof(batchEntry.batch.req) >= qb.cfg.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:]
@@ -145,16 +136,16 @@ func (qb *defaultBatcher) Consume(ctx context.Context, req request.Request, done
 		if qb.sizer.Sizeof(lastReq) < qb.cfg.MinSize {
 			// Do not flush the last item and add it to the current batch.
 			reqList = reqList[:len(reqList)-1]
-			qb.currentBatch = &batch{
-				ctx:  ctx,
-				req:  lastReq,
-				done: multiDone{done},
+			batchEntry.batch = &batch{
+				ctx:     ctx,
+				req:     lastReq,
+				done:    multiDone{done},
+				created: time.Now(),
 			}
-			qb.resetTimer()
 		}
 	}
 
-	qb.currentBatchMu.Unlock()
+	batchEntry.mu.Unlock()
 	if firstBatch != nil {
 		qb.flush(firstBatch.ctx, firstBatch.req, firstBatch.done)
 	}
@@ -172,8 +163,8 @@ func (qb *defaultBatcher) startTimeBasedFlushingGoroutine() {
 			select {
 			case <-qb.shutdownCh:
 				return
-			case <-qb.timer.C:
-				qb.flushCurrentBatchIfNecessary()
+			case <-qb.ticker.C:
+				qb.flushCurrentBatchIfNecessary(false)
 			}
 		}
 	}()
@@ -182,27 +173,32 @@ func (qb *defaultBatcher) startTimeBasedFlushingGoroutine() {
 // Start starts the goroutine that reads from the queue and flushes asynchronously.
 func (qb *defaultBatcher) Start(_ context.Context, _ component.Host) error {
 	if qb.cfg.FlushTimeout > 0 {
-		qb.timer = time.NewTimer(qb.cfg.FlushTimeout)
+		qb.ticker = time.NewTicker(qb.cfg.FlushTimeout)
 		qb.startTimeBasedFlushingGoroutine()
 	}
 
 	return nil
 }
 
 // flushCurrentBatchIfNecessary sends out the current request batch if it is not nil
-func (qb *defaultBatcher) flushCurrentBatchIfNecessary() {
-	qb.currentBatchMu.Lock()
-	if qb.currentBatch == nil {
-		qb.currentBatchMu.Unlock()
-		return
-	}
-	batchToFlush := qb.currentBatch
-	qb.currentBatch = nil
-	qb.currentBatchMu.Unlock()
+func (qb *defaultBatcher) flushCurrentBatchIfNecessary(forceFlush bool) {
+	qb.batchManager.forEachBatch(func(batchEntry *batchEntry) {
+		batchEntry.mu.Lock()
+		if batchEntry.batch == nil {
+			batchEntry.mu.Unlock()
+			return
+		}
+		if !forceFlush && time.Since(batchEntry.created) < qb.cfg.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)
-	qb.resetTimer()
+		// 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.
@@ -224,7 +220,7 @@ func (qb *defaultBatcher) flush(ctx context.Context, req request.Request, done D
 func (qb *defaultBatcher) Shutdown(_ context.Context) error {
 	close(qb.shutdownCh)
 	// Make sure execute one last flush if necessary.
-	qb.flushCurrentBatchIfNecessary()
+	qb.flushCurrentBatchIfNecessary(true)
 	qb.stopWG.Wait()
 	return nil
 }

exporter/exporterhelper/internal/queuebatch/key_func.go

@@ -0,0 +1,17 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package queuebatch // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/queuebatch"
+
+import "context"
+
+// BatchKeyGetter is an interface that returns the the batch key of the given element.
+type BatchKeyGetter[T any] interface {
+	GetKey(T) string
+}
+
+type KeyFunc[T any] func(T) string
+
+func (f KeyFunc[T]) GetKey(ctx context.Context, t T) string {
+	return f(t)
+}

exporter/exporterhelper/internal/queuebatch/queue_batch_test.go

@@ -590,8 +590,8 @@ func TestQueueBatchTimerFlush(t *testing.T) {
 	assert.LessOrEqual(t, 1, sink.RequestsCount())
 	assert.Equal(t, 8, sink.ItemsCount())
 
-	// Confirm that it is flushed after 100ms (using 60+50=110 here to be safe)
-	time.Sleep(50 * time.Millisecond)
+	// Confirm that it is flushed after 100ms (using 100+50=150 here to be safe)
+	time.Sleep(100 * time.Millisecond)
 	assert.LessOrEqual(t, 2, sink.RequestsCount())
 	assert.Equal(t, 12, sink.ItemsCount())
 	require.NoError(t, qb.Shutdown(context.Background()))