Perf: avoid repeated orchestration-history scans for tracing

src/Shared/Grpc/Tracing/TraceHelper.cs

@@ -20,6 +20,17 @@ static class TraceHelper
 
     static readonly ActivitySource ActivityTraceSource = new ActivitySource(Source);
 
+    /// <summary>
+    /// Gets a value indicating whether any listener is currently registered for the Durable Task
+    /// <see cref="ActivitySource"/>.
+    /// </summary>
+    /// <remarks>
+    /// This is a cheap check that callers can use to skip trace-event lookup work (such as scanning
+    /// orchestration history to correlate scheduling events) when no listener is registered and any resulting
+    /// <see cref="Activity"/> would be discarded anyway.
+    /// </remarks>
+    public static bool HasListeners => ActivityTraceSource.HasListeners();
+
     /// <summary>
     /// Starts a new trace activity for scheduling an orchestration from the client.
     /// </summary>

src/Shared/Grpc/Tracing/TraceHistoryEventLookup.cs

@@ -0,0 +1,177 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using P = Microsoft.DurableTask.Protobuf;
+
+namespace Microsoft.DurableTask.Tracing;
+
+/// <summary>
+/// Provides indexed lookups of past history events by event ID, used to correlate new completion/failure
+/// events (e.g. "TaskCompleted") back to the history event that scheduled them (e.g. "TaskScheduled").
+/// </summary>
+/// <remarks>
+/// The indexes contain only scheduling event IDs referenced by the current work item's new completion/failure
+/// events. They are built together in one lazy scan and cached for the lifetime of the orchestrator work item.
+/// This avoids both re-scanning the full set of past events for every new event and retaining unrelated history
+/// events, bounding lookup storage by the number of correlation IDs in the current work item.
+/// </remarks>
+sealed class TraceHistoryEventLookup
+{
+    readonly IEnumerable<P.HistoryEvent> pastEvents;
+
+    readonly Dictionary<int, P.HistoryEvent?>? taskScheduledEventsByEventId;
+    readonly Dictionary<int, P.HistoryEvent?>? subOrchestrationInstanceCreatedEventsByEventId;
+
+    HashSet<int>? duplicateTaskScheduledEventIds;
+    HashSet<int>? duplicateSubOrchestrationInstanceCreatedEventIds;
+    bool indexesBuilt;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="TraceHistoryEventLookup"/> class.
+    /// </summary>
+    /// <param name="pastEvents">The past history events for the current orchestrator work item.</param>
+    /// <param name="newEvents">The new history events whose correlation IDs may be looked up.</param>
+    public TraceHistoryEventLookup(
+        IEnumerable<P.HistoryEvent> pastEvents,
+        IEnumerable<P.HistoryEvent> newEvents)
+    {
+        this.pastEvents = pastEvents;
+
+        foreach (P.HistoryEvent newEvent in newEvents)
+        {
+            switch (newEvent.EventTypeCase)
+            {
+                case P.HistoryEvent.EventTypeOneofCase.TaskCompleted:
+                    this.taskScheduledEventsByEventId ??= new();
+                    this.taskScheduledEventsByEventId[newEvent.TaskCompleted.TaskScheduledId] = null;
+                    break;
+
+                case P.HistoryEvent.EventTypeOneofCase.TaskFailed:
+                    this.taskScheduledEventsByEventId ??= new();
+                    this.taskScheduledEventsByEventId[newEvent.TaskFailed.TaskScheduledId] = null;
+                    break;
+
+                case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCompleted:
+                    this.subOrchestrationInstanceCreatedEventsByEventId ??= new();
+                    this.subOrchestrationInstanceCreatedEventsByEventId[
+                        newEvent.SubOrchestrationInstanceCompleted.TaskScheduledId] = null;
+                    break;
+
+                case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceFailed:
+                    this.subOrchestrationInstanceCreatedEventsByEventId ??= new();
+                    this.subOrchestrationInstanceCreatedEventsByEventId[
+                        newEvent.SubOrchestrationInstanceFailed.TaskScheduledId] = null;
+                    break;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Gets the "TaskScheduled" history event with the given event ID, if any.
+    /// </summary>
+    /// <param name="eventId">The event ID to look up.</param>
+    /// <returns>The matching event, or <see langword="null"/> if none is found.</returns>
+    /// <exception cref="InvalidOperationException">
+    /// Thrown when more than one "TaskScheduled" event has the same event ID.
+    /// </exception>
+    public P.HistoryEvent? GetTaskScheduledEvent(int eventId)
+    {
+        this.BuildIndexes();
+        ThrowIfDuplicate(
+            this.duplicateTaskScheduledEventIds,
+            eventId,
+            P.HistoryEvent.EventTypeOneofCase.TaskScheduled);
+
+        return this.taskScheduledEventsByEventId is not null
+            && this.taskScheduledEventsByEventId.TryGetValue(eventId, out P.HistoryEvent? historyEvent)
+            ? historyEvent
+            : null;
+    }
+
+    /// <summary>
+    /// Gets the "SubOrchestrationInstanceCreated" history event with the given event ID, if any.
+    /// </summary>
+    /// <param name="eventId">The event ID to look up.</param>
+    /// <returns>The matching event, or <see langword="null"/> if none is found.</returns>
+    /// <exception cref="InvalidOperationException">
+    /// Thrown when more than one "SubOrchestrationInstanceCreated" event has the same event ID.
+    /// </exception>
+    public P.HistoryEvent? GetSubOrchestrationInstanceCreatedEvent(int eventId)
+    {
+        this.BuildIndexes();
+        ThrowIfDuplicate(
+            this.duplicateSubOrchestrationInstanceCreatedEventIds,
+            eventId,
+            P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCreated);
+
+        return this.subOrchestrationInstanceCreatedEventsByEventId is not null
+            && this.subOrchestrationInstanceCreatedEventsByEventId.TryGetValue(
+                eventId, out P.HistoryEvent? historyEvent)
+            ? historyEvent
+            : null;
+    }
+
+    static void IndexRequestedEvent(
+        Dictionary<int, P.HistoryEvent?>? index,
+        ref HashSet<int>? duplicateEventIds,
+        P.HistoryEvent historyEvent)
+    {
+        if (index is null
+            || !index.TryGetValue(historyEvent.EventId, out P.HistoryEvent? existingEvent))
+        {
+            return;
+        }
+
+        if (existingEvent is null)
+        {
+            index[historyEvent.EventId] = historyEvent;
+        }
+        else
+        {
+            duplicateEventIds ??= new();
+            duplicateEventIds.Add(historyEvent.EventId);
+        }
+    }
+
+    static void ThrowIfDuplicate(
+        HashSet<int>? duplicateEventIds,
+        int eventId,
+        P.HistoryEvent.EventTypeOneofCase eventType)
+    {
+        if (duplicateEventIds?.Contains(eventId) == true)
+        {
+            throw new InvalidOperationException(
+                $"Past orchestration history contains multiple '{eventType}' events with event ID '{eventId}'.");
+        }
+    }
+
+    void BuildIndexes()
+    {
+        if (this.indexesBuilt)
+        {
+            return;
+        }
+
+        foreach (P.HistoryEvent historyEvent in this.pastEvents)
+        {
+            switch (historyEvent.EventTypeCase)
+            {
+                case P.HistoryEvent.EventTypeOneofCase.TaskScheduled:
+                    IndexRequestedEvent(
+                        this.taskScheduledEventsByEventId,
+                        ref this.duplicateTaskScheduledEventIds,
+                        historyEvent);
+                    break;
+
+                case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCreated:
+                    IndexRequestedEvent(
+                        this.subOrchestrationInstanceCreatedEventsByEventId,
+                        ref this.duplicateSubOrchestrationInstanceCreatedEventIds,
+                        historyEvent);
+                    break;
+            }
+        }
+
+        this.indexesBuilt = true;
+    }
+}

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -602,107 +602,110 @@ 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)
+            // Avoid the cost of scanning orchestration history for tracing purposes (potentially O(new events x
+            // past events) for work items with many new events) when no listener is registered for the Durable
+            // Task ActivitySource. In that case any Activity created below would be discarded anyway.
+            Activity? traceActivity = null;
+            if (TraceHelper.HasListeners)
             {
-                P.HistoryEvent? GetSuborchestrationInstanceCreatedEvent(int eventId)
-                {
-                    var subOrchestrationEvent =
-                        request
-                            .PastEvents
-                            .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCreated)
-                            .FirstOrDefault(x => x.EventId == eventId);
+                P.ExecutionStartedEvent? executionStartedEvent = FindExecutionStartedEvent(request);
 
-                    return subOrchestrationEvent;
-                }
+                traceActivity = TraceHelper.StartTraceActivityForOrchestrationExecution(
+                    executionStartedEvent,
+                    request.OrchestrationTraceContext);
 
-                P.HistoryEvent? GetTaskScheduledEvent(int eventId)
+                if (executionStartedEvent is not null)
                 {
-                    var taskScheduledEvent =
-                        request
-                            .PastEvents
-                            .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.TaskScheduled)
-                            .LastOrDefault(x => x.EventId == eventId);
+                    // Index only correlation IDs referenced by this work item's new events in one lazy history pass.
+                    TraceHistoryEventLookup historyLookup = new(request.PastEvents, request.NewEvents);
 
-                    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;
-                            }
+                        switch (newEvent.EventTypeCase)
+                        {
+                            case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceCompleted:
+                                {
+                                    P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
+                                        historyLookup.GetSubOrchestrationInstanceCreatedEvent(
+                                            newEvent.SubOrchestrationInstanceCompleted.TaskScheduledId);
+
+                                    TraceHelper.EmitTraceActivityForSubOrchestrationCompleted(
+                                        request.InstanceId,
+                                        subOrchestrationInstanceCreatedEvent,
+                                        subOrchestrationInstanceCreatedEvent?.SubOrchestrationInstanceCreated);
+                                    break;
+                                }
+
+                            case P.HistoryEvent.EventTypeOneofCase.SubOrchestrationInstanceFailed:
+                                {
+                                    P.HistoryEvent? subOrchestrationInstanceCreatedEvent =
+                                        historyLookup.GetSubOrchestrationInstanceCreatedEvent(
+                                            newEvent.SubOrchestrationInstanceFailed.TaskScheduledId);
+
+                                    TraceHelper.EmitTraceActivityForSubOrchestrationFailed(
+                                        request.InstanceId,
+                                        subOrchestrationInstanceCreatedEvent,
+                                        subOrchestrationInstanceCreatedEvent?.SubOrchestrationInstanceCreated,
+                                        newEvent.SubOrchestrationInstanceFailed);
+                                    break;
+                                }
+
+                            case P.HistoryEvent.EventTypeOneofCase.TaskCompleted:
+                                {
+                                    P.HistoryEvent? taskScheduledEvent =
+                                        historyLookup.GetTaskScheduledEvent(newEvent.TaskCompleted.TaskScheduledId);
 
-                        case P.HistoryEvent.EventTypeOneofCase.TaskCompleted:
-                            {
-                                P.HistoryEvent? taskScheduledEvent =
-                                    GetTaskScheduledEvent(newEvent.TaskCompleted.TaskScheduledId);
+                                    TraceHelper.EmitTraceActivityForTaskCompleted(
+                                        request.InstanceId,
+                                        taskScheduledEvent,
+                                        taskScheduledEvent?.TaskScheduled);
+                                    break;
+                                }
 
-                                TraceHelper.EmitTraceActivityForTaskCompleted(
+                            case P.HistoryEvent.EventTypeOneofCase.TaskFailed:
+                                {
+                                    P.HistoryEvent? taskScheduledEvent =
+                                        historyLookup.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);
+                                    executionStartedEvent.Name,
+                                    newEvent.Timestamp.ToDateTime(),
+                                    newEvent.TimerFired);
                                 break;
-                            }
+                        }
+                    }
+                }
+            }
 
-                        case P.HistoryEvent.EventTypeOneofCase.TaskFailed:
-                            {
-                                P.HistoryEvent? taskScheduledEvent =
-                                    GetTaskScheduledEvent(newEvent.TaskFailed.TaskScheduledId);
+            static P.ExecutionStartedEvent? FindExecutionStartedEvent(P.OrchestratorRequest request)
+            {
+                foreach (P.HistoryEvent newEvent in request.NewEvents)
+                {
+                    if (newEvent.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.ExecutionStarted)
+                    {
+                        return newEvent.ExecutionStarted;
+                    }
+                }
 
-                                TraceHelper.EmitTraceActivityForTaskFailed(
-                                    request.InstanceId,
-                                    taskScheduledEvent,
-                                    taskScheduledEvent?.TaskScheduled,
-                                    newEvent.TaskFailed);
-                                break;
-                            }
-
-                        case P.HistoryEvent.EventTypeOneofCase.TimerFired:
-                            TraceHelper.EmitTraceActivityForTimer(
-                                request.InstanceId,
-                                executionStartedEvent.Name,
-                                newEvent.Timestamp.ToDateTime(),
-                                newEvent.TimerFired);
-                            break;
+                foreach (P.HistoryEvent pastEvent in request.PastEvents)
+                {
+                    if (pastEvent.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.ExecutionStarted)
+                    {
+                        return pastEvent.ExecutionStarted;
                     }
                 }
+
+                return null;
             }
 
             OrchestratorExecutionResult? result = null;