Bytes based batching for traces

Signed-off-by: Israel Blancas <[email protected]>

Author: iblancasa

exporter/exporterhelper/internal/sizer/traces_sizer.go

@@ -0,0 +1,71 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package sizer // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
+
+import (
+	math_bits "math/bits"
+
+	"go.opentelemetry.io/collector/pdata/ptrace"
+)
+
+type TracesSizer interface {
+	TracesSize(ld ptrace.Traces) int
+	ResourceSpansSize(rs ptrace.ResourceSpans) int
+	ScopeSpansSize(ss ptrace.ScopeSpans) int
+	SpanSize(span ptrace.Span) int
+	// DeltaSize() returns the delta size when a span is added.
+	DeltaSize(newItemSize int) int
+}
+
+// TracesBytesSizer returns the byte size of serialized protos.
+type TracesBytesSizer struct {
+	ptrace.ProtoMarshaler
+}
+
+// DeltaSize() returns the delta size of a proto slice when a new item is added.
+// Example:
+//
+//	prevSize := proto1.Size()
+//	proto1.RepeatedField().AppendEmpty() = proto2
+//
+// Then currSize of proto1 can be calculated as
+//
+//	currSize := (prevSize + sizer.DeltaSize(proto2.Size()))
+//
+// This is derived from pdata/internal/data/protogen/trace/v1/trace.pb.go
+// which is generated with gogo/protobuf.
+func (s *TracesBytesSizer) DeltaSize(newItemSize int) int {
+	return 1 + newItemSize + sov(uint64(newItemSize)) //nolint:gosec // disable G115
+}
+
+// TracesCountSizer returns the number of spans in the traces.
+type TracesCountSizer struct{}
+
+func (s *TracesCountSizer) TracesSize(td ptrace.Traces) int {
+	return td.SpanCount()
+}
+
+func (s *TracesCountSizer) ResourceSpansSize(rs ptrace.ResourceSpans) int {
+	count := 0
+	for k := 0; k < rs.ScopeSpans().Len(); k++ {
+		count += rs.ScopeSpans().At(k).Spans().Len()
+	}
+	return count
+}
+
+func (s *TracesCountSizer) ScopeSpansSize(ss ptrace.ScopeSpans) int {
+	return ss.Spans().Len()
+}
+
+func (s *TracesCountSizer) SpanSize(_ ptrace.Span) int {
+	return 1
+}
+
+func (s *TracesCountSizer) DeltaSize(newItemSize int) int {
+	return newItemSize
+}
+
+func sov(x uint64) int {
+	return (math_bits.Len64(x|1) + 6) / 7
+}

exporter/exporterhelper/internal/sizer/traces_sizer_test.go

@@ -0,0 +1,63 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+package sizer
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"go.opentelemetry.io/collector/pdata/testdata"
+)
+
+func TestTracesCountSizer(t *testing.T) {
+	td := testdata.GenerateTraces(5)
+	sizer := TracesCountSizer{}
+	require.Equal(t, 5, sizer.TracesSize(td))
+
+	rs := td.ResourceSpans().At(0)
+	require.Equal(t, 5, sizer.ResourceSpansSize(rs))
+
+	ss := rs.ScopeSpans().At(0)
+	require.Equal(t, 5, sizer.ScopeSpansSize(ss))
+
+	require.Equal(t, 1, sizer.SpanSize(ss.Spans().At(0)))
+	require.Equal(t, 1, sizer.SpanSize(ss.Spans().At(1)))
+	require.Equal(t, 1, sizer.SpanSize(ss.Spans().At(2)))
+	require.Equal(t, 1, sizer.SpanSize(ss.Spans().At(3)))
+	require.Equal(t, 1, sizer.SpanSize(ss.Spans().At(4)))
+
+	prevSize := sizer.ScopeSpansSize(ss)
+	span := ss.Spans().At(2)
+	span.CopyTo(ss.Spans().AppendEmpty())
+	require.Equal(t, sizer.ScopeSpansSize(ss), prevSize+sizer.DeltaSize(sizer.SpanSize(span)))
+}
+
+func TestTracesBytesSizer(t *testing.T) {
+	td := testdata.GenerateTraces(2)
+	sizer := TracesBytesSizer{}
+	require.Equal(t, 338, sizer.TracesSize(td))
+
+	rs := td.ResourceSpans().At(0)
+	require.Equal(t, 335, sizer.ResourceSpansSize(rs))
+
+	ss := rs.ScopeSpans().At(0)
+	require.Equal(t, 290, sizer.ScopeSpansSize(ss))
+
+	require.Equal(t, 187, sizer.SpanSize(ss.Spans().At(0)))
+	require.Equal(t, 96, sizer.SpanSize(ss.Spans().At(1)))
+
+	prevSize := sizer.ScopeSpansSize(ss)
+	span := ss.Spans().At(1)
+	spanSize := sizer.SpanSize(span)
+	span.CopyTo(ss.Spans().AppendEmpty())
+	ds := sizer.DeltaSize(spanSize)
+	require.Equal(t, prevSize+ds, sizer.ScopeSpansSize(ss))
+}
+
+func TestLogsBytesDeltaSize(t *testing.T) {
+	sizer := TracesBytesSizer{}
+	require.Equal(t, 129, sizer.DeltaSize(127))
+	require.Equal(t, 131, sizer.DeltaSize(128))
+	require.Equal(t, 242, sizer.DeltaSize(239))
+}

exporter/exporterhelper/traces.go

@@ -14,6 +14,7 @@ import (
 	"go.opentelemetry.io/collector/consumer/consumererror"
 	"go.opentelemetry.io/collector/exporter"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 	"go.opentelemetry.io/collector/pipeline"
 )
@@ -24,16 +25,16 @@ var (
 )
 
 type tracesRequest struct {
-	td               ptrace.Traces
-	pusher           consumer.ConsumeTracesFunc
-	cachedItemsCount int
+	td         ptrace.Traces
+	pusher     consumer.ConsumeTracesFunc
+	cachedSize int
 }
 
 func newTracesRequest(td ptrace.Traces, pusher consumer.ConsumeTracesFunc) Request {
 	return &tracesRequest{
-		td:               td,
-		pusher:           pusher,
-		cachedItemsCount: td.SpanCount(),
+		td:         td,
+		pusher:     pusher,
+		cachedSize: -1,
 	}
 }
 
@@ -66,11 +67,18 @@ func (req *tracesRequest) Export(ctx context.Context) error {
 }
 
 func (req *tracesRequest) ItemsCount() int {
-	return req.cachedItemsCount
+	return req.td.SpanCount()
 }
 
-func (req *tracesRequest) setCachedItemsCount(count int) {
-	req.cachedItemsCount = count
+func (req *tracesRequest) Size(sizer sizer.TracesSizer) int {
+	if req.cachedSize == -1 {
+		req.cachedSize = sizer.TracesSize(req.td)
+	}
+	return req.cachedSize
+}
+
+func (req *tracesRequest) setCachedSize(size int) {
+	req.cachedSize = size
 }
 
 type tracesExporter struct {

exporter/exporterhelper/traces_batch.go

@@ -8,106 +8,159 @@ import (
 	"errors"
 
 	"go.opentelemetry.io/collector/exporter/exporterbatcher"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 )
 
 // MergeSplit splits and/or merges the provided traces request and the current request into one or more requests
 // conforming with the MaxSizeConfig.
 func (req *tracesRequest) MergeSplit(_ context.Context, cfg exporterbatcher.SizeConfig, r2 Request) ([]Request, error) {
+	var sz sizer.TracesSizer
+	switch cfg.Sizer {
+	case exporterbatcher.SizerTypeItems:
+		sz = &sizer.TracesCountSizer{}
+	case exporterbatcher.SizerTypeBytes:
+		sz = &sizer.TracesBytesSizer{}
+	default:
+		return nil, errors.New("unknown sizer type")
+	}
+
 	if r2 != nil {
 		req2, ok := r2.(*tracesRequest)
 		if !ok {
 			return nil, errors.New("invalid input type")
 		}
-		req2.mergeTo(req)
+		req2.mergeTo(req, sz)
 	}
 
 	// If no limit we can simply merge the new request into the current and return.
 	if cfg.MaxSize == 0 {
 		return []Request{req}, nil
 	}
-	return req.split(cfg)
+	return req.split(cfg.MaxSize, sz), nil
 }
 
-func (req *tracesRequest) mergeTo(dst *tracesRequest) {
-	dst.setCachedItemsCount(dst.ItemsCount() + req.ItemsCount())
-	req.setCachedItemsCount(0)
+func (req *tracesRequest) mergeTo(dst *tracesRequest, sz sizer.TracesSizer) {
+	if sz != nil {
+		dst.setCachedSize(dst.Size(sz) + req.Size(sz))
+		req.setCachedSize(0)
+	}
 	req.td.ResourceSpans().MoveAndAppendTo(dst.td.ResourceSpans())
 }
 
-func (req *tracesRequest) split(cfg exporterbatcher.SizeConfig) ([]Request, error) {
+func (req *tracesRequest) split(maxSize int, sz sizer.TracesSizer) []Request {
 	var res []Request
-	for req.ItemsCount() > cfg.MaxSize {
-		td := extractTraces(req.td, cfg.MaxSize)
-		size := td.SpanCount()
-		req.setCachedItemsCount(req.ItemsCount() - size)
-		res = append(res, &tracesRequest{td: td, pusher: req.pusher, cachedItemsCount: size})
+	for req.Size(sz) > maxSize {
+		td, rmSize := extractTraces(req.td, maxSize, sz)
+		req.setCachedSize(req.Size(sz) - rmSize)
+		res = append(res, newTracesRequest(td, req.pusher))
 	}
 	res = append(res, req)
-	return res, nil
+	return res
 }
 
 // extractTraces extracts a new traces with a maximum number of spans.
-func extractTraces(srcTraces ptrace.Traces, count int) ptrace.Traces {
+func extractTraces(srcTraces ptrace.Traces, capacity int, sz sizer.TracesSizer) (ptrace.Traces, int) {
 	destTraces := ptrace.NewTraces()
+	capacityLeft := capacity - sz.TracesSize(destTraces)
+	removedSize := 0
 	srcTraces.ResourceSpans().RemoveIf(func(srcRS ptrace.ResourceSpans) bool {
-		if count == 0 {
+		// If the no more capacity left just return.
+		if capacityLeft == 0 {
 			return false
 		}
-		needToExtract := resourceTracesCount(srcRS) > count
-		if needToExtract {
-			srcRS = extractResourceSpans(srcRS, count)
+		rawRsSize := sz.ResourceSpansSize(srcRS)
+		rsSize := sz.DeltaSize(rawRsSize)
+
+		if rsSize > capacityLeft {
+			extSrcRS, extRsSize := extractResourceSpans(srcRS, capacityLeft, sz)
+			// This cannot make it to exactly 0 for the bytes,
+			// force it to be 0 since that is the stopping condition.
+			capacityLeft = 0
+			removedSize += extRsSize
+			// There represents the delta between the delta sizes.
+			removedSize += rsSize - rawRsSize - (sz.DeltaSize(rawRsSize-extRsSize) - (rawRsSize - extRsSize))
+			// It is possible that for the bytes scenario, the extracted field contains no log records.
+			// Do not add it to the destination if that is the case.
+			if extSrcRS.ScopeSpans().Len() > 0 {
+				extSrcRS.MoveTo(destTraces.ResourceSpans().AppendEmpty())
+			}
+			return extSrcRS.ScopeSpans().Len() != 0
 		}
-		count -= resourceTracesCount(srcRS)
+		capacityLeft -= rsSize
+		removedSize += rsSize
+
 		srcRS.MoveTo(destTraces.ResourceSpans().AppendEmpty())
-		return !needToExtract
+		return true
 	})
-	return destTraces
+	return destTraces, removedSize
 }
 
 // extractResourceSpans extracts spans and returns a new resource spans with the specified number of spans.
-func extractResourceSpans(srcRS ptrace.ResourceSpans, count int) ptrace.ResourceSpans {
+func extractResourceSpans(srcRS ptrace.ResourceSpans, capacity int, sz sizer.TracesSizer) (ptrace.ResourceSpans, int) {
 	destRS := ptrace.NewResourceSpans()
 	destRS.SetSchemaUrl(srcRS.SchemaUrl())
 	srcRS.Resource().CopyTo(destRS.Resource())
+	// Take into account that this can have max "capacity", so when added to the parent will need space for the extra delta size.
+	capacityLeft := capacity - (sz.DeltaSize(capacity) - capacity) - sz.ResourceSpansSize(destRS)
+	removedSize := 0
 	srcRS.ScopeSpans().RemoveIf(func(srcSS ptrace.ScopeSpans) bool {
-		if count == 0 {
+		// If the no more capacity left just return.
+		if capacityLeft == 0 {
 			return false
 		}
-		needToExtract := srcSS.Spans().Len() > count
-		if needToExtract {
-			srcSS = extractScopeSpans(srcSS, count)
+
+		rawSlSize := sz.ScopeSpansSize(srcSS)
+		ssSize := sz.DeltaSize(rawSlSize)
+		if ssSize > capacityLeft {
+			extSrcSS, extSsSize := extractScopeSpans(srcSS, capacityLeft, sz)
+			// This cannot make it to exactly 0 for the bytes,
+			// force it to be 0 since that is the stopping condition.
+			capacityLeft = 0
+			removedSize += extSsSize
+			// There represents the delta between the delta sizes.
+			removedSize += ssSize - rawSlSize - (sz.DeltaSize(rawSlSize-extSsSize) - (rawSlSize - extSsSize))
+			// It is possible that for the bytes scenario, the extracted field contains no log records.
+			// Do not add it to the destination if that is the case.
+			if extSrcSS.Spans().Len() > 0 {
+				extSrcSS.MoveTo(destRS.ScopeSpans().AppendEmpty())
+			}
+			return extSrcSS.Spans().Len() != 0
 		}
-		count -= srcSS.Spans().Len()
+		capacityLeft -= ssSize
+		removedSize += ssSize
+
 		srcSS.MoveTo(destRS.ScopeSpans().AppendEmpty())
-		return !needToExtract
+		return true
 	})
-	srcRS.Resource().CopyTo(destRS.Resource())
-	return destRS
+	return destRS, removedSize
 }
 
 // extractScopeSpans extracts spans and returns a new scope spans with the specified number of spans.
-func extractScopeSpans(srcSS ptrace.ScopeSpans, count int) ptrace.ScopeSpans {
+func extractScopeSpans(srcSS ptrace.ScopeSpans, capacity int, sz sizer.TracesSizer) (ptrace.ScopeSpans, int) {
 	destSS := ptrace.NewScopeSpans()
 	destSS.SetSchemaUrl(srcSS.SchemaUrl())
 	srcSS.Scope().CopyTo(destSS.Scope())
+	// Take into account that this can have max "capacity", so when added to the parent will need space for the extra delta size.
+	capacityLeft := capacity - (sz.DeltaSize(capacity) - capacity) - sz.ScopeSpansSize(destSS)
+	removedSize := 0
 	srcSS.Spans().RemoveIf(func(srcSpan ptrace.Span) bool {
-		if count == 0 {
+		// If the no more capacity left just return.
+		if capacityLeft == 0 {
+			return false
+		}
+		rsSize := sz.DeltaSize(sz.SpanSize(srcSpan))
+		if rsSize > capacityLeft {
+			// This cannot make it to exactly 0 for the bytes,
+			// force it to be 0 since that is the stopping condition.
+			capacityLeft = 0
 			return false
 		}
+
+		capacityLeft -= rsSize
+		removedSize += rsSize
 		srcSpan.MoveTo(destSS.Spans().AppendEmpty())
-		count--
 		return true
 	})
-	return destSS
-}
-
-// resourceTracesCount calculates the total number of spans in the pdata.ResourceSpans.
-func resourceTracesCount(rs ptrace.ResourceSpans) int {
-	count := 0
-	rs.ScopeSpans().RemoveIf(func(ss ptrace.ScopeSpans) bool {
-		count += ss.Spans().Len()
-		return false
-	})
-	return count
+	return destSS, removedSize
 }