Change StartServerlessTraceAgent to use args struct.

DataDog · Nov 4, 2024 · e39c815 · e39c815
1 parent 48eaa5d
commit e39c815
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 15 deletions.
diff --git a/cmd/serverless-init/main.go b/cmd/serverless-init/main.go
@@ -151,7 +151,11 @@ func setup(_ mode.Conf, tagger tagger.Component) (cloudservice.CloudService, *se
 	return cloudService, agentLogConfig, traceAgent, metricAgent, logsAgent
 }
 func setupTraceAgent(tags map[string]string, tagger tagger.Component) trace.ServerlessTraceAgent {
-	traceAgent := trace.StartServerlessTraceAgent(pkgconfigsetup.Datadog().GetBool("apm_config.enabled"), &trace.LoadConfig{Path: datadogConfigPath, Tagger: tagger}, nil, random.Random.Uint64())
+	traceAgent := trace.StartServerlessTraceAgent(trace.StartServerlessTraceAgentArgs{
+		Enabled:         pkgconfigsetup.Datadog().GetBool("apm_config.enabled"),
+		LoadConfig:      &trace.LoadConfig{Path: datadogConfigPath, Tagger: tagger},
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	traceAgent.SetTags(tags)
 	go func() {
 		for range time.Tick(3 * time.Second) {

diff --git a/cmd/serverless/main.go b/cmd/serverless/main.go
@@ -337,7 +337,12 @@ func startOtlpAgent(wg *sync.WaitGroup, metricAgent *metrics.ServerlessMetricAge
 
 func startTraceAgent(wg *sync.WaitGroup, lambdaSpanChan chan *pb.Span, coldStartSpanId uint64, serverlessDaemon *daemon.Daemon, tagger tagger.Component) {
 	defer wg.Done()
-	traceAgent := trace.StartServerlessTraceAgent(pkgconfigsetup.Datadog().GetBool("apm_config.enabled"), &trace.LoadConfig{Path: datadogConfigPath, Tagger: tagger}, lambdaSpanChan, coldStartSpanId)
+	traceAgent := trace.StartServerlessTraceAgent(trace.StartServerlessTraceAgentArgs{
+		Enabled:         pkgconfigsetup.Datadog().GetBool("apm_config.enabled"),
+		LoadConfig:      &trace.LoadConfig{Path: datadogConfigPath, Tagger: tagger},
+		LambdaSpanChan:  lambdaSpanChan,
+		ColdStartSpanId: coldStartSpanId,
+	})
 	serverlessDaemon.SetTraceAgent(traceAgent)
 }
 

diff --git a/pkg/serverless/daemon/daemon_test.go b/pkg/serverless/daemon/daemon_test.go
@@ -171,7 +171,12 @@ func TestSetTraceTagOk(t *testing.T) {
 	}
 	t.Setenv("DD_API_KEY", "x")
 	t.Setenv("DD_RECEIVER_PORT", strconv.Itoa(testutil.FreeTCPPort(t)))
-	agent := trace.StartServerlessTraceAgent(true, &trace.LoadConfig{Path: "/does-not-exist.yml"}, make(chan *pb.Span), random.Random.Uint64())
+	agent := trace.StartServerlessTraceAgent(trace.StartServerlessTraceAgentArgs{
+		Enabled:         true,
+		LoadConfig:      &trace.LoadConfig{Path: "/does-not-exist.yml"},
+		LambdaSpanChan:  make(chan *pb.Span),
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	defer agent.Stop()
 	d := Daemon{
 		TraceAgent: agent,

diff --git a/pkg/serverless/daemon/routes_test.go b/pkg/serverless/daemon/routes_test.go
@@ -441,7 +441,11 @@ func BenchmarkStartEndInvocation(b *testing.B) {
 func startAgents() *Daemon {
 	d := StartDaemon(fmt.Sprint("127.0.0.1:", testutil.FreeTCPPort(nil)))
 
-	ta := trace.StartServerlessTraceAgent(true, &trace.LoadConfig{Path: "/some/path/datadog.yml"}, nil, 123)
+	ta := trace.StartServerlessTraceAgent(trace.StartServerlessTraceAgentArgs{
+		Enabled:         true,
+		LoadConfig:      &trace.LoadConfig{Path: "/some/path/datadog.yml"},
+		ColdStartSpanId: 123,
+	})
 	d.SetTraceAgent(ta)
 
 	ma := &metrics.ServerlessMetricAgent{

diff --git a/pkg/serverless/otlp/otlp_test.go b/pkg/serverless/otlp/otlp_test.go
@@ -65,7 +65,10 @@ func TestServerlessOTLPAgentReceivesTraces(t *testing.T) {
 	t.Setenv("DD_OTLP_CONFIG_RECEIVER_PROTOCOLS_GRPC_ENDPOINT", grpcEndpoint)
 
 	// setup trace agent
-	traceAgent := trace.StartServerlessTraceAgent(true, &trace.LoadConfig{Path: "./testdata/valid.yml"}, nil, 0)
+	traceAgent := trace.StartServerlessTraceAgent(trace.StartServerlessTraceAgentArgs{
+		Enabled:    true,
+		LoadConfig: &trace.LoadConfig{Path: "./testdata/valid.yml"},
+	})
 	defer traceAgent.Stop()
 	assert.NotNil(traceAgent)
 	traceChan := make(chan struct{})

diff --git a/pkg/serverless/trace/trace.go b/pkg/serverless/trace/trace.go
@@ -91,17 +91,24 @@ func (l *LoadConfig) Load() (*config.AgentConfig, error) {
 	return comptracecfg.LoadConfigFile(l.Path, c, l.Tagger)
 }
 
+type StartServerlessTraceAgentArgs struct {
+	Enabled         bool
+	LoadConfig      Load
+	LambdaSpanChan  chan<- *pb.Span
+	ColdStartSpanId uint64
+}
+
 // Start starts the agent
 //
 //nolint:revive // TODO(SERV) Fix revive linter
-func StartServerlessTraceAgent(enabled bool, loadConfig Load, lambdaSpanChan chan<- *pb.Span, coldStartSpanId uint64) ServerlessTraceAgent {
-	if enabled {
+func StartServerlessTraceAgent(args StartServerlessTraceAgentArgs) ServerlessTraceAgent {
+	if args.Enabled {
 		// Set the serverless config option which will be used to determine if
 		// hostname should be resolved. Skipping hostname resolution saves >1s
 		// in load time between gRPC calls and agent commands.
 		pkgconfigsetup.Datadog().Set("serverless.enabled", true, model.SourceAgentRuntime)
 
-		tc, confErr := loadConfig.Load()
+		tc, confErr := args.LoadConfig.Load()
 		if confErr != nil {
 			log.Errorf("Unable to load trace agent config: %s", confErr)
 		} else {
@@ -110,8 +117,8 @@ func StartServerlessTraceAgent(enabled bool, loadConfig Load, lambdaSpanChan cha
 			tc.SynchronousFlushing = true
 			ta := agent.NewAgent(context, tc, telemetry.NewNoopCollector(), &statsd.NoOpClient{}, zstd.NewComponent())
 			ta.SpanModifier = &spanModifier{
-				coldStartSpanId: coldStartSpanId,
-				lambdaSpanChan:  lambdaSpanChan,
+				coldStartSpanId: args.ColdStartSpanId,
+				lambdaSpanChan:  args.LambdaSpanChan,
 				ddOrigin:        getDDOrigin(),
 			}
 

diff --git a/pkg/serverless/trace/trace_test.go b/pkg/serverless/trace/trace_test.go
@@ -35,7 +35,10 @@ func TestStartEnabledFalse(t *testing.T) {
 	setupTraceAgentTest(t)
 
 	lambdaSpanChan := make(chan *pb.Span)
-	agent := StartServerlessTraceAgent(false, nil, lambdaSpanChan, random.Random.Uint64())
+	agent := StartServerlessTraceAgent(StartServerlessTraceAgentArgs{
+		LambdaSpanChan:  lambdaSpanChan,
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	defer agent.Stop()
 	assert.NotNil(t, agent)
 	assert.IsType(t, noopTraceAgent{}, agent)
@@ -53,7 +56,12 @@ func TestStartEnabledTrueInvalidConfig(t *testing.T) {
 	setupTraceAgentTest(t)
 
 	lambdaSpanChan := make(chan *pb.Span)
-	agent := StartServerlessTraceAgent(true, &LoadConfigMocked{}, lambdaSpanChan, random.Random.Uint64())
+	agent := StartServerlessTraceAgent(StartServerlessTraceAgentArgs{
+		Enabled:         true,
+		LoadConfig:      &LoadConfigMocked{},
+		LambdaSpanChan:  lambdaSpanChan,
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	defer agent.Stop()
 	assert.NotNil(t, agent)
 	assert.IsType(t, noopTraceAgent{}, agent)
@@ -65,7 +73,12 @@ func TestStartEnabledTrueValidConfigInvalidPath(t *testing.T) {
 	lambdaSpanChan := make(chan *pb.Span)
 
 	t.Setenv("DD_API_KEY", "x")
-	agent := StartServerlessTraceAgent(true, &LoadConfig{Path: "invalid.yml"}, lambdaSpanChan, random.Random.Uint64())
+	agent := StartServerlessTraceAgent(StartServerlessTraceAgentArgs{
+		Enabled:         true,
+		LoadConfig:      &LoadConfig{Path: "invalid.yml"},
+		LambdaSpanChan:  lambdaSpanChan,
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	defer agent.Stop()
 	assert.NotNil(t, agent)
 	assert.IsType(t, &serverlessTraceAgent{}, agent)
@@ -76,7 +89,12 @@ func TestStartEnabledTrueValidConfigValidPath(t *testing.T) {
 
 	lambdaSpanChan := make(chan *pb.Span)
 
-	agent := StartServerlessTraceAgent(true, &LoadConfig{Path: "./testdata/valid.yml"}, lambdaSpanChan, random.Random.Uint64())
+	agent := StartServerlessTraceAgent(StartServerlessTraceAgentArgs{
+		Enabled:         true,
+		LoadConfig:      &LoadConfig{Path: "./testdata/valid.yml"},
+		LambdaSpanChan:  lambdaSpanChan,
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	defer agent.Stop()
 	assert.NotNil(t, agent)
 	assert.IsType(t, &serverlessTraceAgent{}, agent)
@@ -89,7 +107,12 @@ func TestLoadConfigShouldBeFast(t *testing.T) {
 	startTime := time.Now()
 	lambdaSpanChan := make(chan *pb.Span)
 
-	agent := StartServerlessTraceAgent(true, &LoadConfig{Path: "./testdata/valid.yml"}, lambdaSpanChan, random.Random.Uint64())
+	agent := StartServerlessTraceAgent(StartServerlessTraceAgentArgs{
+		Enabled:         true,
+		LoadConfig:      &LoadConfig{Path: "./testdata/valid.yml"},
+		LambdaSpanChan:  lambdaSpanChan,
+		ColdStartSpanId: random.Random.Uint64(),
+	})
 	defer agent.Stop()
 	assert.True(t, time.Since(startTime) < time.Second)
 }