[Exception Replay] Fixed capturing issue of async methods with await in finally block + added missing snapshot attributes + better frame matching algorithm

tracer/src/Datadog.Trace/Debugger/ExceptionAutoInstrumentation/ExceptionCaseInstrumentationManager.cs

@@ -9,6 +9,7 @@
 using System.Diagnostics;
 using System.Linq;
 using System.Reflection;
+using System.Reflection.Emit;
 using System.Threading;
 using System.Threading.Tasks;
 using Datadog.Trace.Debugger.Configurations.Models;
@@ -33,11 +34,13 @@ internal class ExceptionCaseInstrumentationManager
         private static readonly ConcurrentDictionary<MethodUniqueIdentifier, ExceptionDebuggingProbe> MethodToProbe = new();
         private static readonly int MaxFramesToCapture = ExceptionDebugging.Settings.MaximumFramesToCapture;
 
-        internal static ExceptionCase Instrument(ExceptionIdentifier exceptionId)
+        internal static ExceptionCase Instrument(ExceptionIdentifier exceptionId, string exceptionToString)
         {
             Log.Information("Instrumenting {ExceptionId}", exceptionId);
 
-            var participatingUserMethods = GetMethodsToRejit(exceptionId.StackTrace);
+            var parsedFramesFromExceptionToString = StackTraceProcessor.ParseFrames(exceptionToString).ToArray();
+            var stackTrace = exceptionId.StackTrace.Where(frame => parsedFramesFromExceptionToString.Any(f => MethodMatcher.IsMethodMatch(f, frame.Method))).ToArray();
+            var participatingUserMethods = GetMethodsToRejit(stackTrace);
 
             var uniqueMethods = participatingUserMethods
                                .Distinct(EqualityComparer<MethodUniqueIdentifier>.Default)
@@ -65,7 +68,7 @@ internal static ExceptionCase Instrument(ExceptionIdentifier exceptionId)
                 }
             }
 
-            var newCase = new ExceptionCase(exceptionId, probes);
+            var newCase = new ExceptionCase(exceptionId.ExceptionTypes, probes);
 
             foreach (var method in uniqueMethods)
             {
@@ -84,6 +87,8 @@ bool ShouldInstrumentFrameAtIndex(int i)
         private static List<MethodUniqueIdentifier> GetMethodsToRejit(ParticipatingFrame[] allFrames)
         {
             var methodsToRejit = new List<MethodUniqueIdentifier>();
+            MethodUniqueIdentifier? lastMethod = null;
+            var wasLastMisleading = false;
 
             foreach (var frame in allFrames)
             {
@@ -102,11 +107,24 @@ private static List<MethodUniqueIdentifier> GetMethodsToRejit(ParticipatingFrame
                         continue;
                     }
 
-                    methodsToRejit.Add(frame.MethodIdentifier);
+                    var currentMethod = frame.MethodIdentifier;
+                    var isCurrentMisleading = currentMethod.IsMisleadMethod();
+
+                    // Add the method if either:
+                    // 1. It's not misleading (we keep all non-misleading methods)
+                    // 2. It's misleading but different from the last misleading method we saw
+                    // 3. It's the first misleading method after non-misleading methods
+                    if (!isCurrentMisleading || currentMethod != lastMethod || !wasLastMisleading)
+                    {
+                        methodsToRejit.Add(currentMethod);
+                    }
+
+                    lastMethod = currentMethod;
+                    wasLastMisleading = isCurrentMisleading;
                 }
                 catch (Exception ex)
                 {
-                    Log.Error(ex, "Failed to instrument frame the frame: {FrameToRejit}", frame);
+                    Log.Error(ex, "Failed to instrument the frame: {FrameToRejit}", frame);
                 }
             }
 

tracer/src/Datadog.Trace/Debugger/ExceptionAutoInstrumentation/ExceptionDebuggingProbe.cs

@@ -6,6 +6,10 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Reflection;
+using System.Reflection.Emit;
+using System.Runtime.CompilerServices;
+using System.Runtime.ExceptionServices;
 using System.Threading;
 using Datadog.Trace.Debugger.Expressions;
 using Datadog.Trace.Debugger.Helpers;
@@ -93,7 +97,7 @@ private void ProcessCase(ExceptionCase @case)
                     var parentProbes = probes.Take(index).ToArray();
                     var childProbes = probes.Skip(index + 1).ToArray();
 
-                    var processor = new ExceptionProbeProcessor(probe, @case.ExceptionId.ExceptionTypes, parentProbes: parentProbes, childProbes: childProbes);
+                    var processor = new ExceptionProbeProcessor(probe, @case.ExceptionTypes, parentProbes: parentProbes, childProbes: childProbes);
                     @case.Processors.TryAdd(processor, 0);
                     ExceptionDebuggingProcessor?.AddProbeProcessor(processor);
                 }

tracer/src/Datadog.Trace/Debugger/ExceptionAutoInstrumentation/ExceptionDebuggingProcessor.cs

@@ -19,6 +19,7 @@ internal class ExceptionDebuggingProcessor : IProbeProcessor
         private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(ExceptionDebuggingProcessor));
         private readonly object _lock = new();
         private readonly int _maxFramesToCapture;
+        private readonly bool _isMisleadingMethod;
         private ExceptionProbeProcessor[] _processors;
 
         internal ExceptionDebuggingProcessor(string probeId, MethodUniqueIdentifier method)
@@ -27,6 +28,7 @@ internal ExceptionDebuggingProcessor(string probeId, MethodUniqueIdentifier meth
             ProbeId = probeId;
             Method = method;
             _maxFramesToCapture = ExceptionDebugging.Settings.MaximumFramesToCapture;
+            _isMisleadingMethod = method.IsMisleadMethod();
         }
 
         public string ProbeId { get; }
@@ -68,7 +70,17 @@ public bool Process<TCapture>(ref CaptureInfo<TCapture> info, IDebuggerSnapshotC
                     case MethodState.EntryStart:
                     case MethodState.EntryAsync:
                         shadowStack = ShadowStackHolder.EnsureShadowStackEnabled();
-                        snapshotCreator.EnterHash = shadowStack.CurrentStackFrameNode?.EnterSequenceHash ?? Fnv1aHash.FnvOffsetBias;
+                        var currentFrame = shadowStack.CurrentStackFrameNode;
+                        snapshotCreator.EnterHash = Fnv1aHash.Combine(info.Method.MetadataToken, shadowStack.CurrentStackFrameNode?.EnterSequenceHash ?? Fnv1aHash.FnvOffsetBias);
+
+                        if (currentFrame?.Method == info.Method && _isMisleadingMethod)
+                        {
+                            // Methods marked as `misleading` are methods we tolerate being in the shadow stack multiple times.
+                            // We flatten those methods due to `ExceptionDispatchInfo.Capture(X).Throw()` API causing frames to appear twice
+                            // while in reality they were involved only once.
+                            snapshotCreator.TrackedStackFrameNode = shadowStack.EnterFake(info.Method);
+                            return true;
+                        }
 
                         var shouldProcess = false;
                         foreach (var processor in snapshotCreator.Processors)
@@ -113,7 +125,15 @@ public bool Process<TCapture>(ref CaptureInfo<TCapture> info, IDebuggerSnapshotC
 
                         var exception = info.Value as Exception;
                         snapshotCreator.TrackedStackFrameNode.LeavingException = exception;
-                        snapshotCreator.LeaveHash = shadowStack.CurrentStackFrameNode?.LeaveSequenceHash ?? Fnv1aHash.FnvOffsetBias;
+                        snapshotCreator.LeaveHash = shadowStack.CurrentStackFrameNode!.LeaveSequenceHash;
+
+                        if (snapshotCreator.TrackedStackFrameNode is FakeTrackedStackFrameNode)
+                        {
+                            shadowStack.Leave(snapshotCreator.TrackedStackFrameNode, exception);
+                            snapshotCreator.TrackedStackFrameNode.CapturingStrategy = SnapshotCapturingStrategy.FullSnapshot;
+                            snapshotCreator.TrackedStackFrameNode.AddScopeMember(info.Name, info.Type, info.Value, info.MemberKind);
+                            return true;
+                        }
 
                         var leavingExceptionType = info.Value.GetType();
 

tracer/src/Datadog.Trace/Debugger/ExceptionAutoInstrumentation/ExceptionNormalizer.cs

@@ -4,6 +4,7 @@
 // </copyright>
 
 using System;
+using System.Collections.Generic;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using System.Text;
@@ -16,6 +17,12 @@ namespace Datadog.Trace.Debugger.ExceptionAutoInstrumentation
 {
     internal class ExceptionNormalizer
     {
+        protected ExceptionNormalizer()
+        {
+        }
+
+        public static ExceptionNormalizer Instance { get; } = new();
+
         /// <summary>
         /// Given the string representation of an exception alongside it's FQN of the outer and (potential) inner exception,
         /// this function cleanse the stack trace from error messages, customized information attached to the exception and PDB line info if present.
@@ -95,5 +102,85 @@ protected virtual int HashLine(VendoredMicrosoftCode.System.ReadOnlySpan<char> l
 
             return fnvHashCode;
         }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal List<string> ParseFrames(string exceptionString)
+        {
+            if (string.IsNullOrEmpty(exceptionString))
+            {
+                throw new ArgumentException(@"Exception string cannot be null or empty", nameof(exceptionString));
+            }
+
+            var results = new List<string>();
+            var currentSpan = VendoredMicrosoftCode.System.MemoryExtensions.AsSpan(exceptionString);
+
+            while (!currentSpan.IsEmpty)
+            {
+                var lineEndIndex = currentSpan.IndexOfAny('\r', '\n');
+                VendoredMicrosoftCode.System.ReadOnlySpan<char> line;
+
+                if (lineEndIndex >= 0)
+                {
+                    line = currentSpan.Slice(0, lineEndIndex);
+                    currentSpan = currentSpan.Slice(lineEndIndex + 1);
+                    if (!currentSpan.IsEmpty && currentSpan[0] == '\n')
+                    {
+                        currentSpan = currentSpan.Slice(1);
+                    }
+                }
+                else
+                {
+                    line = currentSpan;
+                    currentSpan = default;
+                }
+
+                ProcessLine(line, results);
+            }
+
+            return results;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static void ProcessLine(VendoredMicrosoftCode.System.ReadOnlySpan<char> line, List<string> results)
+        {
+            line = line.TrimStart();
+            if (line.IsEmpty)
+            {
+                return;
+            }
+
+            // Check if it's a stack frame line (starts with "at ")
+            if (!VendoredMicrosoftCode.System.MemoryExtensions.StartsWith(line, VendoredMicrosoftCode.System.MemoryExtensions.AsSpan("at "), StringComparison.Ordinal))
+            {
+                return;
+            }
+
+            // Skip the "at " prefix
+            line = line.Slice(3);
+
+            // Skip lambda and Datadog frames early
+            if (ContainsAny(line, "lambda_", "at Datadog."))
+            {
+                return;
+            }
+
+            // Find the " in " marker and truncate if found
+            var inIndex = VendoredMicrosoftCode.System.MemoryExtensions.IndexOf(line, VendoredMicrosoftCode.System.MemoryExtensions.AsSpan(" in "), StringComparison.Ordinal);
+
+            if (inIndex > 0)
+            {
+                line = line.Slice(0, inIndex);
+            }
+
+            // Only create a string when we're sure we want to keep this frame
+            results.Add(line.ToString());
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static bool ContainsAny(VendoredMicrosoftCode.System.ReadOnlySpan<char> source, string first, string second)
+        {
+            return VendoredMicrosoftCode.System.MemoryExtensions.Contains(source, VendoredMicrosoftCode.System.MemoryExtensions.AsSpan(first), StringComparison.Ordinal) ||
+                   VendoredMicrosoftCode.System.MemoryExtensions.Contains(source, VendoredMicrosoftCode.System.MemoryExtensions.AsSpan(second), StringComparison.Ordinal);
+        }
     }
 }

tracer/src/Datadog.Trace/Debugger/ExceptionAutoInstrumentation/ExceptionProbeProcessor.cs

@@ -116,19 +116,17 @@ private bool EnsureLeaveHashComputed()
                 Array.Reverse(installedProbes);
             }
 
-            if (!installedProbes.Any())
-            {
-                return Fnv1aHash.FnvOffsetBias;
-            }
-
             var hash = Fnv1aHash.FnvOffsetBias;
 
-            foreach (var probe in installedProbes)
+            if (installedProbes.Any())
             {
-                hash = Fnv1aHash.Combine(probe.Method.MethodToken, hash);
+                foreach (var probe in installedProbes)
+                {
+                    hash = Fnv1aHash.Combine(probe.Method.Method.MetadataToken, hash);
+                }
             }
 
-            return hash;
+            return Fnv1aHash.Combine(ExceptionDebuggingProcessor.Method.Method.MetadataToken, hash);
         }
 
         internal void InvalidateEnterLeave()

tracer/src/Datadog.Trace/Debugger/ExceptionAutoInstrumentation/ExceptionReplaySnapshotCreator.cs

@@ -0,0 +1,50 @@
+// <copyright file="ExceptionReplaySnapshotCreator.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Datadog.Trace.Debugger.Expressions;
+using Datadog.Trace.Debugger.Snapshots;
+using ProbeLocation = Datadog.Trace.Debugger.Expressions.ProbeLocation;
+
+#nullable enable
+namespace Datadog.Trace.Debugger.ExceptionAutoInstrumentation
+{
+    internal class ExceptionReplaySnapshotCreator : DebuggerSnapshotCreator
+    {
+        public ExceptionReplaySnapshotCreator(bool isFullSnapshot, ProbeLocation location, bool hasCondition, string[] tags, CaptureLimitInfo limitInfo)
+            : base(isFullSnapshot, location, hasCondition, tags, limitInfo)
+        {
+        }
+
+        public ExceptionReplaySnapshotCreator(bool isFullSnapshot, ProbeLocation location, bool hasCondition, string[] tags, MethodScopeMembers methodScopeMembers, CaptureLimitInfo limitInfo)
+            : base(isFullSnapshot, location, hasCondition, tags, methodScopeMembers, limitInfo)
+        {
+        }
+
+        internal static string ExceptionHash { get; } = Guid.NewGuid().ToString();
+
+        internal static string ExceptionCaptureId { get; } = Guid.NewGuid().ToString();
+
+        internal static string FrameIndex { get; } = Guid.NewGuid().ToString();
+
+        internal override DebuggerSnapshotCreator EndSnapshot()
+        {
+            JsonWriter.WritePropertyName("exceptionHash");
+            JsonWriter.WriteValue(ExceptionHash);
+
+            JsonWriter.WritePropertyName("exceptionId");
+            JsonWriter.WriteValue(ExceptionCaptureId);
+
+            JsonWriter.WritePropertyName("frameIndex");
+            JsonWriter.WriteValue(FrameIndex);
+
+            return base.EndSnapshot();
+        }
+    }
+}