Add metric origins for sketches for MarshalSplitCompress path (#20990)

pkg/aggregator/time_sampler.go

@@ -120,6 +120,8 @@ func (s *TimeSampler) newSketchSeries(ck ckey.ContextKey, points []metrics.Sketc
 		Interval:   s.interval,
 		Points:     points,
 		ContextKey: ck,
+		Source:     ctx.source,
+		NoIndex:    ctx.noIndex,
 	}
 
 	return ss

pkg/metrics/sketch_series.go

@@ -9,9 +9,10 @@ import (
 	"bytes"
 	"encoding/json"
 
+	"github.com/DataDog/opentelemetry-mapping-go/pkg/quantile"
+
 	"github.com/DataDog/datadog-agent/pkg/aggregator/ckey"
 	"github.com/DataDog/datadog-agent/pkg/tagset"
-	"github.com/DataDog/opentelemetry-mapping-go/pkg/quantile"
 )
 
 // A SketchSeries is a timeseries of quantile sketches.
@@ -22,6 +23,8 @@ type SketchSeries struct {
 	Interval   int64                `json:"interval"`
 	Points     []SketchPoint        `json:"points"`
 	ContextKey ckey.ContextKey      `json:"-"`
+	NoIndex    bool                 `json:"-"` // This is only used by api V2
+	Source     MetricSource         `json:"-"` // This is only used by api V2
 }
 
 // String returns the JSON representation of a SketchSeries as a string

pkg/metrics/test_helper.go

@@ -18,9 +18,10 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/DataDog/opentelemetry-mapping-go/pkg/quantile"
+
 	"github.com/DataDog/datadog-agent/pkg/aggregator/ckey"
 	"github.com/DataDog/datadog-agent/pkg/tagset"
-	"github.com/DataDog/opentelemetry-mapping-go/pkg/quantile"
 )
 
 // AssertPointsEqual evaluate if two list of point are equal (order doesn't matters).
@@ -156,6 +157,14 @@ type SketchesSourceTest struct {
 	currentIndex int
 }
 
+// NewSketchesSourceTestWithSketch populates values with a single test sketch
+func NewSketchesSourceTestWithSketch() *SketchesSourceTest {
+	return &SketchesSourceTest{
+		currentIndex: -1,
+		values:       SketchSeriesList{&SketchSeries{Name: "fakename", Host: "fakehost"}},
+	}
+}
+
 func NewSketchesSourceTest() *SketchesSourceTest {
 	return &SketchesSourceTest{
 		currentIndex: -1,

pkg/serializer/internal/metrics/sketch_series_list.go

@@ -67,6 +67,7 @@ func (sl SketchSeriesList) MarshalSplitCompress(bufferContext *marshaler.BufferC
 	// const sketchDistributions = 3
 	const sketchTags = 4
 	const sketchDogsketches = 7
+	const sketchMetadata = 8
 	// const distributionTs = 1
 	// const distributionCnt = 2
 	// const distributionMin = 3
@@ -86,6 +87,30 @@ func (sl SketchSeriesList) MarshalSplitCompress(bufferContext *marshaler.BufferC
 	const dogsketchK = 7
 	const dogsketchN = 8
 
+	const sketchMetadataOrigin = 1
+	//         |------| 'Metadata' message
+	//                 |-----| 'origin' field index
+	const sketchMetadataOriginMetricType = 3
+	//         |------| 'Metadata' message
+	//                 |----| 'origin' message
+	//                       |--------| 'metric_type' field index
+	const metryTypeNotIndexed = 9
+	//    |-----------------| 'metric_type_agent_hidden' field index
+
+	const sketchMetadataOriginOriginProduct = 4
+	//                 |----|  'Origin' message
+	//                       |-----------| 'origin_product' field index
+	const sketchMetadataOriginOriginCategory = 5
+	//                 |----|  'Origin' message
+	//                       |-----------| 'origin_category' field index
+	const sketchMetadataOriginOriginService = 6
+	//                 |----|  'Origin' message
+	//                       |-----------| 'origin_service' field index
+	const serieMetadataOriginOriginProductAgentType = 10
+	//                 |----|  'Origin' message
+	//                       |-----------| 'OriginProduct' enum
+	//                                    |-------| 'Agent' enum value
+
 	// the backend accepts payloads up to specific compressed / uncompressed
 	// sizes, but prefers small uncompressed payloads.
 	maxPayloadSize := config.Datadog.GetInt("serializer_max_payload_size")
@@ -216,6 +241,28 @@ func (sl SketchSeriesList) MarshalSplitCompress(bufferContext *marshaler.BufferC
 					return err
 				}
 			}
+			err = ps.Embedded(sketchMetadata, func(ps *molecule.ProtoStream) error {
+				return ps.Embedded(sketchMetadataOrigin, func(ps *molecule.ProtoStream) error {
+					if ss.NoIndex {
+						err = ps.Int32(sketchMetadataOriginMetricType, metryTypeNotIndexed)
+						if err != nil {
+							return err
+						}
+					}
+					err = ps.Int32(sketchMetadataOriginOriginProduct, serieMetadataOriginOriginProductAgentType)
+					if err != nil {
+						return err
+					}
+					err = ps.Int32(sketchMetadataOriginOriginCategory, MetricSourceToOriginCategory(ss.Source))
+					if err != nil {
+						return err
+					}
+					return ps.Int32(sketchMetadataOriginOriginService, MetricSourceToOriginService(ss.Source))
+				})
+			})
+			if err != nil {
+				return err
+			}
 
 			return nil
 		})

pkg/serializer/internal/metrics/sketch_series_test.go

@@ -155,8 +155,6 @@ func TestSketchSeriesMarshalSplitCompress(t *testing.T) {
 		sl.Append(Makeseries(i))
 	}
 
-	serializer1 := SketchSeriesList{SketchesSource: sl}
-	payload, _ := serializer1.Marshal()
 	sl.Reset()
 	serializer2 := SketchSeriesList{SketchesSource: sl}
 	payloads, err := serializer2.MarshalSplitCompress(marshaler.NewBufferContext())
@@ -169,9 +167,6 @@ func TestSketchSeriesMarshalSplitCompress(t *testing.T) {
 	decompressed, _ := io.ReadAll(r)
 	r.Close()
 
-	// Check that we encoded the protobuf correctly
-	assert.Equal(t, decompressed, payload)
-
 	pl := new(gogen.SketchPayload)
 	err = pl.Unmarshal(decompressed)
 	require.NoError(t, err)

pkg/serializer/serializer_test.go

@@ -322,11 +322,14 @@ func TestSendSeries(t *testing.T) {
 func TestSendSketch(t *testing.T) {
 	f := &forwarder.MockedForwarder{}
 
-	matcher := createProtoscopeMatcher(`2: {}`)
+	matcher := createProtoscopeMatcher(`
+		1: { 1: {"fakename"} 2: {"fakehost"} 8: { 1: { 4: 10 }}}
+		2: {}
+		`)
 	f.On("SubmitSketchSeries", matcher, protobufExtraHeadersWithCompression).Return(nil).Times(1)
 
 	s := NewSerializer(f, nil)
-	err := s.SendSketch(metrics.NewSketchesSourceTest())
+	err := s.SendSketch(metrics.NewSketchesSourceTestWithSketch())
 	require.Nil(t, err)
 	f.AssertExpectations(t)
 }