compute stats for serverless-init (#23391)

DataDog · Mar 4, 2024 · 9507b46 · 9507b46
1 parent f235403
commit 9507b46
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 17 deletions.
diff --git a/cmd/serverless-init/tag/tag.go b/cmd/serverless-init/tag/tag.go
@@ -57,6 +57,7 @@ func GetBaseTagsMapWithMetadata(metadata map[string]string) map[string]string {
 	}
 
 	tagsMap["datadog_init_version"] = tags.GetExtensionVersion()
+	tagsMap[tags.ComputeStatsKey] = tags.ComputeStatsValue
 
 	return tagsMap
 }

diff --git a/cmd/serverless-init/tag/tag_test.go b/cmd/serverless-init/tag/tag_test.go
@@ -17,7 +17,7 @@ import (
 )
 
 func TestGetBaseTagsArrayNoEnvNoMetadata(t *testing.T) {
-	assert.Equal(t, 1, len(GetBaseTagsArrayWithMetadataTags(make(map[string]string, 0))))
+	assert.Equal(t, 2, len(GetBaseTagsArrayWithMetadataTags(make(map[string]string, 0))))
 }
 
 func TestGetBaseTagsArrayWithMetadataTagsNoMetadata(t *testing.T) {
@@ -28,11 +28,12 @@ func TestGetBaseTagsArrayWithMetadataTagsNoMetadata(t *testing.T) {
 	t.Setenv("DD_VERSION", "123.4")
 	tags := GetBaseTagsArrayWithMetadataTags(make(map[string]string, 0))
 	sort.Strings(tags)
-	assert.Equal(t, 4, len(tags))
-	assert.Contains(t, tags[0], "datadog_init_version")
-	assert.Equal(t, "env:myenv", tags[1])
-	assert.Equal(t, "service:superservice", tags[2])
-	assert.Equal(t, "version:123.4", tags[3])
+	assert.Equal(t, 5, len(tags))
+	assert.Contains(t, tags[0], "_dd.compute_stats:1")
+	assert.Contains(t, tags[1], "datadog_init_version")
+	assert.Equal(t, "env:myenv", tags[2])
+	assert.Equal(t, "service:superservice", tags[3])
+	assert.Equal(t, "version:123.4", tags[4])
 }
 
 func TestGetTagFound(t *testing.T) {
@@ -49,7 +50,7 @@ func TestGetTagNotFound(t *testing.T) {
 }
 
 func TestGetBaseTagsMapNoEnvNoMetadata(t *testing.T) {
-	assert.Equal(t, 1, len(GetBaseTagsMapWithMetadata(make(map[string]string, 0))))
+	assert.Equal(t, 2, len(GetBaseTagsMapWithMetadata(make(map[string]string, 0))))
 }
 
 func TestGetBaseTagsMapNoMetadata(t *testing.T) {
@@ -59,7 +60,7 @@ func TestGetBaseTagsMapNoMetadata(t *testing.T) {
 	t.Setenv("DD_SERVICE", "superService")
 	t.Setenv("DD_VERSION", "123.4")
 	tags := GetBaseTagsMapWithMetadata(make(map[string]string, 0))
-	assert.Equal(t, 4, len(tags))
+	assert.Equal(t, 5, len(tags))
 	assert.Equal(t, "myenv", tags["env"])
 	assert.Equal(t, "superservice", tags["service"])
 	assert.Equal(t, "123.4", tags["version"])
@@ -71,7 +72,7 @@ func TestGetBaseTagsMapWithMetadata(t *testing.T) {
 		"location":      "mysuperlocation",
 		"othermetadata": "mysuperothermetadatavalue",
 	})
-	assert.Equal(t, 3, len(tags))
+	assert.Equal(t, 4, len(tags))
 	assert.Equal(t, "mysuperlocation", tags["location"])
 	assert.Equal(t, "mysuperothermetadatavalue", tags["othermetadata"])
 }
@@ -83,10 +84,11 @@ func TestGetBaseTagsArrayWithMetadataTags(t *testing.T) {
 		"othermetadata": "mysuperothermetadatavalue",
 	})
 	sort.Strings(tags)
-	assert.Equal(t, 3, len(tags))
-	assert.Contains(t, tags[0], "datadog_init_version")
-	assert.Equal(t, "location:mysuperlocation", tags[1])
-	assert.Equal(t, "othermetadata:mysuperothermetadatavalue", tags[2])
+	assert.Equal(t, 4, len(tags))
+	assert.Contains(t, tags[0], "_dd.compute_stats:1")
+	assert.Contains(t, tags[1], "datadog_init_version")
+	assert.Equal(t, "location:mysuperlocation", tags[2])
+	assert.Equal(t, "othermetadata:mysuperothermetadatavalue", tags[3])
 }
 
 func TestDdTags(t *testing.T) {

diff --git a/pkg/serverless/tags/tags.go b/pkg/serverless/tags/tags.go
@@ -56,8 +56,10 @@ const (
 	traceOriginMetadataKey   = "_dd.origin"
 	traceOriginMetadataValue = "lambda"
 
-	computeStatsKey   = "_dd.compute_stats"
-	computeStatsValue = "1"
+	// ComputeStatsKey is the tag key indicating whether trace stats should be computed
+	ComputeStatsKey = "_dd.compute_stats"
+	// ComputeStatsValue is the tag value indicating trace stats should be computed
+	ComputeStatsValue = "1"
 
 	extensionVersionKey = "dd_extension_version"
 
@@ -97,7 +99,7 @@ func BuildTagMap(arn string, configTags []string) map[string]string {
 	tags = MergeWithOverwrite(tags, ArrayToMap(configTags))
 
 	tags = setIfNotEmpty(tags, traceOriginMetadataKey, traceOriginMetadataValue)
-	tags = setIfNotEmpty(tags, computeStatsKey, computeStatsValue)
+	tags = setIfNotEmpty(tags, ComputeStatsKey, ComputeStatsValue)
 	tags = setIfNotEmpty(tags, FunctionARNKey, arn)
 	tags = setIfNotEmpty(tags, extensionVersionKey, GetExtensionVersion())
 
@@ -156,7 +158,7 @@ func MergeWithOverwrite(tags map[string]string, overwritingTags map[string]strin
 
 // BuildTagsFromMap builds an array of tag based on map of tags
 func BuildTagsFromMap(tags map[string]string) []string {
-	tagsMap := buildTags(tags, []string{traceOriginMetadataKey, computeStatsKey})
+	tagsMap := buildTags(tags, []string{traceOriginMetadataKey, ComputeStatsKey})
 	return MapToArray(tagsMap)
 }