refactor helper functions into span_pointers.go

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]
-}

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.

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))$`)

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",

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
+}

pkg/serverless/invocationlifecycle/init_test.go → 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))
 		})
 	}
 }