Pass TraceContext to out-of-proc workers

Author: yojagad

src/WebJobs.Script/Description/Rpc/ScriptInvocationContext.cs

@@ -14,6 +14,12 @@ public class ScriptInvocationContext
 
         public ExecutionContext ExecutionContext { get; set; }
 
+        public string Traceparent { get; set; }
+
+        public string Tracestate { get; set; }
+
+        public IEnumerable<KeyValuePair<string, string>> Attributes { get; set; }
+
         public IEnumerable<(string name, DataType type, object val)> Inputs { get; set; }
 
         public Dictionary<string, object> BindingData { get; set; }

src/WebJobs.Script/Description/Rpc/WorkerLanguageInvoker.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reactive.Linq;
@@ -67,6 +68,8 @@ protected override async Task<object> InvokeCore(object[] parameters, FunctionIn
                 Inputs = inputs,
                 ResultSource = new TaskCompletionSource<ScriptInvocationResult>(),
                 AsyncExecutionContext = System.Threading.ExecutionContext.Capture(),
+                Traceparent = Activity.Current?.Id,
+                Attributes = Activity.Current?.Tags,
 
                 // TODO: link up cancellation token to parameter descriptors
                 CancellationToken = CancellationToken.None,

src/WebJobs.Script/Rpc/LanguageWorkerChannel.cs

@@ -5,10 +5,10 @@
 using System.Collections;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reactive.Linq;
-using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
 using Microsoft.AspNetCore.Http;
@@ -300,34 +300,7 @@ internal void SendInvocationRequest(ScriptInvocationContext context)
                         context.ResultSource.SetCanceled();
                         return;
                     }
-
-                    var functionMetadata = context.FunctionMetadata;
-
-                    InvocationRequest invocationRequest = new InvocationRequest()
-                    {
-                        FunctionId = functionMetadata.FunctionId,
-                        InvocationId = context.ExecutionContext.InvocationId.ToString(),
-                    };
-                    foreach (var pair in context.BindingData)
-                    {
-                        if (pair.Value != null)
-                        {
-                            if ((pair.Value is HttpRequest) && IsTriggerMetadataPopulatedByWorker())
-                            {
-                                continue;
-                            }
-                            invocationRequest.TriggerMetadata.Add(pair.Key, pair.Value.ToRpc(_workerChannelLogger, _workerCapabilities));
-                        }
-                    }
-                    foreach (var input in context.Inputs)
-                    {
-                        invocationRequest.InputData.Add(new ParameterBinding()
-                        {
-                            Name = input.name,
-                            Data = input.val.ToRpc(_workerChannelLogger, _workerCapabilities)
-                        });
-                    }
-
+                    InvocationRequest invocationRequest = context.ToRpcInvocationRequest(IsTriggerMetadataPopulatedByWorker(), _workerChannelLogger, _workerCapabilities);
                     _executingInvocations.TryAdd(invocationRequest.InvocationId, context);
 
                     SendStreamingMessage(new StreamingMessage

src/WebJobs.Script/Rpc/MessageExtensions/ScriptInvocationContextExtensions.cs

@@ -0,0 +1,70 @@
+// 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.Linq;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.Description;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.Azure.WebJobs.Script.Rpc
+{
+    internal static class ScriptInvocationContextExtensions
+    {
+        public static InvocationRequest ToRpcInvocationRequest(this ScriptInvocationContext context, bool isTriggerMetadataPopulatedByWorker, ILogger logger, Capabilities capabilities)
+        {
+            InvocationRequest invocationRequest = new InvocationRequest()
+            {
+                FunctionId = context.FunctionMetadata.FunctionId,
+                InvocationId = context.ExecutionContext.InvocationId.ToString(),
+                TraceContext = GetRpcTraceContext(context.Traceparent, context.Tracestate, context.Attributes, logger),
+            };
+
+            foreach (var pair in context.BindingData)
+            {
+                if (pair.Value != null)
+                {
+                    if ((pair.Value is HttpRequest) && isTriggerMetadataPopulatedByWorker)
+                    {
+                        continue;
+                    }
+                    invocationRequest.TriggerMetadata.Add(pair.Key, pair.Value.ToRpc(logger, capabilities));
+                }
+            }
+            foreach (var input in context.Inputs)
+            {
+                invocationRequest.InputData.Add(new ParameterBinding()
+                {
+                    Name = input.name,
+                    Data = input.val.ToRpc(logger, capabilities)
+                });
+            }
+
+            return invocationRequest;
+        }
+
+        internal static RpcTraceContext GetRpcTraceContext(string activityId, string traceStateString, IEnumerable<KeyValuePair<string, string>> tags, ILogger logger)
+        {
+            RpcTraceContext traceContext = new RpcTraceContext
+            {
+                TraceParent = activityId ?? string.Empty,
+                TraceState = traceStateString ?? string.Empty,
+            };
+
+            foreach (KeyValuePair<string, string> tag in tags ?? Enumerable.Empty<KeyValuePair<string, string>>())
+            {
+                if (string.IsNullOrEmpty(tag.Value))
+                {
+                    logger?.LogDebug($"Excluding {tag.Key} from being added to TraceContext.Attributes since it's value is null/empty");
+                }
+                else
+                {
+                    traceContext.Attributes.Add(tag.Key, tag.Value);
+                }
+            }
+
+            return traceContext;
+        }
+    }
+}

test/WebJobs.Script.Tests/Rpc/ScriptInvocationContextExtensionsTests.cs

@@ -0,0 +1,42 @@
+// 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.Linq;
+using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
+using Microsoft.Azure.WebJobs.Script.Rpc;
+using Microsoft.Extensions.Logging.Abstractions;
+using Xunit;
+
+namespace Microsoft.Azure.WebJobs.Script.Tests.Rpc
+{
+    public class ScriptInvocationContextExtensionsTests
+    {
+        [Theory]
+        [InlineData("someTraceParent", "someTraceState", null)]
+        [InlineData("", "", null)]
+        [InlineData(null, null, null)]
+        public void TestGetRpcTraceContext_WithExpectedValues(string traceparent, string tracestate, IEnumerable<KeyValuePair<string, string>> attributes)
+        {
+            if (!string.IsNullOrEmpty(traceparent))
+            {
+                attributes = new List<KeyValuePair<string, string>> { new KeyValuePair<string, string>("key1", "value1") };
+            }
+
+            RpcTraceContext traceContext = ScriptInvocationContextExtensions.GetRpcTraceContext(traceparent, tracestate, attributes, NullLogger.Instance);
+
+            Assert.Equal(traceparent ?? string.Empty, traceContext.TraceParent);
+            Assert.Equal(tracestate ?? string.Empty, traceContext.TraceState);
+
+            if (attributes != null)
+            {
+                Assert.True(attributes.SequenceEqual(traceContext.Attributes));
+            }
+            else
+            {
+                Assert.Equal(0, traceContext.Attributes.Count);
+            }
+        }
+    }
+}