Optimize orchestration history scans for tracing performance

src/Shared/Grpc/Tracing/TraceHelper.cs

@@ -20,6 +20,12 @@ static class TraceHelper
 
     static readonly ActivitySource ActivityTraceSource = new ActivitySource(Source);
 
+    /// <summary>
+    /// Gets whether any listener is subscribed to Durable Task tracing activities.
+    /// </summary>
+    /// <returns><see langword="true"/> when the activity source has at least one listener; otherwise, <see langword="false"/>.</returns>
+    public static bool HasListeners() => ActivityTraceSource.HasListeners();
+
     /// <summary>
     /// Starts a new trace activity for scheduling an orchestration from the client.
     /// </summary>

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -659,13 +659,19 @@ async Task OnRunOrchestratorAsync(
             }
 
             IReadOnlyList<P.HistoryEvent> pastEvents = materializedPastEvents ?? request.PastEvents;
-            var executionStartedEvent =
-                request
-                    .NewEvents
-                    .Concat(pastEvents)
-                    .Where(e => e.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.ExecutionStarted)
-                    .Select(e => e.ExecutionStarted)
-                    .FirstOrDefault();
+            bool hasTraceListeners = TraceHelper.HasListeners();
+
+            P.ExecutionStartedEvent? executionStartedEvent = null;
+            if (hasTraceListeners || isInitialRewind)
+            {
+                executionStartedEvent =
+                    request
+                        .NewEvents
+                        .Concat(pastEvents)
+                        .Where(e => e.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.ExecutionStarted)
+                        .Select(e => e.ExecutionStarted)
+                        .FirstOrDefault();
+            }
 
             if (isInitialRewind)
             {
@@ -684,9 +690,11 @@ async Task OnRunOrchestratorAsync(
             // A rewind starts a new orchestration span instead of continuing the failed execution's stored span.
             P.OrchestrationTraceContext? orchestrationTraceContext =
                 isInitialRewind ? null : request.OrchestrationTraceContext;
-            Activity? traceActivity = TraceHelper.StartTraceActivityForOrchestrationExecution(
-                executionStartedEvent,
-                orchestrationTraceContext);
+            Activity? traceActivity = hasTraceListeners
+                ? TraceHelper.StartTraceActivityForOrchestrationExecution(
+                    executionStartedEvent,
+                    orchestrationTraceContext)
+                : null;
 
             if (isInitialRewind)
             {
@@ -701,27 +709,9 @@ await this.CompleteOrchestratorTaskWithChunkingAsync(
                 return;
             }
 
-            if (executionStartedEvent is not null)
+            if (hasTraceListeners && executionStartedEvent is not null)
             {
-                P.HistoryEvent? GetSuborchestrationInstanceCreatedEvent(int eventId)
-                {
-                    var subOrchestrationEvent =
-                        pastEvents
-                            .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCreated)
-                            .FirstOrDefault(x => x.EventId == eventId);
-
-                    return subOrchestrationEvent;
-                }
-
-                P.HistoryEvent? GetTaskScheduledEvent(int eventId)
-                {
-                    var taskScheduledEvent =
-                        pastEvents
-                            .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.TaskScheduled)
-                            .LastOrDefault(x => x.EventId == eventId);
-
-                    return taskScheduledEvent;
-                }
+                TracingHistoryEventIndex historyEventIndex = new(pastEvents);
 
                 foreach (var newEvent in request.NewEvents)
                 {
@@ -730,7 +720,7 @@ await this.CompleteOrchestratorTaskWithChunkingAsync(
                         case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCompleted:
                             {
                                 P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
-                                    GetSuborchestrationInstanceCreatedEvent(
+                                    historyEventIndex.GetSubOrchestrationInstanceCreatedEvent(
                                         newEvent.SubOrchestrationInstanceCompleted.TaskScheduledId);
 
                                 TraceHelper.EmitTraceActivityForSubOrchestrationCompleted(
@@ -743,7 +733,7 @@ await this.CompleteOrchestratorTaskWithChunkingAsync(
                         case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceFailed:
                             {
                                 P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
-                                    GetSuborchestrationInstanceCreatedEvent(
+                                    historyEventIndex.GetSubOrchestrationInstanceCreatedEvent(
                                         newEvent.SubOrchestrationInstanceFailed.TaskScheduledId);
 
                                 TraceHelper.EmitTraceActivityForSubOrchestrationFailed(
@@ -757,7 +747,7 @@ await this.CompleteOrchestratorTaskWithChunkingAsync(
                         case P.HistoryEvent.EventTypeOneofCase.TaskCompleted:
                             {
                                 P.HistoryEvent? taskScheduledEvent =
-                                    GetTaskScheduledEvent(newEvent.TaskCompleted.TaskScheduledId);
+                                    historyEventIndex.GetTaskScheduledEvent(newEvent.TaskCompleted.TaskScheduledId);
 
                                 TraceHelper.EmitTraceActivityForTaskCompleted(
                                     request.InstanceId,
@@ -769,7 +759,7 @@ await this.CompleteOrchestratorTaskWithChunkingAsync(
                         case P.HistoryEvent.EventTypeOneofCase.TaskFailed:
                             {
                                 P.HistoryEvent? taskScheduledEvent =
-                                    GetTaskScheduledEvent(newEvent.TaskFailed.TaskScheduledId);
+                                    historyEventIndex.GetTaskScheduledEvent(newEvent.TaskFailed.TaskScheduledId);
 
                                 TraceHelper.EmitTraceActivityForTaskFailed(
                                     request.InstanceId,

src/Worker/Grpc/TracingHistoryEventIndex.cs

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using P = Microsoft.DurableTask.Protobuf;
+
+namespace Microsoft.DurableTask.Worker.Grpc;
+
+/// <summary>
+/// Indexes the orchestration history events used to reconstruct tracing spans.
+/// </summary>
+sealed class TracingHistoryEventIndex
+{
+    readonly Dictionary<int, P.HistoryEvent> subOrchestrationCreatedEvents = new();
+    readonly Dictionary<int, P.HistoryEvent> taskScheduledEvents = new();
+
+    public TracingHistoryEventIndex(IEnumerable<P.HistoryEvent> pastEvents)
+    {
+        foreach (P.HistoryEvent historyEvent in pastEvents)
+        {
+            switch (historyEvent.EventTypeCase)
+            {
+                case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCreated:
+                    // Preserve the previous FirstOrDefault semantics for duplicate IDs.
+                    this.subOrchestrationCreatedEvents.TryAdd(historyEvent.EventId, historyEvent);
+                    break;
+
+                case P.HistoryEvent.EventTypeOneofCase.TaskScheduled:
+                    // Preserve the previous LastOrDefault semantics for duplicate IDs.
+                    this.taskScheduledEvents[historyEvent.EventId] = historyEvent;
+                    break;
+            }
+        }
+    }
+
+    public P.HistoryEvent? GetSubOrchestrationInstanceCreatedEvent(int eventId)
+        => this.subOrchestrationCreatedEvents.TryGetValue(eventId, out P.HistoryEvent? historyEvent)
+            ? historyEvent
+            : null;
+
+    public P.HistoryEvent? GetTaskScheduledEvent(int eventId)
+        => this.taskScheduledEvents.TryGetValue(eventId, out P.HistoryEvent? historyEvent)
+            ? historyEvent
+            : null;
+}

test/Worker/Grpc.Tests/TraceHelperTests.cs

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Diagnostics;
+using Microsoft.DurableTask.Tracing;
+
+namespace Microsoft.DurableTask.Worker.Grpc.Tests;
+
+public class TraceHelperTests
+{
+    [Fact]
+    public void HasListeners_TracksMatchingActivityListener()
+    {
+        bool initialHasListeners = TraceHelper.HasListeners();
+        ActivityListener listener = new()
+        {
+            ShouldListenTo = source => source.Name == "Microsoft.DurableTask",
+            Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.None,
+        };
+
+        try
+        {
+            ActivitySource.AddActivityListener(listener);
+            TraceHelper.HasListeners().Should().BeTrue();
+        }
+        finally
+        {
+            listener.Dispose();
+        }
+
+        TraceHelper.HasListeners().Should().Be(initialHasListeners);
+    }
+}

test/Worker/Grpc.Tests/TracingHistoryEventIndexTests.cs

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using P = Microsoft.DurableTask.Protobuf;
+
+namespace Microsoft.DurableTask.Worker.Grpc.Tests;
+
+public class TracingHistoryEventIndexTests
+{
+    [Fact]
+    public void GetSubOrchestrationInstanceCreatedEvent_DuplicateIds_ReturnsFirstEvent()
+    {
+        P.HistoryEvent first = new()
+        {
+            EventId = 7,
+            SubOrchestrationInstanceCreated = new P.SubOrchestrationInstanceCreatedEvent { Name = "first" },
+        };
+        P.HistoryEvent second = new()
+        {
+            EventId = 7,
+            SubOrchestrationInstanceCreated = new P.SubOrchestrationInstanceCreatedEvent { Name = "second" },
+        };
+
+        TracingHistoryEventIndex index = new([first, second]);
+        index.GetSubOrchestrationInstanceCreatedEvent(7).Should().BeSameAs(first);
+    }
+
+    [Fact]
+    public void GetTaskScheduledEvent_DuplicateIds_ReturnsLastEvent()
+    {
+        P.HistoryEvent first = new()
+        {
+            EventId = 11,
+            TaskScheduled = new P.TaskScheduledEvent { Name = "first" },
+        };
+        P.HistoryEvent second = new()
+        {
+            EventId = 11,
+            TaskScheduled = new P.TaskScheduledEvent { Name = "second" },
+        };
+
+        TracingHistoryEventIndex index = new([first, second]);
+        index.GetTaskScheduledEvent(11).Should().BeSameAs(second);
+    }
+
+    [Fact]
+    public void Lookups_MissingIds_ReturnNull()
+    {
+        P.HistoryEvent unrelated = new()
+        {
+            EventId = 3,
+            TimerCreated = new P.TimerCreatedEvent(),
+        };
+
+        TracingHistoryEventIndex index = new([unrelated]);
+        index.GetSubOrchestrationInstanceCreatedEvent(3).Should().BeNull();
+        index.GetTaskScheduledEvent(3).Should().BeNull();
+    }
+}