Send sampling priority to lambda extension (#6719)

Author: nhulston

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs

@@ -47,8 +47,11 @@ WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, bool isE
             request.Headers.Set(HttpHeaderNames.TraceId, span.TraceId128.Lower.ToString(CultureInfo.InvariantCulture));
             request.Headers.Set(HttpHeaderNames.SpanId, span.SpanId.ToString(CultureInfo.InvariantCulture));
 
-            var samplingPriority = span.Context.TraceContext?.GetOrMakeSamplingDecision();
-            request.Headers.Set(HttpHeaderNames.SamplingPriority, SamplingPriorityValues.ToString(samplingPriority));
+            if (span.Context.TraceContext is { } traceContext)
+            {
+                var samplingPriority = traceContext.GetOrMakeSamplingDecision(span);
+                request.Headers.Set(HttpHeaderNames.SamplingPriority, SamplingPriorityValues.ToString(samplingPriority));
+            }
 
             var errorMessage = span.GetTag("error.msg");
             if (errorMessage != null)

tracer/src/Datadog.Trace/SpanContext.cs

@@ -430,7 +430,7 @@ private static TraceId GetTraceId(ISpanContext context, TraceId fallback)
         }
 
         /// <summary>
-        /// If <see cref="TraceContext"/> is not null, returns <see cref="Trace.TraceContext.GetOrMakeSamplingDecision"/>.
+        /// If <see cref="TraceContext"/> is not null, returns <see cref="Trace.TraceContext.GetOrMakeSamplingDecision()"/>.
         /// Otherwise, returns <see cref="SamplingPriority"/>.
         /// </summary>
         internal int? GetOrMakeSamplingDecision() =>

tracer/src/Datadog.Trace/TraceContext.cs

@@ -257,12 +257,12 @@ public int GetOrMakeSamplingDecision()
                 return samplingPriority;
             }
 
-            return GetOrMakeSamplingDecisionSlow();
+            return GetOrMakeSamplingDecision(_rootSpan);
         }
 
-        private int GetOrMakeSamplingDecisionSlow()
+        public int GetOrMakeSamplingDecision(Span? span)
         {
-            if (_rootSpan is null)
+            if (span is null)
             {
                 // we can't make a sampling decision without a root span because:
                 // - we need a trace id, and for now trace id lives in SpanContext, not in TraceContext
@@ -274,7 +274,7 @@ private int GetOrMakeSamplingDecisionSlow()
             }
 
             var samplingDecision = CurrentTraceSettings?.TraceSampler is { } sampler
-                                       ? sampler.MakeSamplingDecision(_rootSpan)
+                                       ? sampler.MakeSamplingDecision(span)
                                        : SamplingDecision.Default;
 
             SetSamplingPriority(samplingDecision.Priority, samplingDecision.Mechanism);

tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs

@@ -55,14 +55,14 @@ public async Task TestGetEndInvocationRequestWithoutError()
         public async Task TestGetEndInvocationRequestWithScope()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap();
             var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
             var request = requestBuilder.GetEndInvocationRequest(scope, isError: false);
             request.Headers.Get("x-datadog-invocation-error").Should().BeNull();
             request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false");
-            request.Headers.Get("x-datadog-sampling-priority").Should().Be("-1");
+            request.Headers.Get("x-datadog-sampling-priority").Should().Be("1");
             request.Headers.Get("x-datadog-trace-id").Should().Be("1234");
             request.Headers.Get("x-datadog-span-id").Should().NotBeNull();
         }