refactor where baggage is stored

Author: satvu

src/DotNetWorker.Core/Context/DefaultTraceContext.cs

@@ -7,17 +7,20 @@ namespace Microsoft.Azure.Functions.Worker
 {
     internal sealed class DefaultTraceContext : TraceContext
     {
-        public DefaultTraceContext(string traceParent, string traceState, IReadOnlyDictionary<string, string> attributes)
+        public DefaultTraceContext(string traceParent, string traceState, IReadOnlyDictionary<string, string> attributes, IReadOnlyDictionary<string, string> baggage)
         {
             TraceParent = traceParent;
             TraceState = traceState;
             Attributes = attributes;
+            Baggage = baggage;
         }
 
         public override string TraceParent { get; }
 
         public override string TraceState { get; }
 
         public override IReadOnlyDictionary<string, string> Attributes { get; }
+
+        public override IReadOnlyDictionary<string, string> Baggage { get;  }
     }
 }

src/DotNetWorker.Core/Context/TraceContext.cs

@@ -25,5 +25,10 @@ public abstract class TraceContext
         /// Gets the attributes associated with the trace.
         /// </summary>
         public virtual IReadOnlyDictionary<string, string> Attributes => ImmutableDictionary<string, string>.Empty;
+
+        /// <summary>
+        /// Gets the baggage associated with the trace.
+        /// </summary>
+        public virtual IReadOnlyDictionary<string, string> Baggage => ImmutableDictionary<string, string>.Empty;
     }
 }

src/DotNetWorker.Core/Diagnostics/TraceConstants.cs

@@ -46,7 +46,6 @@ public static class InternalKeys
         public const string FunctionName = "AzureFunctions_FunctionName";
         public const string HostInstanceId = "HostInstanceId";
         public const string AzFuncLiveLogsSessionId = "#AzFuncLiveLogsSessionId";
-        public const string BaggageKeyName = "AzureFunctions_TraceContext_Baggage";
     }
 
     public static class CapabilityFlags

src/DotNetWorker.Grpc/GrpcFunctionInvocation.cs

@@ -13,7 +13,8 @@ internal sealed class GrpcFunctionInvocation : FunctionInvocation, IExecutionRet
         public GrpcFunctionInvocation(InvocationRequest invocationRequest)
         {
             _invocationRequest = invocationRequest;
-            TraceContext = new DefaultTraceContext(_invocationRequest.TraceContext.TraceParent, _invocationRequest.TraceContext.TraceState, _invocationRequest.TraceContext.Attributes);
+            TraceContext = new DefaultTraceContext(_invocationRequest.TraceContext.TraceParent, 
+                _invocationRequest.TraceContext.TraceState, _invocationRequest.TraceContext.Attributes, _invocationRequest.TraceContext.Baggage);
         }
 
         public override string Id => _invocationRequest.InvocationId;

src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs

@@ -87,13 +87,6 @@ public async Task<InvocationResponse> InvokeAsync(InvocationRequest request)
                     invocationFeatures.Set<IInputConversionFeature>(conversion!);
                 }
 
-                var baggage = request.TraceContext?.Baggage;
-
-                if (baggage is not null)
-                {
-                    context.Items[TraceConstants.InternalKeys.BaggageKeyName] = baggage;
-                }
-
                 await _application.InvokeFunctionAsync(context);
 
                 var serializer = _workerOptions.Serializer!;

src/DotNetWorker.OpenTelemetry/BaggageMiddleware.cs

@@ -1,8 +1,6 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
-using System;
-using System.Collections.Generic;
 using System.Threading.Tasks;
 using Microsoft.Azure.Functions.Worker.Middleware;
 using OpenTelemetry;
@@ -15,13 +13,9 @@ public async Task Invoke(FunctionContext context, FunctionExecutionDelegate next
         {
             try
             {
-                if (context.Items.TryGetValue(TraceConstants.BaggageKeyName, out var value) &&
-                    value is IEnumerable<KeyValuePair<string, string>> dict)
+                foreach (var kv in context.TraceContext.Baggage)
                 {
-                    foreach (var kv in dict)
-                    {
-                        Baggage.SetBaggage(kv.Key, kv.Value);
-                    }
+                    Baggage.SetBaggage(kv.Key, kv.Value);
                 }
 
                 await next(context);

src/DotNetWorker.OpenTelemetry/TraceConstants.cs

@@ -17,17 +17,16 @@ public async Task Invoke_SetsAndClearsBaggage_WhenBaggagePresent()
         {
             var baggageKey = "TestKey";
             var baggageValue = "TestValue";
-            var baggageDict = new List<KeyValuePair<string, string>>
+            var baggageDict = new Dictionary<string, string>
             {
-                new KeyValuePair<string, string>(baggageKey, baggageValue)
+                { baggageKey, baggageValue }
             };
 
+            var traceContextMock = new Mock<TraceContext>();
+            traceContextMock.Setup(t => t.Baggage).Returns(baggageDict);
+
             var contextMock = new Mock<FunctionContext>();
-            var items = new Dictionary<object, object>
-            {
-                { TraceConstants.BaggageKeyName, baggageDict }
-            };
-            contextMock.Setup(c => c.Items).Returns(items);
+            contextMock.Setup(c => c.TraceContext).Returns(traceContextMock.Object);
 
             bool nextCalled = false;
             var middleware = new BaggageMiddleware();
@@ -52,17 +51,16 @@ public async Task Invoke_ClearsBaggage_OnException()
         {
             var baggageKey = "TestKey";
             var baggageValue = "TestValue";
-            var baggageDict = new List<KeyValuePair<string, string>>
-        {
-            new KeyValuePair<string, string>(baggageKey, baggageValue)
-        };
+            var baggageDict = new Dictionary<string, string>
+            {
+                { baggageKey, baggageValue }
+            };
+
+            var traceContextMock = new Mock<TraceContext>();
+            traceContextMock.Setup(t => t.Baggage).Returns(baggageDict);
 
             var contextMock = new Mock<FunctionContext>();
-            var items = new Dictionary<object, object>
-        {
-            { TraceConstants.BaggageKeyName, baggageDict }
-        };
-            contextMock.Setup(c => c.Items).Returns(items);
+            contextMock.Setup(c => c.TraceContext).Returns(traceContextMock.Object);
 
             var middleware = new BaggageMiddleware();
 

test/DotNetWorker.OpenTelemetry.Tests/BaggageMiddlewareTests.cs

@@ -273,16 +273,13 @@ public async Task InvokeAsync_SetsBaggage_WhenPresentInInvocationRequest()
             var testValue = "testValue";
             request.TraceContext.Baggage[testKey] = testValue;
 
-            // Setup to verify baggage is set in context.Items during invocation
             _mockApplication
                 .Setup(m => m.InvokeFunctionAsync(It.IsAny<FunctionContext>()))
                 .Returns(Task.CompletedTask)
                 .Callback<FunctionContext>(ctx =>
                 {
-                    Assert.True(ctx.Items.ContainsKey(CoreTraceConstants.InternalKeys.BaggageKeyName));
-                    var actualBaggage = ctx.Items[CoreTraceConstants.InternalKeys.BaggageKeyName] as System.Collections.Generic.IEnumerable<System.Collections.Generic.KeyValuePair<string, string>>;
-                    Assert.NotNull(actualBaggage);
-                    Assert.Contains(actualBaggage, kv => kv.Key == testKey && kv.Value == testValue);
+                    Assert.True(ctx.TraceContext.Baggage.ContainsKey(testKey));
+                    Assert.Equal(testValue, ctx.TraceContext.Baggage[testKey]);
                 });
 
             var handler = CreateInvocationHandler();

test/DotNetWorker.Tests/Handlers/InvocationHandlerTests.cs

@@ -30,7 +30,12 @@ public TestFunctionInvocation(string id = null, string functionId = null)
                 { TraceConstants.InternalKeys.AzFuncLiveLogsSessionId, Guid.NewGuid().ToString() },
             };
 
-            TraceContext = new DefaultTraceContext(activity.Id, Guid.NewGuid().ToString(), attributes);
+            Dictionary<string, string> baggage = new Dictionary<string, string>
+            {
+                { "TestBaggageKey", "TestBaggageValue" }
+            };
+
+            TraceContext = new DefaultTraceContext(activity.Id, Guid.NewGuid().ToString(), attributes, baggage);
         }
 
         public override string Id { get; } = Guid.NewGuid().ToString();