test changes

DataDog · Nov 8, 2024 · b6f241a · b6f241a
1 parent e26296e
commit b6f241a
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/internal/extension/extension.go b/internal/extension/extension.go
@@ -181,6 +181,8 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi
 		if contextKey, ok := ctx.Value("internal.contextKey").(string); ok {
 			fmt.Printf("[DEBUG lambda]: contextKey%s\n", contextKey)
 		}
+		priority := getSamplingPriority(functionExecutionSpan.Context())
+		fmt.Printf("[DEBUG lambda copied from agent]: priority: %s\n", priority)
 	}
 	fmt.Printf("[DEBUG lambda] full extension.go context: %+v\n", ctx)
 
@@ -200,6 +202,21 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi
 	}
 }
 
+func getSamplingPriority(sc ddtrace.SpanContext) int {
+	// Default priority (matching -128 from the agent code)
+	priority := -128
+
+	// 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
+			fmt.Printf("[DEBUG] [extension.go] Found sampling priority: %d\n", priority)
+			return priority
+		}
+	}
+	return priority
+}
+
 // defaultStackLength specifies the default maximum size of a stack trace.
 const defaultStackLength = 32