Distributed Tracing for Entities (In-Process) (#1190)

* added distributed tracing for entities

* change in how the distributed trace context is propagated in the case that an entity in WebJobs starts an orchestration

* slight style update

* some more tiny style things

* addressing two more PR comments

* forgot to add one annotation

* yet another small style change

* addressing PR comments

* forgot one file

* addressing more PR comments

* forgot two comments

* addressing a few small comments

* refactored most of the tracing into this repo to more accurately reflect the end time of activites (like, when the call/signal request to an entity is actually sent, or when an orchestration is actually created by an entity)

* fixing spacing

* slight change in formatting of the create orchestration trace

* addressing PR comments

* adding back new lines at the ends of files

* trying to fix these line endings

* dealing with new lines again

* tiny change

* missed a null check

* reverting to old design

* missed some

* and missed some more

* missed even more

* will it ever end

* last one i think

* moved activity for entity starting an orchestration back into webjobs

* as always missed some stuff

* added support for an entities enabled flag

* added a null check for client span ID when creating the activity for scheduling a suborchestration

* getting rid of an extra line

* addressing PR comment

* addressing a few PR comments

* addressing some PR comments, slightly changing the logic of the modified TaskHubClient method

---------

Co-authored-by: Sophia Tevosyan <stevosyan@microsoft.com>

Author: sophiatev

src/DurableTask.Core/OrchestrationTags.cs

@@ -36,6 +36,16 @@ public static class OrchestrationTags
         /// </remarks>
         public const string FireAndForget = "FireAndForget";
 
+        /// <summary>
+        /// The ID of the parent trace that created this orchestration instance (see https://www.w3.org/TR/trace-context/#traceparent-header)
+        /// </summary>
+        public const string TraceParent = "MS_Entities_TraceParent";
+
+        /// <summary>
+        /// The trace state of the parent trace that created this orchestration instance (see https://www.w3.org/TR/trace-context/#tracestate-header)
+        /// </summary>
+        public const string TraceState = "MS_Entities_TraceState";
+
         /// <summary>
         /// Check whether the given tags contain the fire and forget tag
         /// </summary>

src/DurableTask.Core/TaskHubClient.cs

@@ -746,6 +746,19 @@ void CreateAndTrackDependencyTelemetry(TraceContextBase? requestTraceContext)
         /// <param name="eventName">Name of the event</param>
         /// <param name="eventData">Data for the event</param>
         public async Task RaiseEventAsync(OrchestrationInstance orchestrationInstance, string eventName, object eventData)
+        {
+            await this.RaiseEventAsync(orchestrationInstance, eventName, eventData, emitTraceActivity: true);
+        }
+
+        /// <summary>
+        ///     Raises an event in the specified orchestration instance, which eventually causes the OnEvent() method in the
+        ///     orchestration to fire.
+        /// </summary>
+        /// <param name="orchestrationInstance">Instance in which to raise the event</param>
+        /// <param name="eventName">Name of the event</param>
+        /// <param name="eventData">Data for the event</param>
+        /// <param name="emitTraceActivity">Whether or not to emit a trace activity for this event.</param>
+        public async Task RaiseEventAsync(OrchestrationInstance orchestrationInstance, string eventName, object eventData, bool emitTraceActivity = true)
         {
 
             if (string.IsNullOrWhiteSpace(orchestrationInstance.InstanceId))
@@ -757,7 +770,11 @@ public async Task RaiseEventAsync(OrchestrationInstance orchestrationInstance, s
 
             // Distributed Tracing
             EventRaisedEvent eventRaisedEvent = new EventRaisedEvent(-1, serializedInput) { Name = eventName };
-            using Activity? traceActivity = TraceHelper.StartActivityForNewEventRaisedFromClient(eventRaisedEvent, orchestrationInstance);
+            Activity? traceActivity = null;
+            if (emitTraceActivity)
+            {
+                traceActivity = TraceHelper.StartActivityForNewEventRaisedFromClient(eventRaisedEvent, orchestrationInstance);
+            }
 
             var taskMessage = new TaskMessage
             {
@@ -776,6 +793,10 @@ public async Task RaiseEventAsync(OrchestrationInstance orchestrationInstance, s
                 TraceHelper.AddErrorDetailsToSpan(traceActivity, e);
                 throw;
             }
+            finally
+            {
+                traceActivity?.Dispose();
+            }
         }
 
         /// <summary>

src/DurableTask.Core/TaskOrchestrationDispatcher.cs

@@ -116,6 +116,11 @@ public async Task StopAsync(bool forced)
         /// </summary>
         public bool IncludeParameters { get; set; }
 
+        /// <summary>
+        /// Gets or sets the flag for whether or not entities are enabled
+        /// </summary>
+        public bool EntitiesEnabled { get; set; }
+
         /// <summary>
         /// Method to get the next work item to process within supplied timeout
         /// </summary>
@@ -1143,9 +1148,6 @@ TaskMessage ProcessCreateSubOrchestrationInstanceDecision(
                 historyEvent.Input = createSubOrchestrationAction.Input;
             }
 
-            ActivitySpanId clientSpanId = ActivitySpanId.CreateRandom();
-            historyEvent.ClientSpanId = clientSpanId.ToString();
-
             runtimeState.AddEvent(historyEvent);
 
             var taskMessage = new TaskMessage();
@@ -1169,8 +1171,22 @@ TaskMessage ProcessCreateSubOrchestrationInstanceDecision(
                 Version = createSubOrchestrationAction.Version
             };
 
-            if (parentTraceActivity != null)
+            // If a parent trace context was provided via the CreateSubOrchestrationAction.Tags, we will use this as the parent trace context of the suborchestration execution Activity rather than Activity.Current.Context.
+            if (createSubOrchestrationAction.Tags != null
+                && createSubOrchestrationAction.Tags.TryGetValue(OrchestrationTags.TraceParent, out string traceParent))
+            {
+                // If a parent trace context was provided but we fail to parse it, we don't want to attach any parent trace context to the start event since that will incorrectly link the trace corresponding to the orchestration execution
+                // as a child of Activity.Current, which is not truly the parent of the request
+                if (createSubOrchestrationAction.Tags.TryGetValue(OrchestrationTags.TraceState, out string traceState)
+                    && ActivityContext.TryParse(traceParent, traceState, out ActivityContext parentTraceContext))
+                {
+                    startedEvent.SetParentTraceContext(parentTraceContext);
+                }
+            }
+            else if (parentTraceActivity != null)
             {
+                ActivitySpanId clientSpanId = ActivitySpanId.CreateRandom();
+                historyEvent.ClientSpanId = clientSpanId.ToString();
                 ActivityContext activityContext = new ActivityContext(parentTraceActivity.TraceId, clientSpanId, parentTraceActivity.ActivityTraceFlags, parentTraceActivity.TraceStateString);
                 startedEvent.SetParentTraceContext(activityContext);
             }
@@ -1204,7 +1220,7 @@ TaskMessage ProcessSendEventDecision(
 
             // Distributed Tracing: start a new trace activity derived from the orchestration
             // for an EventRaisedEvent (external event)
-            using Activity? traceActivity = TraceHelper.StartTraceActivityForEventRaisedFromWorker(eventRaisedEvent, runtimeState.OrchestrationInstance, sendEventAction.Instance?.InstanceId);
+            using Activity? traceActivity = TraceHelper.StartTraceActivityForEventRaisedFromWorker(eventRaisedEvent, runtimeState.OrchestrationInstance, this.EntitiesEnabled, sendEventAction.Instance?.InstanceId);
 
             this.logHelper.RaisingEvent(runtimeState.OrchestrationInstance!, historyEvent);
 

src/DurableTask.Core/Tracing/TraceHelper.cs

@@ -284,7 +284,10 @@ internal static void EmitTraceActivityForTaskFailed(
                 return null;
             }
 
-            activity.SetSpanId(createdEvent.ClientSpanId);
+            if (!string.IsNullOrEmpty(createdEvent.ClientSpanId))
+            {
+                activity.SetSpanId(createdEvent.ClientSpanId);
+            }
 
             activity.SetTag(Schema.Task.Type, TraceActivityConstants.Orchestration);
             activity.SetTag(Schema.Task.Name, createdEvent.Name);
@@ -348,15 +351,23 @@ internal static void EmitTraceActivityForSubOrchestrationFailed(
         /// </summary>
         /// <param name="eventRaisedEvent">The associated <see cref="EventRaisedEvent"/>.</param>
         /// <param name="instance">The associated <see cref="OrchestrationInstance"/>.</param>
+        /// <param name="entitiesEnabled">Whether or not entities are enabled, meaning this event could possibly correspond to entity.</param>
         /// <param name="targetInstanceId">The instance id of the orchestration that will receive the event.</param>
         /// <returns>
         /// Returns a newly started <see cref="Activity"/> with (task) activity and orchestration-specific metadata.
         /// </returns>
         internal static Activity? StartTraceActivityForEventRaisedFromWorker(
             EventRaisedEvent eventRaisedEvent,
             OrchestrationInstance? instance,
+            bool entitiesEnabled,
             string? targetInstanceId)
         {
+            // We don't want to emit tracing for external events when they are related to entities
+            if (entitiesEnabled && (Entities.IsEntityInstance(targetInstanceId ?? string.Empty) || Entities.IsEntityInstance(instance?.InstanceId ?? string.Empty)))
+            {
+                return null;
+            }
+
             Activity? newActivity = ActivityTraceSource.StartActivity(
                 CreateSpanName(TraceActivityConstants.OrchestrationEvent, eventRaisedEvent.Name, null),
                 kind: ActivityKind.Producer,