diff --git a/protos/azure-functions-language-worker-protobuf/src/proto/FunctionRpc.proto b/protos/azure-functions-language-worker-protobuf/src/proto/FunctionRpc.proto index 64162e12d..616451f11 100644 --- a/protos/azure-functions-language-worker-protobuf/src/proto/FunctionRpc.proto +++ b/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 { diff --git a/src/DotNetWorker.Core/Diagnostics/TraceConstants.cs b/src/DotNetWorker.Core/Diagnostics/TraceConstants.cs index d6120c997..4947264f1 100644 --- a/src/DotNetWorker.Core/Diagnostics/TraceConstants.cs +++ b/src/DotNetWorker.Core/Diagnostics/TraceConstants.cs @@ -1,6 +1,9 @@ // 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; +using System.Collections.Immutable; + namespace Microsoft.Azure.Functions.Worker.Diagnostics; internal static class TraceConstants @@ -40,8 +43,21 @@ public static class OTelAttributes_1_37_0 public const string SchemaVersion = "https://opentelemetry.io/schemas/1.37.0"; } + public static class KnownAttributes + { + public static ImmutableHashSet All { get; } = ImmutableHashSet.Create( + 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 = "AzureFunctions_ActivityTags"; public const string FunctionInvocationId = "AzureFunctions_InvocationId"; public const string FunctionName = "AzureFunctions_FunctionName"; public const string HostInstanceId = "HostInstanceId"; diff --git a/src/DotNetWorker.Core/FunctionsApplication.cs b/src/DotNetWorker.Core/FunctionsApplication.cs index 40157f9c0..fa794b89d 100644 --- a/src/DotNetWorker.Core/FunctionsApplication.cs +++ b/src/DotNetWorker.Core/FunctionsApplication.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Concurrent; +using System.Collections.Generic; using System.Diagnostics; using System.Linq; using System.Threading; @@ -74,6 +75,16 @@ 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 ??= new ConcurrentDictionary(); + context.Items[TraceConstants.InternalKeys.FunctionContextItemsKey] = tags; + } } catch (Exception ex) { diff --git a/src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs b/src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs index 07a5300ce..8871e0d9e 100644 --- a/src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs +++ b/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 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,39 @@ 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 = { } + }; + + if (context.Items is null) + { + return traceContext; + } + + var tags = context.Items.TryGetValue(TraceConstants.InternalKeys.FunctionContextItemsKey, out var tagsObj) + ? tagsObj as System.Collections.Generic.IDictionary + : null; + + if (tags is not null) + { + var known = TraceConstants.KnownAttributes.All; + + foreach (var (key, value) in tags) + { + if (!known.Contains(key)) + { + traceContext.Attributes[key] = value; + } + } + } + + return traceContext; + } } } diff --git a/test/DotNetWorker.Tests/Handlers/InvocationHandlerTests.cs b/test/DotNetWorker.Tests/Handlers/InvocationHandlerTests.cs index c7e37dec2..d148340d9 100644 --- a/test/DotNetWorker.Tests/Handlers/InvocationHandlerTests.cs +++ b/test/DotNetWorker.Tests/Handlers/InvocationHandlerTests.cs @@ -1,11 +1,14 @@ -// 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; using Azure.Core.Serialization; +using Microsoft.Azure.Functions.Worker.ApplicationInsights; using Microsoft.Azure.Functions.Worker.Context.Features; using Microsoft.Azure.Functions.Worker.Grpc.Messages; using Microsoft.Azure.Functions.Worker.Handlers; @@ -261,6 +264,57 @@ 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> + { + new("customTag1", "value1"), + new("customTag2", "value2"), + new("customTag1", "value1-latest") + }; + + var expectedTags = activityTags + .GroupBy(kv => kv.Key) + .ToDictionary(g => g.Key, g => g.Last().Value); + + // Set up the application to add tags to the context.Items + _mockApplication + .Setup(m => m.InvokeFunctionAsync(It.IsAny())) + .Callback(ctx => + { + // Ensure Items is initialized before use + ctx.Items ??= new ConcurrentDictionary(); + + // Simulate tags being set in the Items dictionary + ctx.Items[Worker.Diagnostics.TraceConstants.InternalKeys.FunctionContextItemsKey] = + new Dictionary(expectedTags); + }) + .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.TraceContext?.Attributes; + 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 = null) {