Bytes based batching for profiles

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

Author: iblancasa

exporter/exporterbatcher/config.go

@@ -22,7 +22,7 @@ var SizerTypeRequests = exporterhelper.RequestSizerTypeRequests
 // Deprecated: [v0.123.0] use exporterhelper.RequestSizerTypeItems
 var SizerTypeItems = exporterhelper.RequestSizerTypeItems
 
-// Deprecated: [v0.123.0] use exporterhelper.RequestSizerTypeRequests
+// Deprecated: [v0.123.0] use exporterhelper.RequestSizerTypeBytes
 var SizerTypeBytes = exporterhelper.RequestSizerTypeBytes
 
 // Deprecated: [v0.123.0] use exporterhelper.NewDefaultBatcherConfig

exporter/exporterhelper/internal/sizer/profiles_sizer.go

@@ -0,0 +1,53 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package sizer // import "go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
+
+import (
+	"go.opentelemetry.io/collector/pdata/pprofile"
+)
+
+type ProfilesSizer interface {
+	ProfilesSize(pd pprofile.Profiles) int
+	ResourceProfilesSize(rp pprofile.ResourceProfiles) int
+	ScopeProfilesSize(sp pprofile.ScopeProfiles) int
+	ProfileSize(p pprofile.Profile) int
+	DeltaSize(newItemSize int) int
+}
+
+// TracesBytesSizer returns the byte size of serialized protos.
+type ProfilesBytesSizer struct {
+	pprofile.ProtoMarshaler
+	protoDeltaSizer
+}
+
+var _ ProfilesSizer = (*ProfilesBytesSizer)(nil)
+
+// ProfilesCountSizer returns the number of profiles in the profiles.
+type ProfilesCountSizer struct{}
+
+var _ ProfilesSizer = (*ProfilesCountSizer)(nil)
+
+func (s *ProfilesCountSizer) ProfilesSize(pd pprofile.Profiles) int {
+	return pd.SampleCount()
+}
+
+func (s *ProfilesCountSizer) ResourceProfilesSize(rp pprofile.ResourceProfiles) int {
+	count := 0
+	for k := 0; k < rp.ScopeProfiles().Len(); k++ {
+		count += rp.ScopeProfiles().At(k).Profiles().Len()
+	}
+	return count
+}
+
+func (s *ProfilesCountSizer) ScopeProfilesSize(sp pprofile.ScopeProfiles) int {
+	return sp.Profiles().Len()
+}
+
+func (s *ProfilesCountSizer) ProfileSize(_ pprofile.Profile) int {
+	return 1
+}
+
+func (s *ProfilesCountSizer) DeltaSize(newItemSize int) int {
+	return newItemSize
+}

exporter/exporterhelper/xexporterhelper/profiles.go

@@ -17,6 +17,7 @@ import (
 	"go.opentelemetry.io/collector/exporter/exporterhelper"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/request"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
 	"go.opentelemetry.io/collector/exporter/xexporter"
 	"go.opentelemetry.io/collector/pdata/pprofile"
 	"go.opentelemetry.io/collector/pipeline/xpipeline"
@@ -46,14 +47,14 @@ func NewProfilesQueueBatchSettings() exporterhelper.QueueBatchSettings {
 }
 
 type profilesRequest struct {
-	pd               pprofile.Profiles
-	cachedItemsCount int
+	pd         pprofile.Profiles
+	cachedSize int
 }
 
 func newProfilesRequest(pd pprofile.Profiles) exporterhelper.Request {
 	return &profilesRequest{
-		pd:               pd,
-		cachedItemsCount: pd.SampleCount(),
+		pd:         pd,
+		cachedSize: -1,
 	}
 }
 
@@ -80,11 +81,18 @@ func (req *profilesRequest) OnError(err error) exporterhelper.Request {
 }
 
 func (req *profilesRequest) ItemsCount() int {
-	return req.cachedItemsCount
+	return req.pd.SampleCount()
 }
 
-func (req *profilesRequest) setCachedItemsCount(count int) {
-	req.cachedItemsCount = count
+func (req *profilesRequest) size(sizer sizer.ProfilesSizer) int {
+	if req.cachedSize == -1 {
+		req.cachedSize = sizer.ProfilesSize(req.pd)
+	}
+	return req.cachedSize
+}
+
+func (req *profilesRequest) setCachedSize(size int) {
+	req.cachedSize = size
 }
 
 type profileExporter struct {

exporter/exporterhelper/xexporterhelper/profiles_batch.go

@@ -8,108 +8,157 @@ import (
 	"errors"
 
 	"go.opentelemetry.io/collector/exporter/exporterhelper"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
 	"go.opentelemetry.io/collector/pdata/pprofile"
 )
 
 // MergeSplit splits and/or merges the profiles into multiple requests based on the MaxSizeConfig.
-func (req *profilesRequest) MergeSplit(_ context.Context, maxSize int, _ exporterhelper.RequestSizerType, r2 exporterhelper.Request) ([]exporterhelper.Request, error) {
+func (req *profilesRequest) MergeSplit(_ context.Context, maxSize int, szt exporterhelper.RequestSizerType, r2 exporterhelper.Request) ([]exporterhelper.Request, error) {
+	var sz sizer.ProfilesSizer
+	switch szt {
+	case exporterhelper.RequestSizerTypeItems:
+		sz = &sizer.ProfilesCountSizer{}
+	case exporterhelper.RequestSizerTypeBytes:
+		sz = &sizer.ProfilesBytesSizer{}
+	default:
+		return nil, errors.New("unknown sizer type")
+	}
+
 	if r2 != nil {
 		req2, ok := r2.(*profilesRequest)
 		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 maxSize == 0 {
 		return []exporterhelper.Request{req}, nil
 	}
-	return req.split(maxSize)
+	return req.split(maxSize, sz), nil
 }
 
-func (req *profilesRequest) mergeTo(dst *profilesRequest) {
-	dst.setCachedItemsCount(dst.ItemsCount() + req.ItemsCount())
-	req.setCachedItemsCount(0)
+func (req *profilesRequest) mergeTo(dst *profilesRequest, sz sizer.ProfilesSizer) {
+	if sz != nil {
+		dst.setCachedSize(dst.size(sz) + req.size(sz))
+		req.setCachedSize(0)
+	}
 	req.pd.ResourceProfiles().MoveAndAppendTo(dst.pd.ResourceProfiles())
 }
 
-func (req *profilesRequest) split(maxSize int) ([]exporterhelper.Request, error) {
+func (req *profilesRequest) split(maxSize int, sz sizer.ProfilesSizer) []exporterhelper.Request {
 	var res []exporterhelper.Request
-	for req.ItemsCount() > maxSize {
-		pd := extractProfiles(req.pd, maxSize)
-		size := pd.SampleCount()
-		req.setCachedItemsCount(req.ItemsCount() - size)
-		res = append(res, &profilesRequest{pd: pd, cachedItemsCount: size})
+	for req.size(sz) > maxSize {
+		pd, rmSize := extractProfiles(req.pd, maxSize, sz)
+		req.setCachedSize(req.size(sz) - rmSize)
+		res = append(res, newProfilesRequest(pd))
 	}
 	res = append(res, req)
-	return res, nil
+	return res
 }
 
 // extractProfiles extracts a new profiles with a maximum number of samples.
-func extractProfiles(srcProfiles pprofile.Profiles, count int) pprofile.Profiles {
+func extractProfiles(srcProfiles pprofile.Profiles, capacity int, sz sizer.ProfilesSizer) (pprofile.Profiles, int) {
 	destProfiles := pprofile.NewProfiles()
-	srcProfiles.ResourceProfiles().RemoveIf(func(srcRS pprofile.ResourceProfiles) bool {
-		if count == 0 {
+	capacityLeft := capacity - sz.ProfilesSize(destProfiles)
+	removedSize := 0
+	srcProfiles.ResourceProfiles().RemoveIf(func(srcRP pprofile.ResourceProfiles) bool {
+		// If the no more capacity left just return.
+		if capacityLeft == 0 {
 			return false
 		}
-		needToExtract := samplesCount(srcRS) > count
-		if needToExtract {
-			srcRS = extractResourceProfiles(srcRS, count)
+		rawRpSize := sz.ResourceProfilesSize(srcRP)
+		rpSize := sz.DeltaSize(rawRpSize)
+
+		if rpSize > capacityLeft {
+			extSrcRP, extRpSize := extractResourceProfiles(srcRP, 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 += extRpSize
+			// There represents the delta between the delta sizes.
+			removedSize += rpSize - rawRpSize - (sz.DeltaSize(rawRpSize-extRpSize) - (rawRpSize - extRpSize))
+			// It is possible that for the bytes scenario, the extracted field contains no profiles.
+			// Do not add it to the destination if that is the case.
+			if extSrcRP.ScopeProfiles().Len() > 0 {
+				extSrcRP.MoveTo(destProfiles.ResourceProfiles().AppendEmpty())
+			}
+			return extSrcRP.ScopeProfiles().Len() != 0
 		}
-		count -= samplesCount(srcRS)
-		srcRS.MoveTo(destProfiles.ResourceProfiles().AppendEmpty())
-		return !needToExtract
+		capacityLeft -= rpSize
+		removedSize += rpSize
+		srcRP.MoveTo(destProfiles.ResourceProfiles().AppendEmpty())
+		return true
 	})
-	return destProfiles
+	return destProfiles, removedSize
 }
 
 // extractResourceProfiles extracts profiles and returns a new resource profiles with the specified number of profiles.
-func extractResourceProfiles(srcRS pprofile.ResourceProfiles, count int) pprofile.ResourceProfiles {
-	destRS := pprofile.NewResourceProfiles()
-	destRS.SetSchemaUrl(srcRS.SchemaUrl())
-	srcRS.Resource().CopyTo(destRS.Resource())
-	srcRS.ScopeProfiles().RemoveIf(func(srcSS pprofile.ScopeProfiles) bool {
-		if count == 0 {
+func extractResourceProfiles(srcRP pprofile.ResourceProfiles, capacity int, sz sizer.ProfilesSizer) (pprofile.ResourceProfiles, int) {
+	destRP := pprofile.NewResourceProfiles()
+	destRP.SetSchemaUrl(srcRP.SchemaUrl())
+	srcRP.Resource().CopyTo(destRP.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.ResourceProfilesSize(destRP)
+	removedSize := 0
+
+	srcRP.ScopeProfiles().RemoveIf(func(srcSS pprofile.ScopeProfiles) bool {
+		// If the no more capacity left just return.
+		if capacityLeft == 0 {
 			return false
 		}
-		needToExtract := srcSS.Profiles().Len() > count
-		if needToExtract {
-			srcSS = extractScopeProfiles(srcSS, count)
+
+		rawSlSize := sz.ScopeProfilesSize(srcSS)
+		ssSize := sz.DeltaSize(rawSlSize)
+		if ssSize > capacityLeft {
+			extSrcSS, extSsSize := extractScopeProfiles(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 profiles.
+			// Do not add it to the destination if that is the case.
+			if extSrcSS.Profiles().Len() > 0 {
+				extSrcSS.MoveTo(destRP.ScopeProfiles().AppendEmpty())
+			}
+			return extSrcSS.Profiles().Len() != 0
 		}
-		count -= srcSS.Profiles().Len()
-		srcSS.MoveTo(destRS.ScopeProfiles().AppendEmpty())
-		return !needToExtract
+		capacityLeft -= ssSize
+		removedSize += ssSize
+		srcSS.MoveTo(destRP.ScopeProfiles().AppendEmpty())
+		return true
 	})
-	srcRS.Resource().CopyTo(destRS.Resource())
-	return destRS
+
+	return destRP, removedSize
 }
 
 // extractScopeProfiles extracts profiles and returns a new scope profiles with the specified number of profiles.
-func extractScopeProfiles(srcSS pprofile.ScopeProfiles, count int) pprofile.ScopeProfiles {
+func extractScopeProfiles(srcSS pprofile.ScopeProfiles, capacity int, sz sizer.ProfilesSizer) (pprofile.ScopeProfiles, int) {
 	destSS := pprofile.NewScopeProfiles()
 	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.ScopeProfilesSize(destSS)
+	removedSize := 0
 	srcSS.Profiles().RemoveIf(func(srcProfile pprofile.Profile) bool {
-		if count == 0 {
+		// If the no more capacity left just return.
+		if capacityLeft == 0 {
+			return false
+		}
+		rsSize := sz.DeltaSize(sz.ProfileSize(srcProfile))
+		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
 		srcProfile.MoveTo(destSS.Profiles().AppendEmpty())
-		count--
 		return true
 	})
-	return destSS
-}
-
-// resourceProfilessCount calculates the total number of profiles in the pdata.ResourceProfiles.
-func samplesCount(rs pprofile.ResourceProfiles) int {
-	count := 0
-	rs.ScopeProfiles().RemoveIf(func(ss pprofile.ScopeProfiles) bool {
-		ss.Profiles().RemoveIf(func(sp pprofile.Profile) bool {
-			count += sp.Sample().Len()
-			return false
-		})
-		return false
-	})
-	return count
+	return destSS, removedSize
 }

exporter/exporterhelper/xexporterhelper/profiles_batch_test.go

@@ -12,6 +12,7 @@ import (
 
 	"go.opentelemetry.io/collector/exporter/exporterhelper"
 	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/requesttest"
+	"go.opentelemetry.io/collector/exporter/exporterhelper/internal/sizer"
 	"go.opentelemetry.io/collector/pdata/pprofile"
 	"go.opentelemetry.io/collector/pdata/testdata"
 )
@@ -126,7 +127,11 @@ func TestMergeSplitProfiles(t *testing.T) {
 			require.NoError(t, err)
 			assert.Equal(t, len(tt.expected), len(res))
 			for i, r := range res {
-				assert.Equal(t, tt.expected[i], r)
+				expected := tt.expected[i].(*profilesRequest)
+				actual := r.(*profilesRequest)
+
+				assert.Equal(t, expected.size(&sizer.ProfilesCountSizer{}), actual.size(&sizer.ProfilesCountSizer{}))
+				assert.Equal(t, expected.ItemsCount(), actual.ItemsCount())
 			}
 		})
 	}
@@ -135,7 +140,7 @@ func TestMergeSplitProfiles(t *testing.T) {
 func TestExtractProfiles(t *testing.T) {
 	for i := 0; i < 10; i++ {
 		ld := testdata.GenerateProfiles(10)
-		extractedProfiles := extractProfiles(ld, i)
+		extractedProfiles, _ := extractProfiles(ld, i, &sizer.ProfilesCountSizer{})
 		assert.Equal(t, i, extractedProfiles.SampleCount())
 		assert.Equal(t, 10-i, ld.SampleCount())
 	}

exporter/go.mod

@@ -16,6 +16,7 @@ require (
 	go.opentelemetry.io/collector/extension/extensiontest v0.122.1
 	go.opentelemetry.io/collector/extension/xextension v0.122.1
 	go.opentelemetry.io/collector/pdata v1.28.1
+	go.opentelemetry.io/collector/pdata/pprofile v0.122.1
 	go.opentelemetry.io/collector/pdata/testdata v0.122.1
 	go.opentelemetry.io/collector/pipeline v0.122.1
 	go.opentelemetry.io/otel v1.35.0
@@ -50,7 +51,6 @@ require (
 	go.opentelemetry.io/collector/exporter/xexporter v0.122.1 // indirect
 	go.opentelemetry.io/collector/extension v1.28.1 // indirect
 	go.opentelemetry.io/collector/featuregate v1.28.1 // indirect
-	go.opentelemetry.io/collector/pdata/pprofile v0.122.1 // indirect
 	go.opentelemetry.io/collector/receiver v1.28.1 // indirect
 	go.opentelemetry.io/collector/receiver/receivertest v0.122.1 // indirect
 	go.opentelemetry.io/collector/receiver/xreceiver v0.122.1 // indirect