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 tracing work can produce activities; 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

@@ -602,105 +602,89 @@ async Task OnRunOrchestratorAsync(
             string completionToken,
             CancellationToken cancellationToken)
         {
-            var executionStartedEvent =
-                request
-                    .NewEvents
-                    .Concat(request.PastEvents)
-                    .Where(e => e.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.ExecutionStarted)
-                    .Select(e => e.ExecutionStarted)
-                    .FirstOrDefault();
-
-            Activity? traceActivity = TraceHelper.StartTraceActivityForOrchestrationExecution(
-                executionStartedEvent,
-                request.OrchestrationTraceContext);
-
-            if (executionStartedEvent is not null)
+            Activity? traceActivity = null;
+            if (TraceHelper.HasListeners())
             {
-                P.HistoryEvent? GetSuborchestrationInstanceCreatedEvent(int eventId)
+                var executionStartedEvent =
+                    request
+                        .NewEvents
+                        .Concat(request.PastEvents)
+                        .Where(e => e.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.ExecutionStarted)
+                        .Select(e => e.ExecutionStarted)
+                        .FirstOrDefault();
+
+                traceActivity = TraceHelper.StartTraceActivityForOrchestrationExecution(
+                    executionStartedEvent,
+                    request.OrchestrationTraceContext);
+
+                if (executionStartedEvent is not null)
                 {
-                    var subOrchestrationEvent =
-                        request
-                            .PastEvents
-                            .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCreated)
-                            .FirstOrDefault(x => x.EventId == eventId);
+                    TracingHistoryEventIndex historyEventIndex = new(request.PastEvents);
 
-                    return subOrchestrationEvent;
-                }
-
-                P.HistoryEvent? GetTaskScheduledEvent(int eventId)
-                {
-                    var taskScheduledEvent =
-                        request
-                            .PastEvents
-                            .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.TaskScheduled)
-                            .LastOrDefault(x => x.EventId == eventId);
-
-                    return taskScheduledEvent;
-                }
-
-                foreach (var newEvent in request.NewEvents)
-                {
-                    switch (newEvent.EventTypeCase)
+                    foreach (var newEvent in request.NewEvents)
                     {
-                        case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCompleted:
-                            {
-                                P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
-                                    GetSuborchestrationInstanceCreatedEvent(
-                                        newEvent.SubOrchestrationInstanceCompleted.TaskScheduledId);
-
-                                TraceHelper.EmitTraceActivityForSubOrchestrationCompleted(
-                                    request.InstanceId,
-                                    subOrchestrationInstanceCreatedEvent,
-                                    subOrchestrationInstanceCreatedEvent?.SubOrchestrationInstanceCreated);
-                                break;
-                            }
-
-                        case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceFailed:
-                            {
-                                P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
-                                    GetSuborchestrationInstanceCreatedEvent(
-                                        newEvent.SubOrchestrationInstanceFailed.TaskScheduledId);
-
-                                TraceHelper.EmitTraceActivityForSubOrchestrationFailed(
-                                    request.InstanceId,
-                                    subOrchestrationInstanceCreatedEvent,
-                                    subOrchestrationInstanceCreatedEvent?.SubOrchestrationInstanceCreated,
-                                    newEvent.SubOrchestrationInstanceFailed);
-                                break;
-                            }
-
-                        case P.HistoryEvent.EventTypeOneofCase.TaskCompleted:
-                            {
-                                P.HistoryEvent? taskScheduledEvent =
-                                    GetTaskScheduledEvent(newEvent.TaskCompleted.TaskScheduledId);
-
-                                TraceHelper.EmitTraceActivityForTaskCompleted(
-                                    request.InstanceId,
-                                    taskScheduledEvent,
-                                    taskScheduledEvent?.TaskScheduled);
-                                break;
-                            }
+                        switch (newEvent.EventTypeCase)
+                        {
+                            case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCompleted:
+                                {
+                                    P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
+                                        historyEventIndex.GetSubOrchestrationInstanceCreatedEvent(
+                                            newEvent.SubOrchestrationInstanceCompleted.TaskScheduledId);
+
+                                    TraceHelper.EmitTraceActivityForSubOrchestrationCompleted(
+                                        request.InstanceId,
+                                        subOrchestrationInstanceCreatedEvent,
+                                        subOrchestrationInstanceCreatedEvent?.SubOrchestrationInstanceCreated);
+                                    break;
+                                }
+
+                            case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceFailed:
+                                {
+                                    P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
+                                        historyEventIndex.GetSubOrchestrationInstanceCreatedEvent(
+                                            newEvent.SubOrchestrationInstanceFailed.TaskScheduledId);
+
+                                    TraceHelper.EmitTraceActivityForSubOrchestrationFailed(
+                                        request.InstanceId,
+                                        subOrchestrationInstanceCreatedEvent,
+                                        subOrchestrationInstanceCreatedEvent?.SubOrchestrationInstanceCreated,
+                                        newEvent.SubOrchestrationInstanceFailed);
+                                    break;
+                                }
+
+                            case P.HistoryEvent.EventTypeOneofCase.TaskCompleted:
+                                {
+                                    P.HistoryEvent? taskScheduledEvent =
+                                        historyEventIndex.GetTaskScheduledEvent(newEvent.TaskCompleted.TaskScheduledId);
 
-                        case P.HistoryEvent.EventTypeOneofCase.TaskFailed:
-                            {
-                                P.HistoryEvent? taskScheduledEvent =
-                                    GetTaskScheduledEvent(newEvent.TaskFailed.TaskScheduledId);
+                                    TraceHelper.EmitTraceActivityForTaskCompleted(
+                                        request.InstanceId,
+                                        taskScheduledEvent,
+                                        taskScheduledEvent?.TaskScheduled);
+                                    break;
+                                }
 
-                                TraceHelper.EmitTraceActivityForTaskFailed(
+                            case P.HistoryEvent.EventTypeOneofCase.TaskFailed:
+                                {
+                                    P.HistoryEvent? taskScheduledEvent =
+                                        historyEventIndex.GetTaskScheduledEvent(newEvent.TaskFailed.TaskScheduledId);
+
+                                    TraceHelper.EmitTraceActivityForTaskFailed(
+                                        request.InstanceId,
+                                        taskScheduledEvent,
+                                        taskScheduledEvent?.TaskScheduled,
+                                        newEvent.TaskFailed);
+                                    break;
+                                }
+
+                            case P.HistoryEvent.EventTypeOneofCase.TimerFired:
+                                TraceHelper.EmitTraceActivityForTimer(
                                     request.InstanceId,
-                                    taskScheduledEvent,
-                                    taskScheduledEvent?.TaskScheduled,
-                                    newEvent.TaskFailed);
+                                    executionStartedEvent.Name,
+                                    newEvent.Timestamp.ToDateTime(),
+                                    newEvent.TimerFired);
                                 break;
-                            }
-
-                        case P.HistoryEvent.EventTypeOneofCase.TimerFired:
-                            TraceHelper.EmitTraceActivityForTimer(
-                                request.InstanceId,
-                                executionStartedEvent.Name,
-                                newEvent.Timestamp.ToDateTime(),
-                                newEvent.TimerFired);
-                            break;
+                        }
                     }
                 }
             }

src/Worker/Grpc/TracingHistoryEventIndex.cs

@@ -0,0 +1,48 @@
+// 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.
+                    if (!this.subOrchestrationCreatedEvents.ContainsKey(historyEvent.EventId))
+                    {
+                        this.subOrchestrationCreatedEvents.Add(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",
+        };
+
+        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,62 @@
+// 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();
+    }
+}