[Dynamic Instrumentation] Perf improvements + Added probe metadata payload (#3725)

* Deleted all the ConcurrentDictionary that played part in the hot path of our DI's instrumentation

* Added a test with probe that is getting changed between different phases + fixed an issue we had with duplicated probes in the native side

* Cleaned singleton enforcement of the registries

* Fixed a bug we had where we changed the TaskLocal of the continuation and did not restore it upon rate limiting hit

* Fixed misinstrumetation of async methods with stack-allocated return values

* attached span_id & trace_id to debugger snapshots

* attaching tags from ProbeDefintion into Debugger Snapshots

* Addressed PR comments

Author: GreenMatan

tracer/src/Datadog.Trace/Debugger/Expressions/ProbeExpressionsProcessor.cs

@@ -35,23 +35,6 @@ internal static ProbeExpressionsProcessor Instance
             }
         }
 
-        internal ProbeInfo? GetProbeInfo(string probeId)
-        {
-            return Get(probeId)?.ProbeInfo;
-        }
-
-        internal bool Process<T>(string probeId, ref CaptureInfo<T> info, DebuggerSnapshotCreator snapshotCreator)
-        {
-            var probeProcessor = Get(probeId);
-            if (probeProcessor == null)
-            {
-                Log.Error("Probe processor has not found. Probe Id: " + probeId);
-                return false;
-            }
-
-            return probeProcessor.Process(ref info, snapshotCreator);
-        }
-
         internal void AddProbeProcessor(ProbeDefinition probe)
         {
             try
@@ -72,7 +55,7 @@ internal void Remove(string probeId)
             _processors.TryRemove(probeId, out _);
         }
 
-        private ProbeProcessor Get(string probeId)
+        internal ProbeProcessor Get(string probeId)
         {
             _processors.TryGetValue(probeId, out var probeProcessor);
             return probeProcessor;

tracer/src/Datadog.Trace/Debugger/Expressions/ProbeInfo.cs

@@ -12,7 +12,8 @@ internal readonly record struct ProbeInfo(
         ProbeType ProbeType,
         ProbeLocation ProbeLocation,
         EvaluateAt EvaluateAt,
-        bool HasCondition)
+        bool HasCondition,
+        string[] Tags)
     {
         internal string ProbeId { get; } = ProbeId;
 
@@ -25,5 +26,7 @@ internal readonly record struct ProbeInfo(
         internal bool IsFullSnapshot { get; } = ProbeType == ProbeType.Snapshot;
 
         internal bool HasCondition { get; } = HasCondition;
+
+        internal string[] Tags { get; } = Tags;
     }
 }

tracer/src/Datadog.Trace/Debugger/Expressions/ProbeProcessor.cs

@@ -58,7 +58,8 @@ internal ProbeProcessor(ProbeDefinition probe)
                 probeType,
                 location,
                 evaluateAt,
-                HasCondition());
+                HasCondition(),
+                probe.Tags);
         }
 
         internal ProbeInfo ProbeInfo { get; }

tracer/src/Datadog.Trace/Debugger/Instrumentation/AsyncLineDebuggerInvoker.cs

@@ -8,6 +8,7 @@
 using System.Runtime.CompilerServices;
 using Datadog.Trace.Debugger.Expressions;
 using Datadog.Trace.Debugger.Helpers;
+using Datadog.Trace.Debugger.Instrumentation.Registry;
 using Datadog.Trace.Debugger.RateLimiting;
 using Datadog.Trace.Logging;
 
@@ -53,7 +54,8 @@ public static void LogLocal<TLocal>(ref TLocal local, int index, ref AsyncLineDe
                 }
 
                 var captureInfo = new CaptureInfo<TLocal>(value: local, type: typeof(TLocal), methodState: MethodState.LogLocal, name: localName, memberKind: ScopeMemberKind.Local);
-                if (!ProbeExpressionsProcessor.Instance.Process(state.ProbeId, ref captureInfo, state.SnapshotCreator))
+
+                if (!state.ProbeData.Processor.Process(ref captureInfo, state.SnapshotCreator))
                 {
                     state.IsActive = false;
                 }
@@ -104,6 +106,7 @@ public static void LogException(Exception exception, ref AsyncLineDebuggerState
         /// </summary>
         /// <typeparam name="TTarget">Target type</typeparam>
         /// <param name="probeId">The id of the probe</param>
+        /// <param name="probeMetadataIndex">The index used to lookup for the <see cref="ProbeData"/></param>
         /// <param name="instance">Instance value</param>
         /// <param name="methodHandle">The handle of the executing method</param>
         /// <param name="typeHandle">The handle of the type</param>
@@ -112,7 +115,7 @@ public static void LogException(Exception exception, ref AsyncLineDebuggerState
         /// <param name="probeFilePath">The path to the file of the probe</param>
         /// <returns>Live debugger state</returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public static AsyncLineDebuggerState BeginLine<TTarget>(string probeId, TTarget instance, RuntimeMethodHandle methodHandle, RuntimeTypeHandle typeHandle, int methodMetadataIndex, int lineNumber, string probeFilePath)
+        public static AsyncLineDebuggerState BeginLine<TTarget>(string probeId, int probeMetadataIndex, TTarget instance, RuntimeMethodHandle methodHandle, RuntimeTypeHandle typeHandle, int methodMetadataIndex, int lineNumber, string probeFilePath)
         {
             try
             {
@@ -125,7 +128,7 @@ public static AsyncLineDebuggerState BeginLine<TTarget>(string probeId, TTarget
                 }
 
                 // Assess if we have metadata associated with the given index
-                if (!MethodMetadataProvider.IsIndexExists(methodMetadataIndex))
+                if (!MethodMetadataCollection.Instance.IsIndexExists(methodMetadataIndex))
                 {
                     // State machine is null when we run in Optimized code and the original async method was generic,
                     // in which case the state machine is a generic value type.
@@ -137,25 +140,32 @@ public static AsyncLineDebuggerState BeginLine<TTarget>(string probeId, TTarget
                         Log.Warning($"{nameof(BeginLine)}: hoisted 'this' has not found. {kickoffInfo.KickoffParentType.Name}.{kickoffInfo.KickoffMethod.Name}");
                     }
 
-                    if (!MethodMetadataProvider.TryCreateAsyncMethodMetadataIfNotExists(instance, methodMetadataIndex, in methodHandle, in typeHandle, kickoffInfo))
+                    if (!MethodMetadataCollection.Instance.TryCreateAsyncMethodMetadataIfNotExists(instance, methodMetadataIndex, in methodHandle, in typeHandle, kickoffInfo))
                     {
-                        Log.Warning($"BeginMethod_StartMarker: Failed to receive the InstrumentedMethodInfo associated with the executing method. type = {typeof(TTarget)}, instance type name = {instance.GetType().Name}, methodMetadataId = {methodMetadataIndex}");
+                        Log.Warning($"([Async]BeginLine: Failed to receive the InstrumentedMethodInfo associated with the executing method. type = {typeof(TTarget)}, instance type name = {instance.GetType().Name}, methodMetadataId = {methodMetadataIndex}, probeId = {probeId}");
                         return CreateInvalidatedAsyncLineDebuggerState();
                     }
                 }
 
+                if (!ProbeMetadataCollection.Instance.TryCreateProbeMetadataIfNotExists(probeMetadataIndex, probeId))
+                {
+                    Log.Warning($"[Async]BeginLine: Failed to receive the ProbeData associated with the executing probe. type = {typeof(TTarget)}, instance type name = {instance?.GetType().Name}, probeMetadataIndex = {probeMetadataIndex}, probeId = {probeId}");
+                    return CreateInvalidatedAsyncLineDebuggerState();
+                }
+
                 var kickoffParentObject = AsyncHelper.GetAsyncKickoffThisObject(instance);
-                var state = new AsyncLineDebuggerState(probeId, scope: default, DateTimeOffset.UtcNow, methodMetadataIndex, lineNumber, probeFilePath, instance, kickoffParentObject);
+                var state = new AsyncLineDebuggerState(probeId, scope: default, DateTimeOffset.UtcNow, methodMetadataIndex, probeMetadataIndex, lineNumber, probeFilePath, instance, kickoffParentObject);
 
                 if (!state.SnapshotCreator.ProbeHasCondition &&
-                    !ProbeRateLimiter.Instance.Sample(probeId))
+                    !state.ProbeData.Sampler.Sample())
                 {
                     return CreateInvalidatedAsyncLineDebuggerState();
                 }
 
                 var asyncInfo = new AsyncCaptureInfo(state.MoveNextInvocationTarget, state.KickoffInvocationTarget, state.MethodMetadataInfo.KickoffInvocationTargetType, hoistedLocals: state.MethodMetadataInfo.AsyncMethodHoistedLocals, hoistedArgs: state.MethodMetadataInfo.AsyncMethodHoistedArguments);
                 var captureInfo = new CaptureInfo<Type>(value: null, type: state.MethodMetadataInfo.DeclaringType, methodState: MethodState.BeginLineAsync, localsCount: state.MethodMetadataInfo.LocalVariableNames.Length, argumentsCount: state.MethodMetadataInfo.ParameterNames.Length, lineCaptureInfo: new LineCaptureInfo(lineNumber, probeFilePath), asyncCaptureInfo: asyncInfo);
-                if (!ProbeExpressionsProcessor.Instance.Process(probeId, ref captureInfo, state.SnapshotCreator))
+
+                if (!state.ProbeData.Processor.Process(ref captureInfo, state.SnapshotCreator))
                 {
                     state.IsActive = false;
                 }
@@ -191,7 +201,7 @@ public static void EndLine(ref AsyncLineDebuggerState state)
                 state.HasLocalsOrReturnValue = false;
                 var asyncCaptureInfo = new AsyncCaptureInfo(state.MoveNextInvocationTarget, state.KickoffInvocationTarget, state.MethodMetadataInfo.KickoffInvocationTargetType, kickoffMethod: state.MethodMetadataInfo.KickoffMethod, hoistedArgs: state.MethodMetadataInfo.AsyncMethodHoistedArguments, hoistedLocals: state.MethodMetadataInfo.AsyncMethodHoistedLocals);
                 var captureInfo = new CaptureInfo<object>(memberKind: ScopeMemberKind.This, methodState: MethodState.EndLineAsync, hasLocalOrArgument: hasArgumentsOrLocals, asyncCaptureInfo: asyncCaptureInfo, lineCaptureInfo: new LineCaptureInfo(state.LineNumber, state.ProbeFilePath));
-                ProbeExpressionsProcessor.Instance.Process(state.ProbeId, ref captureInfo, state.SnapshotCreator);
+                state.ProbeData.Processor.Process(ref captureInfo, state.SnapshotCreator);
             }
             catch (Exception e)
             {

tracer/src/Datadog.Trace/Debugger/Instrumentation/AsyncLineDebuggerState.cs

@@ -1,11 +1,12 @@
- // <copyright file="AsyncLineDebuggerState.cs" company="Datadog">
+// <copyright file="AsyncLineDebuggerState.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.ComponentModel;
 using System.Runtime.CompilerServices;
+using Datadog.Trace.Debugger.Instrumentation.Registry;
 using Datadog.Trace.Debugger.Snapshots;
 
 namespace Datadog.Trace.Debugger.Instrumentation
@@ -24,7 +25,7 @@ public ref struct AsyncLineDebuggerState
         private readonly int _lineNumber;
 
         /// <summary>
-        /// Used to perform a fast lookup to grab the proper <see cref="Instrumentation.MethodMetadataInfo"/>.
+        /// Used to perform a fast lookup to grab the proper <see cref="Registry.MethodMetadataInfo"/>.
         /// This index is hard-coded into the method's instrumented bytecode.
         /// </summary>
         private readonly int _methodMetadataIndex;
@@ -43,12 +44,13 @@ public ref struct AsyncLineDebuggerState
         /// <param name="probeId">The id of the probe</param>
         /// <param name="scope">Scope instance</param>
         /// <param name="startTime">The timestamp captured when the relevant line of code was hit</param>
-        /// <param name="methodMetadataIndex">The unique index of the method's <see cref="Instrumentation.MethodMetadataInfo"/></param>
+        /// <param name="methodMetadataIndex">The unique index of the method's <see cref="Registry.MethodMetadataInfo"/></param>
+        /// <param name="probeMetadataIndex">The unique index of the probe <see cref="ProbeData"/></param>
         /// <param name="lineNumber">The line number where the probe is located on</param>
         /// <param name="probeFilePath">The path to the file of the probe</param>
         /// <param name="invocationTarget">The instance object (or null for static methods)</param>
         /// <param name="kickoffInvocationTarget">The instance object (or null for static methods) of the kickoff method</param>
-        internal AsyncLineDebuggerState(string probeId, Scope scope, DateTimeOffset? startTime, int methodMetadataIndex, int lineNumber, string probeFilePath, object invocationTarget, object kickoffInvocationTarget)
+        internal AsyncLineDebuggerState(string probeId, Scope scope, DateTimeOffset? startTime, int methodMetadataIndex, int probeMetadataIndex, int lineNumber, string probeFilePath, object invocationTarget, object kickoffInvocationTarget)
         {
             _probeId = probeId;
             _scope = scope;
@@ -57,12 +59,15 @@ internal AsyncLineDebuggerState(string probeId, Scope scope, DateTimeOffset? sta
             _lineNumber = lineNumber;
             _probeFilePath = probeFilePath;
             HasLocalsOrReturnValue = false;
-            SnapshotCreator = DebuggerSnapshotCreator.BuildSnapshotCreator(probeId);
+            ProbeData = ProbeMetadataCollection.Instance.Get(probeMetadataIndex);
+            SnapshotCreator = DebuggerSnapshotCreator.BuildSnapshotCreator(ProbeData.Processor);
             _moveNextInvocationTarget = invocationTarget;
             _kickoffInvocationTarget = kickoffInvocationTarget;
         }
 
-        internal ref MethodMetadataInfo MethodMetadataInfo => ref MethodMetadataProvider.Get(_methodMetadataIndex);
+        internal ref MethodMetadataInfo MethodMetadataInfo => ref MethodMetadataCollection.Instance.Get(_methodMetadataIndex);
+
+        internal ProbeData ProbeData { get; }
 
         /// <summary>
         /// Gets the LiveDebugger SnapshotCreator