Adding a null check before initiating the internal Activity (#2765)

* Adding a null check for the internal Activity.
Azure · Oct 14, 2024 · f2c700d · f2c700d
1 parent ad547ba
commit f2c700d
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 10 deletions.
diff --git a/src/DotNetWorker.Core/FunctionsApplication.cs b/src/DotNetWorker.Core/FunctionsApplication.cs
@@ -67,18 +67,23 @@ public void LoadFunction(FunctionDefinition definition)
 
         public async Task InvokeFunctionAsync(FunctionContext context)
         {
-            // This will act as an internal activity that represents remote Host activity. This cannot be tracked as this is not associate to an ActivitySource. 
-            using Activity activity = new Activity(nameof(InvokeFunctionAsync));
-            activity.Start();
+            Activity? activity = null;
 
-            if (ActivityContext.TryParse(context.TraceContext.TraceParent, context.TraceContext.TraceState, true, out ActivityContext activityContext))
+            if (Activity.Current is null)
             {
-                activity.SetId(context.TraceContext.TraceParent);
-                activity.SetSpanId(activityContext.SpanId.ToString());
-                activity.SetTraceId(activityContext.TraceId.ToString());
-                activity.SetRootId(activityContext.TraceId.ToString());
-                activity.ActivityTraceFlags = activityContext.TraceFlags;
-                activity.TraceStateString = activityContext.TraceState;
+                // This will act as an internal activity that represents remote Host activity. This cannot be tracked as this is not associate to an ActivitySource. 
+                activity = new Activity(nameof(InvokeFunctionAsync));
+                activity.Start();
+
+                if (ActivityContext.TryParse(context.TraceContext.TraceParent, context.TraceContext.TraceState, true, out ActivityContext activityContext))
+                {
+                    activity.SetId(context.TraceContext.TraceParent);
+                    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);
@@ -98,6 +103,9 @@ public async Task InvokeFunctionAsync(FunctionContext context)
 
                 throw;
             }
+
+            invokeActivity?.Stop();
+            activity?.Stop();
         }
     }
 }
diff --git a/test/DotNetWorker.OpenTelemetry.Tests/EndToEndTests.cs b/test/DotNetWorker.OpenTelemetry.Tests/EndToEndTests.cs
@@ -91,6 +91,24 @@ public async Task ContextPropagation()
         }   
     }
 
+    [Fact]
+    public async Task ContextPropagationWithTriggerInstrumentation()
+    {
+        using var host = InitializeHost();
+        var context = CreateContext(host);
+        using Activity testActivity = new Activity("ASPNetCoreMockActivity");
+        testActivity.Start();
+        await _application.InvokeFunctionAsync(context);
+        var activity = OtelFunctionDefinition.LastActivity;
+
+        Assert.Equal(activity.Id, testActivity.Id);
+        Assert.Equal(activity.OperationName, testActivity.OperationName);
+        Assert.Equal(activity.SpanId, testActivity.SpanId);
+        Assert.Equal(activity.TraceId, testActivity.TraceId);
+        Assert.Equal(activity.ActivityTraceFlags, testActivity.ActivityTraceFlags);
+        Assert.Equal(activity.TraceStateString, testActivity.TraceStateString);
+    }
+
     [Fact]
     public void ResourceDetectorLocalDevelopment()
     {