[Azure Functions] Fix span parenting in ASP.NET Core integration

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/Isolated/FunctionExecutionMiddlewareInvokeIntegration.cs

@@ -45,6 +45,13 @@ internal static CallTargetState OnMethodBegin<TTarget, TFunctionContext>(TTarget
     [PreserveContext]
     internal static TReturn OnAsyncMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception exception, in CallTargetState state)
     {
+        // The worker's FunctionExecutionMiddleware catches this exception internally,
+        // so the aspnet_core.request span otherwise records status 200. Annotate it here.
+        if (exception is not null && state.State is Scope aspNetCoreScope)
+        {
+            AzureFunctionsCommon.SetExceptionOnAspNetCoreScope(aspNetCoreScope, exception, Tracer.Instance);
+        }
+
         state.Scope?.DisposeWithException(exception);
         return returnValue;
     }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/Isolated/IFunctionContext.cs

@@ -18,6 +18,8 @@ internal interface IFunctionContext
     FunctionDefinitionStruct FunctionDefinition { get; }
 
     IEnumerable<KeyValuePair<Type, object?>>? Features { get; }
+
+    IDictionary<object, object?>? Items { get; }
 }
 
 #endif

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/Isolated/IHttpContextItems.cs

@@ -0,0 +1,23 @@
+// <copyright file="IHttpContextItems.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#if !NETFRAMEWORK
+#nullable enable
+
+using System.Collections.Generic;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.Functions;
+
+/// <summary>
+/// Duck type for Microsoft.AspNetCore.Http.HttpContext,
+/// used to avoid a hard assembly reference to Microsoft.AspNetCore.Http.Abstractions
+/// which may not be available in non-ASP.NET Core Azure Functions workers.
+/// </summary>
+internal interface IHttpContextItems
+{
+    IDictionary<object, object?> Items { get; }
+}
+
+#endif

tracer/src/Datadog.Trace/ClrProfiler/Instrumentation.cs

@@ -85,7 +85,7 @@ private static void PropagateStableConfiguration()
             var tracerSettings = tracer.Settings;
             var mutableSettings = tracerSettings.Manager.InitialMutableSettings;
 
-            NativeInterop.SharedConfig config = new NativeInterop.SharedConfig
+            var config = new NativeInterop.SharedConfig
             {
                 ProfilingEnabled = profilerSettings.ProfilerState switch
                 {
@@ -525,10 +525,53 @@ private static AspNetCoreDiagnosticObserver GetAspNetCoreDiagnosticObserver()
         [Pure]
         private static bool SkipAspNetCoreDiagnosticObserver()
         {
-            // this is extremely simple now, but will get more complex soon...
-            return AzureInfo.Instance.IsAzureFunction;
+            // Enable AspNetCoreDiagnosticObserver in:
+            // - outside Azure Functions
+            // - Isolated functions worker processes with extension v4
+            //   (to create aspnet_core.request spans that azure_functions.invoke can parent to)
+
+            // Skip AspNetCoreDiagnosticObserver in Azure Functions:
+            // - In-process functions (due to AssemblyLoadContext issues)
+            // - Isolated functions host process (to avoid duplicate spans)
+            // - Isolated functions worker process with extension v1 (FUNCTIONS_EXTENSION_VERSION="~1")
+
+            if (!AzureInfo.Instance.IsAzureFunction)
+            {
+                // We only skip AspNetCoreDiagnosticObserver in Azure Functions.
+                // Don't skip it outside Azure Functions.
+                return false;
+            }
+
+            // FUNCTIONS_WORKER_RUNTIME == "dotnet-isolated"
+            if (!AzureInfo.Instance.IsIsolatedFunction)
+            {
+                // Skip AspNetCoreDiagnosticObserver in in-process Azure Functions
+                Log.Debug("Skipping AspNetCoreDiagnosticObserver: running in an in-process Azure Function.");
+                return true;
+            }
+
+            if (AzureInfo.Instance.IsIsolatedFunctionHostProcess)
+            {
+                // Skip AspNetCoreDiagnosticObserver in Azure Functions _host_ processes
+                Log.Debug("Skipping AspNetCoreDiagnosticObserver: running in an isolated Azure Function host process.");
+                return true;
+            }
+
+            // FUNCTIONS_EXTENSION_VERSION
+            var azureFunctionsExtensionVersion = AzureInfo.Instance.AzureFunctionsExtensionVersion;
+
+            if (azureFunctionsExtensionVersion != "~4")
+            {
+                // Skip AspNetCoreDiagnosticObserver in v1 isolated functions (v2 and v3 are not supported at all)
+                // to keep the previous behavior
+                Log.Debug("Skipping AspNetCoreDiagnosticObserver: running in Azure Function with extension version {AzureFunctionsExtensionVersion}.", azureFunctionsExtensionVersion);
+                return true;
+            }
+
+            // do not skip when running in an isolated Azure Functions worker process with extension v4
+            return false;
         }
-#endif
+#endif // #if !NETFRAMEWORK
 
         private static void InitializeDebugger(TracerSettings tracerSettings)
         {

tracer/src/Datadog.Trace/PlatformHelpers/AspNetCoreHttpRequestHandler.cs

@@ -5,10 +5,8 @@
 
 #if !NETFRAMEWORK
 using System;
-using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using Datadog.Trace.Activity;
 using Datadog.Trace.Activity.DuckTypes;
 using Datadog.Trace.Activity.Helpers;
@@ -20,21 +18,22 @@
 using Datadog.Trace.DataStreamsMonitoring.TransactionTracking;
 using Datadog.Trace.DiagnosticListeners;
 using Datadog.Trace.DuckTyping;
-using Datadog.Trace.ExtensionMethods;
 using Datadog.Trace.Headers;
-using Datadog.Trace.Iast;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Propagators;
+using Datadog.Trace.Serverless;
 using Datadog.Trace.Tagging;
 using Datadog.Trace.Util;
 using Datadog.Trace.Util.Http;
+using Datadog.Trace.Vendors.Serilog.Events;
 using Microsoft.AspNetCore.Http;
 
 namespace Datadog.Trace.PlatformHelpers
 {
     internal sealed class AspNetCoreHttpRequestHandler
     {
         internal const string HttpContextTrackingKey = "__Datadog.AspNetCoreHttpRequestHandler.Tracking";
+        internal const string HttpContextActiveScopeKey = "__Datadog.AspNetCoreHttpRequestHandler.ActiveScope";
 
         private readonly IDatadogLogger _log;
         private readonly IntegrationId _integrationId;
@@ -174,6 +173,21 @@ private Scope StartAspNetCorePipelineScope(Tracer tracer, Security security, Ias
             httpContext.Items[HttpContextTrackingKey] = new RequestTrackingFeature(originalPath, scope, proxyContext?.Scope);
 #endif
 
+            if (AzureInfo.Instance.IsAzureFunction)
+            {
+                // Store scope in HttpContext.Items for Azure Functions middleware to retrieve
+                httpContext.Items[HttpContextActiveScopeKey] = scope;
+
+                if (_log.IsEnabled(LogEventLevel.Debug) && scope.Span.Context is { } spanContext)
+                {
+                    _log.Debug(
+                        "AspNetCore: Stored scope in HttpContext.Items, {TraceId}-{SpanId}, path: {Path}",
+                        spanContext.RawTraceId,
+                        spanContext.RawSpanId,
+                        request.Path);
+                }
+            }
+
             if (tracer.Settings.IpHeaderEnabled || security.AppsecEnabled)
             {
                 var peerIp = new Headers.Ip.IpInfo(httpContext.Connection.RemoteIpAddress?.ToString(), httpContext.Connection.RemotePort);

tracer/src/Datadog.Trace/Tagging/AzureFunctionsTags.cs

@@ -34,14 +34,13 @@ internal sealed partial class AzureFunctionsTags : InstrumentationTags
         public string TriggerType { get; set; } = "Unknown";
 
         internal static void SetRootSpanTags(
-            Span span,
+            ITags tags,
             string shortName,
             string fullName,
             string bindingSource,
             string triggerType)
         {
-            var tags = span.Tags;
-            if (span.Tags is AspNetCoreTags aspNetTags)
+            if (tags is AspNetCoreTags aspNetTags)
             {
                 aspNetTags.InstrumentationName = ComponentName;
             }
@@ -51,10 +50,20 @@ internal static void SetRootSpanTags(
                 tags.SetTag(Tags.InstrumentationName, ComponentName);
             }
 
-            tags.SetTag(ShortNameTagName, shortName);
-            tags.SetTag(FullNameTagName, fullName);
-            tags.SetTag(BindingSourceTagName, bindingSource);
-            tags.SetTag(TriggerTypeTagName, triggerType);
+            if (tags is AzureFunctionsTags azureFunctionsTags)
+            {
+                azureFunctionsTags.ShortName = shortName;
+                azureFunctionsTags.FullName = fullName;
+                azureFunctionsTags.BindingSource = bindingSource;
+                azureFunctionsTags.TriggerType = triggerType;
+            }
+            else
+            {
+                tags.SetTag(ShortNameTagName, shortName);
+                tags.SetTag(FullNameTagName, fullName);
+                tags.SetTag(BindingSourceTagName, bindingSource);
+                tags.SetTag(TriggerTypeTagName, triggerType);
+            }
         }
     }
 }

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AzureFunctionsTests.cs

@@ -145,15 +145,15 @@ public InProcessRuntimeV3(ITestOutputHelper output)
         public async Task SubmitsTraces()
         {
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
+
             using (await RunAzureFunctionAndWaitForExit(agent))
             {
                 const int expectedSpanCount = 21;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
                 var filteredSpans = spans.Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase)).ToImmutableList();
 
                 using var s = new AssertionScope();
-                filteredSpans.Count.Should().Be(expectedSpanCount);
-
+                filteredSpans.Should().HaveCount(expectedSpanCount);
                 await AssertInProcessSpans(filteredSpans);
             }
         }
@@ -179,14 +179,15 @@ public InProcessRuntimeV4(ITestOutputHelper output)
         public async Task SubmitsTraces()
         {
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true, useStatsD: true);
+
             using (await RunAzureFunctionAndWaitForExit(agent, framework: "net6.0"))
             {
                 const int expectedSpanCount = 21;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
                 var filteredSpans = spans.Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase)).ToImmutableList();
 
                 using var s = new AssertionScope();
-
+                filteredSpans.Should().HaveCount(expectedSpanCount);
                 await AssertInProcessSpans(filteredSpans);
             }
         }
@@ -213,13 +214,15 @@ public IsolatedRuntimeV4SdkV1(ITestOutputHelper output)
         public async Task SubmitsTraces()
         {
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
+
             using (await RunAzureFunctionAndWaitForExit(agent, expectedExitCode: -1))
             {
                 const int expectedSpanCount = 21;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
                 var filteredSpans = spans.Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase)).ToImmutableList();
-                using var s = new AssertionScope();
 
+                using var s = new AssertionScope();
+                filteredSpans.Should().HaveCount(expectedSpanCount);
                 await AssertIsolatedSpans(filteredSpans, $"{nameof(AzureFunctionsTests)}.Isolated.V4.Sdk1");
             }
         }
@@ -245,16 +248,14 @@ public async Task SubmitsTraces()
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
             using (await RunAzureFunctionAndWaitForExit(agent, expectedExitCode: -1))
             {
-                const int expectedSpanCount = 26;
+                const int expectedSpanCount = 31;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
 
                 var filteredSpans = FilterOutSocketsHttpHandler(spans);
 
                 using var s = new AssertionScope();
-
+                spans.Should().HaveCount(expectedSpanCount);
                 await AssertIsolatedSpans(filteredSpans.ToImmutableList(), $"{nameof(AzureFunctionsTests)}.Isolated.V4.AspNetCore1");
-
-                spans.Count.Should().Be(expectedSpanCount);
             }
         }
     }
@@ -283,27 +284,27 @@ public async Task SubmitsTraces()
             // so we will enable them with a lot of logging
             SetEnvironmentVariable("DD_LOGS_DIRECT_SUBMISSION_AZURE_FUNCTIONS_HOST_ENABLED", "true");
             SetEnvironmentVariable("DD_LOGS_DIRECT_SUBMISSION_MINIMUM_LEVEL", "VERBOSE");
-            var hostName = "integration_ilogger_az_tests";
+            const string hostName = "integration_ilogger_az_tests";
+
             using var logsIntake = new MockLogsIntake();
             EnableDirectLogSubmission(logsIntake.Port, nameof(IntegrationId.ILogger), hostName);
+
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
+
             using (await RunAzureFunctionAndWaitForExit(agent, expectedExitCode: -1))
             {
                 const int expectedSpanCount = 21;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
                 var filteredSpans = spans.Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase)).ToImmutableList();
 
                 using var s = new AssertionScope();
+                filteredSpans.Should().HaveCount(expectedSpanCount);
                 await AssertIsolatedSpans(filteredSpans);
 
-                filteredSpans.Count.Should().Be(expectedSpanCount);
-
-                var logs = logsIntake.Logs;
-
                 // ~327 (ish) logs but we kill func.exe so some logs are lost
                 // and since sometimes the batch of logs can be 100+ it can be a LOT of logs that we lose
                 // so just check that we have much more than when we have host logs disabled
-                logs.Should().HaveCountGreaterThanOrEqualTo(200);
+                logsIntake.Logs.Should().HaveCountGreaterThanOrEqualTo(200);
             }
         }
     }
@@ -331,25 +332,27 @@ public async Task SubmitsTraces()
         {
             SetEnvironmentVariable("DD_LOGS_DIRECT_SUBMISSION_AZURE_FUNCTIONS_HOST_ENABLED", "false");
             SetEnvironmentVariable("DD_LOGS_DIRECT_SUBMISSION_MINIMUM_LEVEL", "VERBOSE");
-            var hostName = "integration_ilogger_az_tests";
+            const string hostName = "integration_ilogger_az_tests";
+
             using var logsIntake = new MockLogsIntake();
             EnableDirectLogSubmission(logsIntake.Port, nameof(IntegrationId.ILogger), hostName);
 
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
+
             using (await RunAzureFunctionAndWaitForExit(agent, expectedExitCode: -1))
             {
                 const int expectedSpanCount = 21;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
-
                 var filteredSpans = spans.Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase)).ToImmutableList();
 
+                using var s = new AssertionScope();
+                filteredSpans.Should().HaveCount(expectedSpanCount);
                 await AssertIsolatedSpans(filteredSpans, filename: $"{nameof(AzureFunctionsTests)}.Isolated.V4.HostLogsDisabled");
-                filteredSpans.Count.Should().Be(expectedSpanCount);
 
-                var logs = logsIntake.Logs;
                 // we expect some logs still from the worker process
                 // this just seems flaky I THINK because of killing the func.exe process (even though we aren't using the host logs)
                 // commonly see 13, 14, 15, 16 logs, but IF we were logging the host logs we'd see 300+
+                var logs = logsIntake.Logs;
                 logs.Should().HaveCountGreaterThan(10);
                 logs.Should().HaveCountLessThanOrEqualTo(20);
             }
@@ -374,9 +377,10 @@ public IsolatedRuntimeV4AspNetCore(ITestOutputHelper output)
         public async Task SubmitsTraces()
         {
             using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
+
             using (await RunAzureFunctionAndWaitForExit(agent, expectedExitCode: -1))
             {
-                const int expectedSpanCount = 26;
+                const int expectedSpanCount = 31;
                 var spans = await agent.WaitForSpansAsync(expectedSpanCount);
 
                 // There are _additional_ spans created for these compared to the non-AspNetCore version
@@ -385,15 +389,13 @@ public async Task SubmitsTraces()
                 // because of this they cause a lot of flake in the snapshots where they shift places
                 // opting to just scrub them from the snapshots - we also don't think that the spans provide much
                 // value so they may be removed from being traced.
-                var filteredSpans = FilterOutSocketsHttpHandler(spans);
-
-                filteredSpans = filteredSpans.Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase)).ToImmutableList();
+                var filteredSpans = FilterOutSocketsHttpHandler(spans)
+                                   .Where(s => !s.Resource.Equals("Timer ExitApp", StringComparison.OrdinalIgnoreCase))
+                                   .ToImmutableList();
 
                 using var s = new AssertionScope();
-
-                await AssertIsolatedSpans(filteredSpans.ToImmutableList(), $"{nameof(AzureFunctionsTests)}.Isolated.V4.AspNetCore");
-
-                spans.Count.Should().Be(expectedSpanCount);
+                spans.Should().HaveCount(expectedSpanCount);
+                await AssertIsolatedSpans(filteredSpans, $"{nameof(AzureFunctionsTests)}.Isolated.V4.AspNetCore");
             }
         }
     }

tracer/test/Datadog.Trace.Tests/ClrProfiler/AutoInstrumentation/Azure/Functions/AzureFunctionsCommonTests.cs

@@ -122,6 +122,8 @@ private class MockFunctionContext : IFunctionContext
             public FunctionDefinitionStruct FunctionDefinition { get; set; }
 
             public IEnumerable<KeyValuePair<Type, object?>>? Features { get; set; }
+
+            public IDictionary<object, object?>? Items { get; }
         }
 
         // This duck types with tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/Isolated/GrpcBindingsFeatureStruct.cs