ActivityTraceFlags propagation (#2608)

* Adding context propagation for ActivityTraceFlags and TraceState.
Azure · Jul 17, 2024 · 18a5b11 · 18a5b11
1 parent 757df9b
commit 18a5b11
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/src/DotNetWorker.Core/Diagnostics/ActivityExtensions.cs b/src/DotNetWorker.Core/Diagnostics/ActivityExtensions.cs
@@ -20,10 +20,12 @@ static ActivityExtensions()
         {
             BindingFlags flags = BindingFlags.NonPublic | BindingFlags.Instance;
             var activityType = typeof(Activity);
-            _setSpanId = activityType.GetField("_spanId", flags).CreateSetter<Activity, string>();
-            _setId = activityType.GetField("_id", flags).CreateSetter<Activity, string>();
-            _setRootId = activityType.GetField("_rootId", flags).CreateSetter<Activity, string>();
-            _setTraceId = activityType.GetField("_traceId", flags).CreateSetter<Activity, string>();
+
+            // Empty setter serves as a safe fallback mechanism to handle cases where the field is not available.
+            _setSpanId = activityType.GetField("_spanId", flags)?.CreateSetter<Activity, string>() ?? ((_, _) => { /* Ignore */ });
+            _setId = activityType.GetField("_id", flags)?.CreateSetter<Activity, string>() ?? ((_, _) => { /* Ignore */ });
+            _setRootId = activityType.GetField("_rootId", flags)?.CreateSetter<Activity, string>() ?? ((_, _) => { /* Ignore */ });
+            _setTraceId = activityType.GetField("_traceId", flags)?.CreateSetter<Activity, string>() ?? ((_, _) => { /* Ignore */ });
         }
 
         /// <summary>

diff --git a/src/DotNetWorker.Core/FunctionsApplication.cs b/src/DotNetWorker.Core/FunctionsApplication.cs
@@ -77,6 +77,8 @@ public async Task InvokeFunctionAsync(FunctionContext context)
                 activity.SetSpanId(activityContext.SpanId.ToString());
                 activity.SetTraceId(activityContext.TraceId.ToString());
                 activity.SetRootId(activityContext.TraceId.ToString());
+                activity.ActivityTraceFlags = activityContext.TraceFlags;
+                activity.TraceStateString = activityContext.TraceState;
             }
 
             var scope = new FunctionInvocationScope(context.FunctionDefinition.Name, context.InvocationId);

diff --git a/test/DotNetWorker.OpenTelemetry.Tests/EndToEndTests.cs b/test/DotNetWorker.OpenTelemetry.Tests/EndToEndTests.cs
@@ -82,6 +82,8 @@ public async Task ContextPropagation()
             Assert.Equal("InvokeFunctionAsync", activity.OperationName);
             Assert.Equal(activity.SpanId, activityContext.SpanId);
             Assert.Equal(activity.TraceId, activityContext.TraceId);
+            Assert.Equal(activity.ActivityTraceFlags, activityContext.TraceFlags);
+            Assert.Equal(activity.TraceStateString, activityContext.TraceState);
         }
         else
         {
@@ -154,7 +156,6 @@ public OtelFunctionDefinition()
         public async Task TestFunction(FunctionContext context)
         {
             LastActivity = Activity.Current;
-            Activity.Current.ActivityTraceFlags = ActivityTraceFlags.Recorded;
             Activity.Current.AddTag("CustomKey", "CustomValue");
 
             var response = new HttpResponseMessage(HttpStatusCode.OK);