telemetry/exporter/exporterhelper/internal/queuebatch/multi_batcher_test.go

// 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.uber.org/zap"

	"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,
		Sizer:        request.SizerTypeItems,
		MinSize:      10,
	}
	sink := requesttest.NewSink()

	type partitionKey struct{}

	ba, err := newMultiBatcher(cfg,
		request.NewItemsSizer(),
		newWorkerPool(1),
		NewPartitioner(func(ctx context.Context, _ request.Request) string {
			return ctx.Value(partitionKey{}).(string)
		}),
		nil,
		sink.Export,
		zap.NewNop(),
	)

	require.NoError(t, err)
	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
	t.Cleanup(func() {
		require.NoError(t, ba.Shutdown(context.Background()))
	})

	done := newFakeDone()
	assert.Equal(t, int64(0), ba.getActivePartitionsCount())
	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 8}, done)
	assert.Equal(t, int64(1), ba.getActivePartitionsCount())
	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p2"), &requesttest.FakeRequest{Items: 6}, done)
	assert.Equal(t, int64(2), ba.getActivePartitionsCount())

	// 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,
		Sizer:        request.SizerTypeItems,
		MinSize:      100,
	}
	sink := requesttest.NewSink()

	type partitionKey struct{}

	ba, err := newMultiBatcher(cfg,
		request.NewItemsSizer(),
		newWorkerPool(1),
		NewPartitioner(func(ctx context.Context, _ request.Request) string {
			return ctx.Value(partitionKey{}).(string)
		}),
		nil,
		sink.Export,
		zap.NewNop(),
	)

	require.NoError(t, err)
	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()))
}

func TestMultiBatcher_PartitionRemovedAfterIdleTimeout(t *testing.T) {
	// Use a short FlushTimeout so the idle threshold (partitionIdleCycles*FlushTimeout) is reached quickly.
	cfg := BatchConfig{
		FlushTimeout: 10 * time.Millisecond,
		Sizer:        request.SizerTypeItems,
		MinSize:      100, // High min size to prevent immediate flush
	}
	sink := requesttest.NewSink()

	type partitionKey struct{}

	ba, err := newMultiBatcher(cfg,
		request.NewItemsSizer(),
		newWorkerPool(1),
		NewPartitioner(func(ctx context.Context, _ request.Request) string {
			return ctx.Value(partitionKey{}).(string)
		}),
		nil,
		sink.Export,
		zap.NewNop(),
	)

	require.NoError(t, err)
	require.NoError(t, ba.Start(context.Background(), componenttest.NewNopHost()))
	t.Cleanup(func() {
		require.NoError(t, ba.Shutdown(context.Background()))
	})

	done := newFakeDone()

	// Create a partition
	ba.Consume(context.WithValue(context.Background(), partitionKey{}, "p1"), &requesttest.FakeRequest{Items: 5}, done)
	assert.Equal(t, int64(1), ba.getActivePartitionsCount())

	// Wait for the batch to flush via timeout
	assert.Eventually(t, func() bool {
		return sink.RequestsCount() == 1
	}, 500*time.Millisecond, 10*time.Millisecond)

	// Wait for idle timeout (partitionIdleCycles * FlushTimeout = 10 * 10ms = 100ms)
	// After this, the partition should be removed from the LRU cache.
	assert.Eventually(t, func() bool {
		return ba.getActivePartitionsCount() == 0
	}, 500*time.Millisecond, 10*time.Millisecond)
}