Respect baggage in Azure messaging integrations

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/EventHubs/AmqpConsumerReceiveAsyncIntegration.cs

@@ -88,9 +88,9 @@ private static void ProcessReceivedEvents(IReadOnlyList<object> eventsList, IAmq
                     eventData.Properties != null)
                 {
                     var extractedContext = AzureMessagingCommon.ExtractContext(eventData.Properties);
-                    if (extractedContext != null)
+                    if (extractedContext.SpanContext != null)
                     {
-                        extractedContexts.Add(extractedContext);
+                        extractedContexts.Add(extractedContext.SpanContext);
                     }
                 }
             }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/AzureFunctionsCommon.cs

@@ -344,7 +344,7 @@ private static PropagationContext ExtractPropagatedContextFromHttp<T>(T context,
             }
         }
 
-        private static PropagationContext ExtractPropagatedContextFromMessaging<T>(T context, string singlePropertyKey, string batchPropertyKey)
+        internal static PropagationContext ExtractPropagatedContextFromMessaging<T>(T context, string singlePropertyKey, string batchPropertyKey)
             where T : IFunctionContext
         {
             try
@@ -370,15 +370,16 @@ private static PropagationContext ExtractPropagatedContextFromMessaging<T>(T con
                 }
 
                 var triggerMetadata = bindingsFeature.Value.TriggerMetadata;
-                var spanContexts = new List<SpanContext>();
+                var extractedContexts = new List<PropagationContext>();
 
                 // Extract from single message properties
                 if (triggerMetadata?.TryGetValue(singlePropertyKey, out var singlePropsObj) == true &&
                     TryParseJson<Dictionary<string, object>>(singlePropsObj, out var singleProps) && singleProps != null)
                 {
-                    if (ExtractSpanContextFromProperties(singleProps) is { } singleContext)
+                    var singleContext = Shared.AzureMessagingCommon.ExtractContext(singleProps);
+                    if (singleContext.SpanContext != null)
                     {
-                        spanContexts.Add(singleContext);
+                        extractedContexts.Add(singleContext);
                     }
                 }
 
@@ -388,27 +389,28 @@ private static PropagationContext ExtractPropagatedContextFromMessaging<T>(T con
                 {
                     foreach (var props in propsArray)
                     {
-                        if (ExtractSpanContextFromProperties(props) is { } batchContext)
+                        var batchContext = Shared.AzureMessagingCommon.ExtractContext(props);
+                        if (batchContext.SpanContext != null)
                         {
-                            spanContexts.Add(batchContext);
+                            extractedContexts.Add(batchContext);
                         }
                     }
                 }
 
-                if (spanContexts.Count == 0)
+                if (extractedContexts.Count == 0)
                 {
                     return default;
                 }
 
-                bool areAllTheSame = spanContexts.Count == 1 ||
-                                     (spanContexts.Count > 1 && AreAllContextsIdentical(spanContexts));
+                bool areAllTheSame = extractedContexts.Count == 1 ||
+                                     (extractedContexts.Count > 1 && AreAllContextsIdentical(extractedContexts));
 
                 if (!areAllTheSame)
                 {
                     Log.Warning("Multiple different contexts found in messages. Using first context for parentship.");
                 }
 
-                return new PropagationContext(spanContexts[0], Baggage.Current, null);
+                return extractedContexts[0];
             }
             catch (Exception ex)
             {
@@ -417,11 +419,6 @@ private static PropagationContext ExtractPropagatedContextFromMessaging<T>(T con
             }
         }
 
-        private static SpanContext? ExtractSpanContextFromProperties(Dictionary<string, object> userProperties)
-        {
-            return Shared.AzureMessagingCommon.ExtractContext(userProperties);
-        }
-
         private static bool TryParseJson<T>(object? jsonObj, [NotNullWhen(true)] out T? result)
             where T : class
         {
@@ -443,17 +440,18 @@ private static bool TryParseJson<T>(object? jsonObj, [NotNullWhen(true)] out T?
             }
         }
 
-        private static bool AreAllContextsIdentical(List<SpanContext> contexts)
+        private static bool AreAllContextsIdentical(List<PropagationContext> contexts)
         {
             if (contexts.Count <= 1)
             {
                 return true;
             }
 
-            var first = contexts[0];
+            var first = contexts[0].SpanContext;
             return contexts.All(ctx =>
-                ctx.TraceId128.Equals(first.TraceId128) &&
-                ctx.SpanId == first.SpanId);
+                ctx.SpanContext != null &&
+                ctx.SpanContext.TraceId128.Equals(first!.TraceId128) &&
+                ctx.SpanContext.SpanId == first.SpanId);
         }
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/ServiceBusReceiverReceiveMessagesAsyncIntegration.cs

@@ -101,9 +101,9 @@ internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, int max
                     serviceBusMessage.ApplicationProperties != null)
                 {
                     var extractedContext = AzureMessagingCommon.ExtractContext(serviceBusMessage.ApplicationProperties);
-                    if (extractedContext != null)
+                    if (extractedContext.SpanContext != null)
                     {
-                        extractedContexts.Add(extractedContext);
+                        extractedContexts.Add(extractedContext.SpanContext);
                     }
                 }
             }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Shared/AzureMessagingCommon.cs

@@ -41,27 +41,25 @@ public static void InjectContext(IDictionary<string, object>? properties, Scope?
         }
 
         /// <summary>
-        /// Extracts trace context from message properties dictionary
+        /// Extracts full propagation context (span context and baggage) from message properties dictionary
         /// </summary>
-        public static SpanContext? ExtractContext(IDictionary<string, object>? properties)
+        public static PropagationContext ExtractContext(IDictionary<string, object>? properties)
         {
             if (properties == null)
             {
-                return null;
+                return default;
             }
 
             try
             {
-                var extractedContext = Tracer.Instance.TracerManager.SpanContextPropagator.Extract(
+                return Tracer.Instance.TracerManager.SpanContextPropagator.Extract(
                     properties,
                     default(DictionaryContextPropagation));
-
-                return extractedContext.SpanContext;
             }
             catch (Exception ex)
             {
-                Log.Warning(ex, "Failed to extract trace context from message properties");
-                return null;
+                Log.Warning(ex, "Failed to extract propagation context from message properties");
+                return default;
             }
         }
 

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

@@ -0,0 +1,137 @@
+// <copyright file="AzureFunctionsCommonTests.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>
+
+#nullable enable
+
+#if !NETFRAMEWORK
+
+using System;
+using System.Collections.Generic;
+using Datadog.Trace.ClrProfiler.AutoInstrumentation.Azure.Functions;
+using Datadog.Trace.Propagators;
+using Datadog.Trace.Vendors.Newtonsoft.Json;
+using FluentAssertions;
+using Xunit;
+
+#pragma warning disable SA1649 // File name should match first type name
+namespace Microsoft.Azure.Functions.Worker.Context.Features
+{
+    internal interface IFunctionBindingsFeature
+    {
+    }
+}
+#pragma warning restore SA1649
+
+#pragma warning disable SA1403 // File may only contain a single namespace
+namespace Datadog.Trace.Tests.ClrProfiler.AutoInstrumentation.Azure.Functions
+{
+#pragma warning restore SA1403
+
+    public class AzureFunctionsCommonTests
+    {
+        [Fact]
+        public void ExtractPropagatedContextFromMessaging_MergesIntoEmptyBaggageCurrent()
+        {
+            var context = CreateMockFunctionContext(
+                propertyKey: "Properties",
+                headerProperties: new Dictionary<string, object>
+                {
+                    ["traceparent"] = $"00-{1:x32}-{1:x16}-01",
+                    ["baggage"] = "user.id=123"
+                });
+
+            Baggage.Current = new Baggage();
+            var extractedContext = AzureFunctionsCommon.ExtractPropagatedContextFromMessaging(
+                context,
+                "Properties",
+                "PropertiesArray");
+
+            extractedContext.MergeBaggageInto(Baggage.Current);
+
+            extractedContext.SpanContext.Should().NotBeNull();
+            extractedContext.Baggage.Should().NotBeNull();
+            extractedContext.Baggage!["user.id"].Should().Be("123");
+
+            Baggage.Current["user.id"].Should().Be("123");
+        }
+
+        [Fact]
+        public void ExtractPropagatedContextFromMessaging_MergesIntoExistingBaggageCurrent()
+        {
+            var context = CreateMockFunctionContext(
+                propertyKey: "Properties",
+                headerProperties: new Dictionary<string, object>
+                {
+                    ["traceparent"] = $"00-{1:x32}-{1:x16}-01",
+                    ["baggage"] = "user.id=123"
+                });
+
+            Baggage.Current = new Baggage
+            {
+                ["existing.key"] = "existing.value",
+                ["user.id"] = "old.value"
+            };
+
+            var extractedContext = AzureFunctionsCommon.ExtractPropagatedContextFromMessaging(
+                context,
+                "Properties",
+                "PropertiesArray");
+
+            extractedContext.MergeBaggageInto(Baggage.Current);
+
+            extractedContext.SpanContext.Should().NotBeNull();
+            extractedContext.Baggage.Should().NotBeNull();
+            extractedContext.Baggage!["user.id"].Should().Be("123");
+
+            Baggage.Current["existing.key"].Should().Be("existing.value");
+            Baggage.Current["user.id"].Should().Be("123");
+            Baggage.Current.Count.Should().Be(2);
+        }
+
+        private static MockFunctionContext CreateMockFunctionContext(string propertyKey, Dictionary<string, object>? headerProperties)
+        {
+            var triggerMetadata = new Dictionary<string, object?>();
+
+            if (headerProperties != null)
+            {
+                var json = JsonConvert.SerializeObject(headerProperties);
+                triggerMetadata[propertyKey] = json;
+            }
+
+            var bindingsFeature = new MockBindingsFeature
+            {
+                TriggerMetadata = triggerMetadata
+            };
+
+            var features = new List<KeyValuePair<Type, object?>>
+            {
+                new(typeof(Microsoft.Azure.Functions.Worker.Context.Features.IFunctionBindingsFeature), bindingsFeature)
+            };
+
+            return new MockFunctionContext
+            {
+                Features = features
+            };
+        }
+
+        // This duck types with tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/Isolated/IFunctionContext.cs
+        private class MockFunctionContext : IFunctionContext
+        {
+            public FunctionDefinitionStruct FunctionDefinition { get; set; }
+
+            public IEnumerable<KeyValuePair<Type, object?>>? Features { get; set; }
+        }
+
+        // This duck types with tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/Functions/Isolated/GrpcBindingsFeatureStruct.cs
+        private class MockBindingsFeature
+        {
+            public IDictionary<string, object?>? TriggerMetadata { get; set; }
+
+            public IDictionary<string, object?>? InputData { get; set; }
+        }
+    }
+}
+
+#endif