test changes

DataDog · Nov 8, 2024 · c1001d4 · c1001d4
1 parent e133343
commit c1001d4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 25 deletions.
diff --git a/internal/extension/extension.go b/internal/extension/extension.go
@@ -14,7 +14,6 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"github.com/DataDog/datadog-lambda-go/internal/logger"
 	"net/http"
 	"os"
 	"reflect"
@@ -23,6 +22,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/DataDog/datadog-lambda-go/internal/logger"
+
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 )
 
@@ -163,14 +164,16 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi
 			req.Header.Set(string(DdSamplingPriority), samplingPriority)
 		}
 	} else {
-		req.Header.Set(string(DdTraceId), fmt.Sprint(functionExecutionSpan.Context().TraceID()))
-		req.Header.Set(string(DdSpanId), fmt.Sprint(functionExecutionSpan.Context().SpanID()))
-
-		// Send sampling priority
-		priority := getSamplingPriority(functionExecutionSpan.Context())
-		fmt.Printf("[DEBUG] [extension.go] [SendEndInvocationRequest()] Sampling priority: %d\n", priority)
-		if priority != nil {
-			req.Header.Set(string(DdSamplingPriority), fmt.Sprint(*priority))
+		spanContext := functionExecutionSpan.Context()
+		req.Header.Set(string(DdTraceId), fmt.Sprint(spanContext.TraceID()))
+		req.Header.Set(string(DdSpanId), fmt.Sprint(spanContext.SpanID()))
+
+		// Try to get sampling priority
+		// Check if the context implements SamplingPriority method
+		if pc, ok := spanContext.(interface{ SamplingPriority() (int, bool) }); ok && pc != nil {
+			if priority, ok := pc.SamplingPriority(); ok {
+				req.Header.Set(string(DdSamplingPriority), fmt.Sprint(priority))
+			}
 		}
 	}
 
@@ -180,17 +183,6 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi
 	}
 }
 
-func getSamplingPriority(sc ddtrace.SpanContext) *int {
-	// Check if the context implements SamplingPriority method
-	if pc, ok := sc.(interface{ SamplingPriority() (int, bool) }); ok && pc != nil {
-		if p, ok := pc.SamplingPriority(); ok {
-			priority := p
-			return &priority
-		}
-	}
-	return nil
-}
-
 // defaultStackLength specifies the default maximum size of a stack trace.
 const defaultStackLength = 32
 

diff --git a/internal/trace/listener.go b/internal/trace/listener.go
@@ -73,8 +73,6 @@ func (l *Listener) HandlerStarted(ctx context.Context, msg json.RawMessage) cont
 		return ctx
 	}
 
-	fmt.Printf("[DEBUG] universal instrumentation: %t\v", l.universalInstrumentation)
-	fmt.Printf("[DEBUG] extension running: %t\v", l.extensionManager.IsExtensionRunning())
 	if l.universalInstrumentation && l.extensionManager.IsExtensionRunning() {
 		ctx = l.extensionManager.SendStartInvocationRequest(ctx, msg)
 	}
@@ -115,8 +113,6 @@ func (l *Listener) HandlerStarted(ctx context.Context, msg json.RawMessage) cont
 	return ctx
 }
 
-//
-
 // HandlerFinished ends the function execution span and stops the tracer
 func (l *Listener) HandlerFinished(ctx context.Context, err error) {
 	if functionExecutionSpan != nil {
@@ -163,7 +159,7 @@ func startFunctionExecutionSpan(ctx context.Context, mergeXrayTraces bool, isDdS
 	span := tracer.StartSpan(
 		"aws.lambda", // This operation name will be replaced with the value of the service tag by the Forwarder
 		tracer.SpanType("serverless"),
-		//tracer.ChildOf(parentSpanContext),
+		tracer.ChildOf(parentSpanContext),
 		tracer.ResourceName(resourceName),
 		tracer.Tag("cold_start", ctx.Value("cold_start")),
 		tracer.Tag("function_arn", functionArn),