From f989f53b101e41f7c459ffa955835788e54ab410 Mon Sep 17 00:00:00 2001 From: Nicholas Hulston Date: Fri, 13 Dec 2024 13:22:29 -0500 Subject: [PATCH] refactor helper functions into span_pointers.go --- pkg/serverless/invocationlifecycle/init.go | 28 +----------- .../invocationlifecycle/lifecycle.go | 10 +---- pkg/serverless/invocationlifecycle/trace.go | 2 +- .../invocationlifecycle/trace_test.go | 7 +-- pkg/serverless/spanpointers/span_pointers.go | 43 +++++++++++++++++++ .../span_pointers_test.go} | 18 ++++---- 6 files changed, 60 insertions(+), 48 deletions(-) create mode 100644 pkg/serverless/spanpointers/span_pointers.go rename pkg/serverless/{invocationlifecycle/init_test.go => spanpointers/span_pointers_test.go} (89%) diff --git a/pkg/serverless/invocationlifecycle/init.go b/pkg/serverless/invocationlifecycle/init.go index 1f7958f9ef8d4..b399000054fd9 100644 --- a/pkg/serverless/invocationlifecycle/init.go +++ b/pkg/serverless/invocationlifecycle/init.go @@ -6,9 +6,8 @@ package invocationlifecycle import ( - "crypto/sha256" - "encoding/hex" "fmt" + "github.com/DataDog/datadog-agent/pkg/serverless/spanpointers" "strings" "time" @@ -24,7 +23,6 @@ import ( const ( tagFunctionTriggerEventSource = "function_trigger.event_source" tagFunctionTriggerEventSourceArn = "function_trigger.event_source_arn" - s3PointerKind = "aws.s3.object" ) func (lp *LifecycleProcessor) initFromAPIGatewayEvent(event events.APIGatewayProxyRequest, region string) { @@ -133,7 +131,7 @@ func (lp *LifecycleProcessor) initFromS3Event(event events.S3Event) { if lp.InferredSpansEnabled { lp.GetInferredSpan().EnrichInferredSpanWithS3Event(event) } - lp.addSpanPointersFromS3Event(event) + lp.requestHandler.spanPointers = spanpointers.GetSpanPointersFromS3Event(event) } lp.requestHandler.event = event @@ -230,25 +228,3 @@ func (lp *LifecycleProcessor) initFromLambdaFunctionURLEvent(event events.Lambda func (lp *LifecycleProcessor) initFromStepFunctionPayload(event events.StepFunctionPayload) { lp.requestHandler.event = event } - -func (lp *LifecycleProcessor) addSpanPointersFromS3Event(event events.S3Event) { - for _, record := range event.Records { - bucketName := record.S3.Bucket.Name - key := record.S3.Object.Key - eTag := strings.Trim(record.S3.Object.ETag, "\"") - - hash := generateSpanPointerHash([]string{bucketName, key, eTag}) - - spanPointer := SpanPointer{ - Hash: hash, - Kind: s3PointerKind, - } - lp.requestHandler.spanPointers = []SpanPointer{spanPointer} - } -} - -func generateSpanPointerHash(components []string) string { - dataToHash := strings.Join(components, "|") - sum := sha256.Sum256([]byte(dataToHash)) - return hex.EncodeToString(sum[:])[:32] -} diff --git a/pkg/serverless/invocationlifecycle/lifecycle.go b/pkg/serverless/invocationlifecycle/lifecycle.go index 0f4317248ea21..9c5dda6f12ce2 100644 --- a/pkg/serverless/invocationlifecycle/lifecycle.go +++ b/pkg/serverless/invocationlifecycle/lifecycle.go @@ -18,6 +18,7 @@ import ( pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/trace" serverlessLog "github.com/DataDog/datadog-agent/pkg/serverless/logs" serverlessMetrics "github.com/DataDog/datadog-agent/pkg/serverless/metrics" + "github.com/DataDog/datadog-agent/pkg/serverless/spanpointers" "github.com/DataDog/datadog-agent/pkg/serverless/trace/inferredspan" "github.com/DataDog/datadog-agent/pkg/serverless/trace/propagation" "github.com/DataDog/datadog-agent/pkg/serverless/trigger" @@ -50,14 +51,7 @@ type RequestHandler struct { inferredSpans [2]*inferredspan.InferredSpan triggerTags map[string]string triggerMetrics map[string]float64 - spanPointers []SpanPointer -} - -// SpanPointer is a struct that stores a hash and span kind to uniquely -// identify a S3 or DynamoDB operation. -type SpanPointer struct { - Hash string - Kind string + spanPointers []spanpointers.SpanPointer } // SetMetaTag sets a meta span tag. A meta tag is a tag whose value type is string. diff --git a/pkg/serverless/invocationlifecycle/trace.go b/pkg/serverless/invocationlifecycle/trace.go index d966cb5a887a6..5e3e1b03fc54a 100644 --- a/pkg/serverless/invocationlifecycle/trace.go +++ b/pkg/serverless/invocationlifecycle/trace.go @@ -30,7 +30,7 @@ import ( const ( functionNameEnvVar = "AWS_LAMBDA_FUNCTION_NAME" spanPointerLinkKind = "span-pointer" - spanPointerUpDirection = "u" + spanPointerUpDirection = "u" // Tracers will handle cases where direction is down ) var /* const */ runtimeRegex = regexp.MustCompile(`^(dotnet|go|java|ruby)(\d+(\.\d+)*|\d+(\.x))$`) diff --git a/pkg/serverless/invocationlifecycle/trace_test.go b/pkg/serverless/invocationlifecycle/trace_test.go index f37e9f391132a..b429ca90defc8 100644 --- a/pkg/serverless/invocationlifecycle/trace_test.go +++ b/pkg/serverless/invocationlifecycle/trace_test.go @@ -7,6 +7,7 @@ package invocationlifecycle import ( "encoding/json" + "github.com/DataDog/datadog-agent/pkg/serverless/spanpointers" "net/http" "testing" "time" @@ -813,12 +814,12 @@ func TestEndExecutionSpanWithStepFunctions(t *testing.T) { func TestEndExecutionSpanWithSpanLinks(t *testing.T) { tests := []struct { name string - spanPointers []SpanPointer + spanPointers []spanpointers.SpanPointer expectedHashes []string }{ { name: "single span pointer", - spanPointers: []SpanPointer{ + spanPointers: []spanpointers.SpanPointer{ { Hash: "abc", Kind: "aws.s3.object", @@ -828,7 +829,7 @@ func TestEndExecutionSpanWithSpanLinks(t *testing.T) { }, { name: "multiple span pointers", - spanPointers: []SpanPointer{ + spanPointers: []spanpointers.SpanPointer{ { Hash: "def", Kind: "aws.s3.object", diff --git a/pkg/serverless/spanpointers/span_pointers.go b/pkg/serverless/spanpointers/span_pointers.go new file mode 100644 index 0000000000000..1aa315538e3b8 --- /dev/null +++ b/pkg/serverless/spanpointers/span_pointers.go @@ -0,0 +1,43 @@ +package spanpointers + +import ( + "crypto/sha256" + "encoding/hex" + "github.com/DataDog/datadog-agent/pkg/serverless/trigger/events" + "strings" +) + +const ( + s3PointerKind = "aws.s3.object" +) + +// SpanPointer is a struct that stores a hash and span kind to uniquely +// identify a S3 or DynamoDB operation. +type SpanPointer struct { + Hash string + Kind string +} + +func generateSpanPointerHash(components []string) string { + dataToHash := strings.Join(components, "|") + sum := sha256.Sum256([]byte(dataToHash)) + return hex.EncodeToString(sum[:])[:32] +} + +func GetSpanPointersFromS3Event(event events.S3Event) []SpanPointer { + var pointers []SpanPointer + for _, record := range event.Records { + bucketName := record.S3.Bucket.Name + key := record.S3.Object.Key + eTag := strings.Trim(record.S3.Object.ETag, "\"") + + hash := generateSpanPointerHash([]string{bucketName, key, eTag}) + + spanPointer := SpanPointer{ + Hash: hash, + Kind: s3PointerKind, + } + pointers = append(pointers, spanPointer) + } + return pointers +} diff --git a/pkg/serverless/invocationlifecycle/init_test.go b/pkg/serverless/spanpointers/span_pointers_test.go similarity index 89% rename from pkg/serverless/invocationlifecycle/init_test.go rename to pkg/serverless/spanpointers/span_pointers_test.go index 9a611ff2cd689..7fc0654204d6f 100644 --- a/pkg/serverless/invocationlifecycle/init_test.go +++ b/pkg/serverless/spanpointers/span_pointers_test.go @@ -3,7 +3,7 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/). // Copyright 2016-present Datadog, Inc. -package invocationlifecycle +package spanpointers import ( "github.com/DataDog/datadog-agent/pkg/serverless/trigger/events" @@ -42,7 +42,7 @@ func TestGenerateSpanPointerHash(t *testing.T) { } } -func TestAddSpanPointersFromS3Event(t *testing.T) { +func TestGetSpanPointersFromS3Event(t *testing.T) { tests := []struct { name string event events.S3Event @@ -93,7 +93,7 @@ func TestAddSpanPointersFromS3Event(t *testing.T) { }, }, { - name: "multiple invocations - should only keep most recent pointer", + name: "multiple invocations", event: events.S3Event{ Records: []events.S3EventRecord{ { @@ -117,6 +117,10 @@ func TestAddSpanPointersFromS3Event(t *testing.T) { }, }, expectedSpanPointers: []SpanPointer{ + { + Hash: "1294423cd905a1041b4cda23022e476a", + Kind: s3PointerKind, + }, { Hash: "e721375466d4116ab551213fdea08413", Kind: s3PointerKind, @@ -132,13 +136,7 @@ func TestAddSpanPointersFromS3Event(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - lp := &LifecycleProcessor{ - requestHandler: &RequestHandler{}, - } - - lp.addSpanPointersFromS3Event(tt.event) - - assert.Equal(t, tt.expectedSpanPointers, lp.requestHandler.spanPointers) + assert.Equal(t, tt.expectedSpanPointers, GetSpanPointersFromS3Event(tt.event)) }) } }