Propagate tags from worker to host

protos/azure-functions-language-worker-protobuf/CODEOWNERS

@@ -11,4 +11,4 @@
 # include the following people in the PRs.
 # Language owners should get notified of any new changes to the proto file.
 
-src/proto/FunctionRpc.proto @vrdmr @gavin-aguiar @YunchuWang @surgupta-msft @satvu @ejizba @alrod @anatolib @kaibocai @shreyas-gopalakrishna @amamounelsayed @Francisco-Gamino
+src/proto/FunctionRpc.proto @vrdmr @gavin-aguiar @YunchuWang @surgupta-msft @satvu @alrod @anatolib @kaibocai @shreyas-gopalakrishna @amamounelsayed @Francisco-Gamino

protos/azure-functions-language-worker-protobuf/src/proto/FunctionRpc.proto

@@ -437,11 +437,15 @@ message InvocationResponse {
   // Output binding data
   repeated ParameterBinding output_data = 2;
 
+  // Status of the invocation (success/failure/canceled)
+  StatusResult result = 3;
+
   // data returned from Function (for $return and triggers with return support)
   TypedData return_value = 4;
 
-  // Status of the invocation (success/failure/canceled)
-  StatusResult result = 3;
+  // Enables propagation of tags from worker -> host.
+  map<string, string> trace_context_attributes = 5;
+
 }
 
 message WorkerWarmupRequest {

release_notes.md

@@ -6,12 +6,12 @@
 
 ### Microsoft.Azure.Functions.Worker (metapackage) <version>
 
-- <entry>
+- Update `Microsoft.Azure.Functions.Worker.Grpc` to <version>
 
 ### Microsoft.Azure.Functions.Worker.Core <version>
 
 - <entry>
 
 ### Microsoft.Azure.Functions.Worker.Grpc <version>
 
-- <entry>
+- Update protobuf version to <version> and add support for propagating tags from the worker to the functions host (#3303).

src/DotNetWorker.Core/Diagnostics/Telemetry/TelemetryProviderV1_17_0.cs

@@ -9,7 +9,7 @@ namespace Microsoft.Azure.Functions.Worker.Diagnostics;
 internal sealed class TelemetryProviderV1_17_0 : TelemetryProvider
 {
     private static readonly KeyValuePair<string, object> SchemaUrlAttribute =
-        new(TraceConstants.OTelAttributes_1_17_0.SchemaUrl, TraceConstants.OTelAttributes_1_17_0.SchemaVersion);
+        new(TraceConstants.OTelAttributes_1_17_0.SchemaUrl, TraceConstants.OTel_1_17_0_SchemaVersion);
 
     protected override OpenTelemetrySchemaVersion SchemaVersion
         => OpenTelemetrySchemaVersion.V1_17_0;

src/DotNetWorker.Core/Diagnostics/Telemetry/TelemetryProviderV1_37_0.cs

@@ -9,7 +9,7 @@ namespace Microsoft.Azure.Functions.Worker.Diagnostics;
 internal sealed class TelemetryProviderV1_37_0 : TelemetryProvider
 {
     private static readonly KeyValuePair<string, object> SchemaUrlAttribute =
-        new(TraceConstants.OTelAttributes_1_37_0.SchemaUrl, TraceConstants.OTelAttributes_1_37_0.SchemaVersion);
+        new(TraceConstants.OTelAttributes_1_37_0.SchemaUrl, TraceConstants.OTel_1_37_0_SchemaVersion);
     protected override OpenTelemetrySchemaVersion SchemaVersion
         => OpenTelemetrySchemaVersion.V1_37_0;
 

src/DotNetWorker.Core/Diagnostics/TraceConstants.cs

@@ -1,11 +1,17 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System.Collections.Immutable;
+using System.Linq;
+
 namespace Microsoft.Azure.Functions.Worker.Diagnostics;
 
 internal static class TraceConstants
 {
-   public static class ActivityAttributes
+    public const string OTel_1_37_0_SchemaVersion = "https://opentelemetry.io/schemas/1.37.0";
+    public const string OTel_1_17_0_SchemaVersion = "https://opentelemetry.io/schemas/1.17.0";
+
+    public static class ActivityAttributes
     {
         public static readonly string Version = typeof(ActivityAttributes).Assembly.GetName().Version?.ToString() ?? string.Empty;
         public const string Name = "Microsoft.Azure.Functions.Worker";
@@ -27,7 +33,8 @@ public static class OTelAttributes_1_17_0
         // v1.17.0
         public const string InvocationId = "faas.execution";
         public const string SchemaUrl = "az.schema_url";
-        public const string SchemaVersion = "https://opentelemetry.io/schemas/1.17.0";
+
+        internal static readonly string[] All = { InvocationId, SchemaUrl };
     }
 
     public static class OTelAttributes_1_37_0
@@ -37,7 +44,20 @@ public static class OTelAttributes_1_37_0
         public const string FunctionName = "faas.name";
         public const string Instance = "faas.instance";
         public const string SchemaUrl = "schema.url";
-        public const string SchemaVersion = "https://opentelemetry.io/schemas/1.37.0";
+
+        internal static readonly string[] All = { InvocationId, FunctionName, Instance, SchemaUrl };
+    }
+
+    public static class KnownAttributes
+    {
+        /// <summary>
+        /// Returns protected attribute names that are set by Azure functions that should not be overriden.
+        /// </summary>
+        public static readonly ImmutableHashSet<string> All =
+            OTelAttributes_1_17_0.All
+                .Concat(OTelAttributes_1_37_0.All)
+                .Concat(InternalKeys.All)
+                .ToImmutableHashSet();
     }
 
     public static class InternalKeys
@@ -46,6 +66,13 @@ public static class InternalKeys
         public const string FunctionName = "AzureFunctions_FunctionName";
         public const string HostInstanceId = "HostInstanceId";
         public const string AzFuncLiveLogsSessionId = "#AzFuncLiveLogsSessionId";
+
+        internal static readonly string[] All = { FunctionInvocationId, FunctionName, HostInstanceId, AzFuncLiveLogsSessionId };
+    }
+
+    public static class FunctionContextKeys
+    {
+        public const string FunctionContextItemsKey = "AzureFunctions_ActivityTags";
     }
 
     public static class CapabilityFlags

src/DotNetWorker.Core/FunctionsApplication.cs

@@ -3,10 +3,12 @@
 
 using System;
 using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Azure;
 using Microsoft.Azure.Functions.Worker.Diagnostics;
 using Microsoft.Azure.Functions.Worker.Middleware;
 using Microsoft.Azure.Functions.Worker.Pipeline;
@@ -83,6 +85,30 @@ public async Task InvokeFunctionAsync(FunctionContext context)
 
                 throw;
             }
+            finally
+            {
+                var tags = invokeActivity?.Tags;
+
+                if (tags is not null && context.Items is not null)
+                {
+                    var known = TraceConstants.KnownAttributes.All;
+                    var validTags = new List<KeyValuePair<string, string>>();
+
+                    foreach (var (key, value) in tags)
+                    {
+                        // avoid overwriting protected attributes
+                        if (!known.Contains(key) && value is not null)
+                        {
+                            validTags.Add(new KeyValuePair<string, string>(key, value));
+                        }
+                    }
+
+                    if (validTags.Count > 0)
+                    {
+                        context.Items[TraceConstants.FunctionContextKeys.FunctionContextItemsKey] = validTags;
+                    }
+                }
+            }
         }
     }
 }

src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs

@@ -1,12 +1,15 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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.Concurrent;
+using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Azure.Functions.Worker.Context.Features;
 using Microsoft.Azure.Functions.Worker.Core;
+using Microsoft.Azure.Functions.Worker.Diagnostics;
 using Microsoft.Azure.Functions.Worker.Grpc;
 using Microsoft.Azure.Functions.Worker.Grpc.Features;
 using Microsoft.Azure.Functions.Worker.Grpc.Messages;
@@ -112,6 +115,8 @@ public async Task<InvocationResponse> InvokeAsync(InvocationRequest request)
                     response.ReturnValue = returnVal;
                 }
 
+                AddTraceContextTags(response, context);
+
                 response.Result.Status = StatusResult.Types.Status.Success;
             }
             catch (Exception ex) when (!ex.IsFatal())
@@ -163,5 +168,27 @@ public bool TryCancel(string invocationId)
 
             return false;
         }
+
+        private static void AddTraceContextTags(InvocationResponse response, FunctionContext context)
+        {
+            if (context.Items is null)
+            {
+                return;
+            }
+
+            var tags = context.Items.TryGetValue(TraceConstants.FunctionContextKeys.FunctionContextItemsKey, out var tagsObj)
+                ? tagsObj as IEnumerable<KeyValuePair<string, string>>
+                : null;
+
+            if (tags is null)
+            {
+                return;
+            }
+
+            foreach (var (key, value) in tags)
+            {
+                response.TraceContextAttributes[key] = value;
+            }
+        }
     }
 }

test/DotNetWorker.Tests/Handlers/InvocationHandlerTests.cs

@@ -1,7 +1,9 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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.Concurrent;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -261,6 +263,58 @@ public async Task SetRetryContextToNull()
             Assert.Null(_context.RetryContext);
         }
 
+        [Fact]
+        public async Task InvokeAsync_TagsInFunctionInvocationContextItems_AreIncludedInInvocationResponse()
+        {
+            // Arrange
+            var invocationId = "tags-test-invocation";
+            var request = TestUtility.CreateInvocationRequest(invocationId);
+
+            // Build tags dictionary the same way as Activity and FunctionsApplication would
+            var activityTags = new List<KeyValuePair<string, string>>
+            {
+                new("customTag1", "value1"),
+                new("customTag2", "value2"),
+                new("customTag1", "value3"), // Duplicate key to test handling
+            };
+
+            var expectedTags = new List<KeyValuePair<string, string>>
+            {
+                new("customTag2", "value2"),
+                new("customTag1", "value3"),
+            };
+
+            // Set up the application to add tags to the context.Items
+            _mockApplication
+                .Setup(m => m.InvokeFunctionAsync(It.IsAny<FunctionContext>()))
+                .Callback<FunctionContext>(ctx =>
+                {
+                    // Ensure Items is initialized before use
+                    ctx.Items ??= new ConcurrentDictionary<object, object>();
+
+                    // Simulate tags being set in the Items dictionary
+                    ctx.Items[Worker.Diagnostics.TraceConstants.FunctionContextKeys.FunctionContextItemsKey] = activityTags;
+                })
+                .Returns(Task.CompletedTask);
+
+            var handler = CreateInvocationHandler();
+
+            // Act
+            var response = await handler.InvokeAsync(request);
+
+            // Assert
+            Assert.Equal(StatusResult.Types.Status.Success, response.Result.Status);
+
+            // Extract tags from the response
+            var tags = response.TraceContextAttributes;
+            Assert.NotNull(tags);
+            foreach (var expectedTag in expectedTags)
+            {
+                var tag = tags.FirstOrDefault(t => t.Key == expectedTag.Key);
+                Assert.Equal(expectedTag.Value, tag.Value);
+            }
+        }
+
         private InvocationHandler CreateInvocationHandler(IFunctionsApplication application = null,
                                                           IOptions<WorkerOptions> workerOptions = null)
         {

test/DotNetWorker.Tests/OpenTelemetryTraceConstantsTests.cs

@@ -0,0 +1,39 @@
+// 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.Linq;
+using System.Reflection;
+using Microsoft.Azure.Functions.Worker.Diagnostics;
+using Xunit;
+
+namespace Microsoft.Azure.Functions.Worker.Tests;
+
+public class OpenTelemetryTraceConstantsTests
+{
+    [Theory]
+    [InlineData(typeof(TraceConstants.OTelAttributes_1_17_0))]
+    [InlineData(typeof(TraceConstants.OTelAttributes_1_37_0))]
+    [InlineData(typeof(TraceConstants.InternalKeys))]
+    public void All_Field_Contains_All_Public_Constants(Type classType)
+    {
+        // Get all public const string fields (excluding the "All" field itself)
+        var publicConstants = classType
+            .GetFields(BindingFlags.Public | BindingFlags.Static)
+            .Where(f => f.IsLiteral && !f.IsInitOnly && f.FieldType == typeof(string) && f.Name != "All")
+            .Select(f => f.GetValue(null) as string)
+            .Where(v => v != null)
+            .OrderBy(v => v)
+            .ToArray();
+
+        // Get the "All" field
+        var allField = classType.GetField("All", BindingFlags.NonPublic | BindingFlags.Static);
+        Assert.NotNull(allField);
+
+        var allValues = (allField.GetValue(null) as string[])?.OrderBy(v => v).ToArray();
+        Assert.NotNull(allValues);
+
+        // Verify that "All" contains exactly the same values as all public constants
+        Assert.Equal(publicConstants, allValues);
+    }
+}

test/TestUtility/TestFunctionContext.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -86,7 +86,7 @@ public TestFunctionContext(FunctionDefinition functionDefinition, FunctionInvoca
 
         public override FunctionDefinition FunctionDefinition { get; }
 
-        public override IDictionary<object, object> Items { get; set; }
+        public override IDictionary<object, object> Items { get; set; } = new Dictionary<object, object>();
 
         public override IInvocationFeatures Features { get; } = new InvocationFeatures(Enumerable.Empty<IInvocationFeatureProvider>());