[SLVS-4017][serverless-init] Fix incorrect _dd.origin tag for Azure…

… + GCP serverless services (#23169) * fix dd.origin tag for serverless-init * ignore lint * SLVS-4017 - determine origin at init + clean up tests * fix test
DataDog · Feb 29, 2024 · eba2b7b · eba2b7b
1 parent 0cb25f3
commit eba2b7b
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 7 deletions.
diff --git a/cmd/serverless-init/cloudservice/cloudrun.go b/cmd/serverless-init/cloudservice/cloudrun.go
@@ -13,7 +13,8 @@ import (
 
 const (
 	revisionNameEnvVar = "K_REVISION"
-	serviceNameEnvVar  = "K_SERVICE"
+	//nolint:revive // TODO(SERV) Fix revive linter
+	ServiceNameEnvVar = "K_SERVICE"
 )
 
 var metadataHelperFunc = helper.GetMetaData
@@ -26,7 +27,7 @@ func (c *CloudRun) GetTags() map[string]string {
 	tags := metadataHelperFunc(helper.GetDefaultConfig()).TagMap()
 
 	revisionName := os.Getenv(revisionNameEnvVar)
-	serviceName := os.Getenv(serviceNameEnvVar)
+	serviceName := os.Getenv(ServiceNameEnvVar)
 
 	if revisionName != "" {
 		tags["revision_name"] = revisionName
@@ -60,6 +61,6 @@ func (c *CloudRun) Init() error {
 }
 
 func isCloudRunService() bool {
-	_, exists := os.LookupEnv(serviceNameEnvVar)
+	_, exists := os.LookupEnv(ServiceNameEnvVar)
 	return exists
 }
diff --git a/cmd/serverless-init/cloudservice/service_test.go b/cmd/serverless-init/cloudservice/service_test.go
@@ -18,11 +18,11 @@ func TestGetCloudServiceType(t *testing.T) {
 	t.Setenv(ContainerAppNameEnvVar, "test-name")
 	assert.Equal(t, "containerapp", GetCloudServiceType().GetOrigin())
 
-	t.Setenv(serviceNameEnvVar, "test-name")
+	t.Setenv(ServiceNameEnvVar, "test-name")
 	assert.Equal(t, "cloudrun", GetCloudServiceType().GetOrigin())
 
 	os.Unsetenv(ContainerAppNameEnvVar)
-	os.Unsetenv(serviceNameEnvVar)
+	os.Unsetenv(ServiceNameEnvVar)
 	t.Setenv(RunZip, "false")
 	assert.Equal(t, "appservice", GetCloudServiceType().GetOrigin())
 }
diff --git a/pkg/serverless/trace/span_modifer_test.go b/pkg/serverless/trace/span_modifer_test.go
@@ -9,11 +9,13 @@ package trace
 
 import (
 	"context"
+	"os"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/DataDog/datadog-agent/cmd/serverless-init/cloudservice"
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
 	"github.com/DataDog/datadog-agent/pkg/trace/agent"
 	"github.com/DataDog/datadog-agent/pkg/trace/api"
@@ -103,7 +105,7 @@ func TestSpanModifierAddsOriginToAllSpans(t *testing.T) {
 	testOriginTags := func(withModifier bool) {
 		agnt := agent.NewAgent(ctx, cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{})
 		if withModifier {
-			agnt.ModifySpan = (&spanModifier{tags: cfg.GlobalTags}).ModifySpan
+			agnt.ModifySpan = (&spanModifier{tags: cfg.GlobalTags, ddOrigin: getDDOrigin()}).ModifySpan
 		}
 		tc := testutil.RandomTraceChunk(2, 1)
 		tc.Priority = 1 // ensure trace is never sampled out
@@ -143,6 +145,68 @@ func TestSpanModifierAddsOriginToAllSpans(t *testing.T) {
 	testOriginTags(false)
 }
 
+func TestSpanModifierDetectsCloudService(t *testing.T) {
+	cfg := config.New()
+	cfg.Endpoints[0].APIKey = "test"
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	testOriginTags := func(withModifier bool, expectedOrigin string) {
+		agnt := agent.NewAgent(ctx, cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{})
+		if withModifier {
+			agnt.ModifySpan = (&spanModifier{ddOrigin: getDDOrigin()}).ModifySpan
+		}
+		tc := testutil.RandomTraceChunk(2, 1)
+		tc.Priority = 1 // ensure trace is never sampled out
+		tp := testutil.TracerPayloadWithChunk(tc)
+
+		agnt.Process(&api.Payload{
+			TracerPayload: tp,
+			Source:        agnt.Receiver.Stats.GetTagStats(info.Tags{}),
+		})
+
+		select {
+		case ss := <-agnt.TraceWriter.In:
+			tp = ss.TracerPayload
+		case <-time.After(100 * time.Millisecond):
+			t.Fatal("timed out")
+		}
+
+		for _, chunk := range tp.Chunks {
+			if chunk.Origin != expectedOrigin {
+				t.Errorf("chunk should have Origin=%s but has %#v", expectedOrigin, chunk.Origin)
+			}
+			for _, span := range chunk.Spans {
+				tags := span.GetMeta()
+				originVal, ok := tags["_dd.origin"]
+				if withModifier != ok {
+					t.Errorf("unexpected span tags, should have _dd.origin tag %#v: tags=%#v",
+						withModifier, tags)
+				}
+				if withModifier && originVal != expectedOrigin {
+					t.Errorf("got the wrong origin tag value: %#v", originVal)
+					t.Errorf("expected: %#v", expectedOrigin)
+				}
+			}
+		}
+	}
+
+	// Test with and without the span modifier between different cloud services
+	cloudServiceToEnvVar := map[string]string{
+		"cloudrun":     cloudservice.ServiceNameEnvVar,
+		"containerapp": cloudservice.ContainerAppNameEnvVar,
+		"appservice":   cloudservice.RunZip,
+		"lambda":       functionNameEnvVar}
+	for origin, cloudServiceEnvVar := range cloudServiceToEnvVar {
+		// Set the appropriate environment variable to simulate a cloud service
+		t.Setenv(cloudServiceEnvVar, "myService")
+		cfg.GlobalTags = map[string]string{"some": "tag", "_dd.origin": origin}
+		testOriginTags(true, origin)
+		testOriginTags(false, origin)
+		os.Unsetenv(cloudServiceEnvVar)
+	}
+}
+
 func TestLambdaSpanChan(t *testing.T) {
 	cfg := config.New()
 	cfg.GlobalTags = map[string]string{

diff --git a/pkg/serverless/trace/span_modifier.go b/pkg/serverless/trace/span_modifier.go
@@ -23,6 +23,7 @@ type spanModifier struct {
 	lambdaSpanChan chan<- *pb.Span
 	//nolint:revive // TODO(SERV) Fix revive linter
 	coldStartSpanId uint64
+	ddOrigin        string
 }
 
 // ModifySpan applies extra logic to the given span
@@ -39,7 +40,7 @@ func (s *spanModifier) ModifySpan(_ *pb.TraceChunk, span *pb.Span) {
 
 	// ensure all spans have tag _dd.origin in addition to span.Origin
 	if origin := span.Meta[ddOriginTagName]; origin == "" {
-		traceutil.SetMeta(span, ddOriginTagName, ddOriginTagValue)
+		traceutil.SetMeta(span, ddOriginTagName, s.ddOrigin)
 	}
 
 	if span.Name == "aws.lambda.load" {

diff --git a/pkg/serverless/trace/trace.go b/pkg/serverless/trace/trace.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/DataDog/datadog-agent/cmd/serverless-init/cloudservice"
 	compcorecfg "github.com/DataDog/datadog-agent/comp/core/config"
 	comptracecfg "github.com/DataDog/datadog-agent/comp/trace/config"
 	ddConfig "github.com/DataDog/datadog-agent/pkg/config"
@@ -101,6 +102,7 @@ func (s *ServerlessTraceAgent) Start(enabled bool, loadConfig Load, lambdaSpanCh
 			s.spanModifier = &spanModifier{
 				coldStartSpanId: coldStartSpanId,
 				lambdaSpanChan:  lambdaSpanChan,
+				ddOrigin:        getDDOrigin(),
 			}
 
 			s.ta.ModifySpan = s.spanModifier.ModifySpan
@@ -208,3 +210,12 @@ func filterSpanFromLambdaLibraryOrRuntime(span *pb.Span) bool {
 	}
 	return false
 }
+
+// getDDOrigin returns the value for the _dd.origin tag based on the cloud service type
+func getDDOrigin() string {
+	origin := ddOriginTagValue
+	if cloudServiceOrigin := cloudservice.GetCloudServiceType().GetOrigin(); cloudServiceOrigin != "local" {
+		origin = cloudServiceOrigin
+	}
+	return origin
+}
diff --git a/pkg/serverless/trace/trace_test.go b/pkg/serverless/trace/trace_test.go
@@ -9,12 +9,14 @@ package trace
 
 import (
 	"fmt"
+	"os"
 	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/DataDog/datadog-agent/cmd/serverless-init/cloudservice"
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace"
 	"github.com/DataDog/datadog-agent/pkg/serverless/random"
 	"github.com/DataDog/datadog-agent/pkg/trace/config"
@@ -179,3 +181,17 @@ func TestFilterServerlessSpanFromTracer(t *testing.T) {
 	}
 	assert.True(t, filterSpanFromLambdaLibraryOrRuntime(&span))
 }
+
+func TestGetDDOriginCloudServices(t *testing.T) {
+	serviceToEnvVar := map[string]string{
+		"cloudrun":     cloudservice.ServiceNameEnvVar,
+		"appservice":   cloudservice.RunZip,
+		"containerapp": cloudservice.ContainerAppNameEnvVar,
+		"lambda":       functionNameEnvVar,
+	}
+	for service, envVar := range serviceToEnvVar {
+		t.Setenv(envVar, "myService")
+		assert.Equal(t, service, getDDOrigin())
+		os.Unsetenv(envVar)
+	}
+}