Propagate tags from worker to host

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

@@ -442,6 +442,8 @@ message InvocationResponse {
 
   // Status of the invocation (success/failure/canceled)
   StatusResult result = 3;
+
+  RpcTraceContext trace_context = 5;
 }
 
 message WorkerWarmupRequest {

src/DotNetWorker.Core/Diagnostics/TraceConstants.cs

@@ -1,6 +1,8 @@
 // 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.Generic;
+
 namespace Microsoft.Azure.Functions.Worker.Diagnostics;
 
 internal static class TraceConstants
@@ -40,8 +42,22 @@ public static class OTelAttributes_1_37_0
         public const string SchemaVersion = "https://opentelemetry.io/schemas/1.37.0";
     }
 
+    public static class KnownAttributes
+    {
+        public static IReadOnlyList<string> All { get; } = new List<string>
+        {
+            OTelAttributes_1_17_0.InvocationId,
+            OTelAttributes_1_17_0.SchemaUrl,
+            OTelAttributes_1_37_0.InvocationId,
+            OTelAttributes_1_37_0.FunctionName,
+            OTelAttributes_1_37_0.Instance,
+            OTelAttributes_1_37_0.SchemaUrl,
+        };
+    }
+
     public static class InternalKeys
     {
+        public const string FunctionContextItemsKey = "TagsForFunctionActivity";
         public const string FunctionInvocationId = "AzureFunctions_InvocationId";
         public const string FunctionName = "AzureFunctions_FunctionName";
         public const string HostInstanceId = "HostInstanceId";

src/DotNetWorker.Core/FunctionsApplication.cs

@@ -74,6 +74,13 @@ public async Task InvokeFunctionAsync(FunctionContext context)
             try
             {
                 await _functionExecutionDelegate(context);
+                var tags = invokeActivity?.Tags
+                    .GroupBy(kv => kv.Key)
+                    .ToDictionary(g => g.Key, g => g.Last().Value);
+                if (tags is not null)
+                {
+                    context.Items[TraceConstants.InternalKeys.FunctionContextItemsKey] = tags;
+                }
             }
             catch (Exception ex)
             {

src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs

@@ -1,12 +1,16 @@
-// 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.Diagnostics;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Google.Protobuf.Collections;
 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 +116,9 @@ public async Task<InvocationResponse> InvokeAsync(InvocationRequest request)
                     response.ReturnValue = returnVal;
                 }
 
+                RpcTraceContext traceContext = AddTraceContextTags(request, context);
+                response.TraceContext = traceContext;
+
                 response.Result.Status = StatusResult.Types.Status.Success;
             }
             catch (Exception ex) when (!ex.IsFatal())
@@ -163,5 +170,29 @@ public bool TryCancel(string invocationId)
 
             return false;
         }
+
+        private RpcTraceContext AddTraceContextTags(InvocationRequest request, FunctionContext context)
+        {
+            RpcTraceContext traceContext = new RpcTraceContext
+            {
+                TraceParent = request.TraceContext.TraceParent,
+                TraceState = request.TraceContext.TraceState,
+                Attributes = { }
+            };
+
+            var invocationTags = context.Items.TryGetValue(TraceConstants.InternalKeys.FunctionContextItemsKey, out var tagsObj)
+                ? tagsObj as System.Collections.Generic.IDictionary<string, string>
+                : null;
+
+            if (invocationTags is not null)
+            {
+                foreach (var tag in invocationTags.Where(tag => !TraceConstants.KnownAttributes.All.Contains(tag.Key)))
+                {
+                    traceContext.Attributes[tag.Key] = tag.Value;
+                }
+            }
+
+            return traceContext;
+        }
     }
 }