Bug/traceactivity attribute not working when assigned to class

src/OpenTelemetry.Instrumentation.Digma.Tests/Stubs/DecoratedService.cs

@@ -28,8 +28,11 @@ public void MethodJaggedAndMultiDimArraysParams(Action stateValidation, out stri
         bool[][][] jaggedArrayOfBools, short[,,,][,][,,] multiDimArrayOfShorts,
         long[,,][][,][] mixMultiDimAndJaggedArraysOfLongs
     );
-}
 
+    void MethodExplicitlyMarkedForTracingWithAttributes(Action action);
+}
+
+[ActivitiesAttributes("att1:value1")]
 public class DecoratedService : IDecoratedService
 {
     [TraceActivity()]
@@ -38,16 +41,25 @@ public void MethodExplicitlyMarkedForTracing(Action stateValidation)
         var v = Activity.Current;
         stateValidation();
     }
+
+    [TraceActivity()]
+    [ActivitiesAttributes("att1:value1")]
+    public void MethodExplicitlyMarkedForTracingWithAttributes(Action stateValidation)
+    {
+        stateValidation();
+    }
 
     [TraceActivity()]
     public async Task AsyncMethodExplicitlyMarkedForTracing(Action stateValidation)
     {
         stateValidation();
     }
 
-
+    [TraceActivity()]
     public async Task AsyncVoid()
     {
+        var v = Activity.Current;
+
         await Task.Delay(100);
     }
 

src/OpenTelemetry.Instrumentation.Digma.Tests/TestTracingDecorator.cs

@@ -16,8 +16,35 @@ namespace OpenTelemetry.Instrumentation.Digma.Tests;
 public class TestTracingDecorator
 {
     private static readonly string ServiceInterfaceFqn =
-        "OpenTelemetry.Instrumentation.Digma.Tests.Stubs.IDecoratedService";
+        "OpenTelemetry.Instrumentation.Digma.Tests.Stubs.DecoratedService";
 
+    [TestMethod]
+    public void Activity_Created_For_Attribute_Marked_Method()
+    {
+        DecoratedService service = new DecoratedService();
+        IDecoratedService tracingDecorator = TraceDecorator<IDecoratedService>.Create(service);
+        tracingDecorator.MethodExplicitlyMarkedForTracing(() =>
+        {
+            Assert.IsNotNull(Activity.Current);
+            AssertHasCommonTags(Activity.Current, ServiceInterfaceFqn,
+                "MethodExplicitlyMarkedForTracing", "Action");
+        });
+    }
+
+    [TestMethod]
+    public void Attributes_Injected_To_Marked_Method()
+    {
+        DecoratedService service = new DecoratedService();
+        IDecoratedService tracingDecorator = TraceDecorator<IDecoratedService>.Create(service);
+        tracingDecorator.MethodExplicitlyMarkedForTracingWithAttributes(() =>
+        {
+            Assert.IsNotNull(Activity.Current);
+            AssertHasCommonTags(Activity.Current, ServiceInterfaceFqn,
+                "MethodExplicitlyMarkedForTracingWithAttributes", "Action");
+            AssertHasTag(Activity.Current, "att1", "value1");
+
+        });
+    }
     private MockProcessor _mockProcessor = new ();
 
     // [TestMethod]
@@ -146,8 +173,9 @@ public async Task Activity_Async_Error()
         // Act #1
         try
         {
-            await decoratedService.AsyncError();
-            Assert.Fail();
+            var task = decoratedService.AsyncError();
+            await task;
+            throw task.Exception!;
         }
         catch (Exception e)
         {
@@ -174,4 +202,10 @@ private void AssertHasCommonTags(Activity? activity,
                 new KeyValuePair<string, string>("code.function.parameter.types", expectedParameterTypes));
         }
     }
+
+    private void AssertHasTag(Activity? activity, string name, string value)
+    {
+        var kvpTags = activity.Tags.ToArray();
+        CollectionAssert.Contains(kvpTags, new KeyValuePair<string, string>(name, value));
+    }
 }

src/OpenTelemetry.Instrumentation.Digma/Helpers/TracingDecorator.cs

@@ -1,5 +1,7 @@
-using System.Diagnostics;
+using System.Collections.Concurrent;
+using System.Diagnostics;
 using System.Reflection;
+using Microsoft.Extensions.Caching.Memory;
 using OpenTelemetry.Instrumentation.Digma.Helpers.Attributes;
 using OpenTelemetry.Trace;
 
@@ -12,6 +14,9 @@ public class TraceDecorator<TDecorated> : DispatchProxy where TDecorated : class
     private IActivityNamingSchema _namingSchema = new MethodNameSchema();
     private bool _decorateAllMethods = true;
 
+    private readonly ConcurrentDictionary<string, NoActivityAttribute?> _methodNoActivityAttributeCache = new();
+    private readonly ConcurrentDictionary<string, ActivitiesAttributesAttribute?> _methodActivitiesTagsCache = new();
+    private readonly ConcurrentDictionary<string, TraceActivityAttribute?> _methodActivityAttributeCache = new();
     /// <summary>
     /// Creates a new TraceDecorator instance wrapping the specific object and implementing the TDecorated interface 
     /// </summary>
@@ -39,71 +44,141 @@ private void SetParameters(TDecorated decorated, IActivityNamingSchema? spanNami
         }
     }
 
-    private object? InvokeDecoratedExecution(Activity? activity, MethodInfo? targetMethod, object?[]? args, bool? recordException)
+    private bool IsAsync(MethodInfo targetMethod)
     {
+        var taskType = typeof(Task);
+        return targetMethod.ReturnType == taskType ||
+               targetMethod.ReturnType.IsSubclassOf(taskType);
+    }
+    private object? InvokeAsyncDecoratedExecution(Activity? activity, MethodInfo? targetMethod, object?[]? args,
+        bool? recordException)
+    {
+        activity.Start();
+        var resultTask = targetMethod.Invoke(_decorated, args) as Task;
+
+        resultTask.ContinueWith(task =>
+        {
+            if (task.Exception != null && (recordException ?? true))
+            {
+                activity?.RecordException(task.Exception);
+            }
+
+            activity?.Stop();
+
+            activity?.Dispose();
+
+
+        },TaskContinuationOptions.AttachedToParent | TaskContinuationOptions.ExecuteSynchronously);
+
+        return resultTask;
+    }
+
+    private object? InvokeDecoratedExecution(Activity? activity, MethodInfo? targetMethod, object?[]? args,
+        bool? recordException)
+    {
+
         object? result;
         try
         {
-            result = targetMethod.Invoke(_decorated, args);
+            using (activity)
+            {
+                result = targetMethod.Invoke(_decorated, args);
+            }
         }
         catch (Exception e)
         {
-            if(recordException ?? true)
+            if (recordException ?? true)
                 activity?.RecordException(e);
-
-            activity?.Dispose();
+
             throw;
         }
 
-        if (result is Task resultTask)
+
+        return result;
+
+
+
+    }
+
+    protected override object? Invoke(MethodInfo? targetMethod, object?[]? args)
+    {
+        if (targetMethod != null)
         {
-            resultTask.ContinueWith(task =>
+
+            var noActivityAttribute = GetMethodNoActivityAttribute(targetMethod);
+            var activityAttribute = GetMethodActivityAttribute(targetMethod);
+
+            if (noActivityAttribute == null && _decorateAllMethods )
             {
-                if (task.Exception != null && (recordException ?? true))
+                var classType = _decorated!.GetType();
+
+                var defaultSpanName = _namingSchema.GetSpanName(classType, targetMethod);
+                var activity = _activity.StartActivity(activityAttribute?.Name ?? defaultSpanName);
+
+                SpanUtils.AddCommonTags(classType,targetMethod, activity);
+                InjectAttributes(targetMethod, activity);
+
+                if (IsAsync(targetMethod))
                 {
-                    activity?.RecordException(task.Exception);
+                    return InvokeAsyncDecoratedExecution(activity, targetMethod, args,
+                        activityAttribute?.RecordExceptions);
                 }
-
-                activity?.Dispose();
-            }, TaskContinuationOptions.AttachedToParent | TaskContinuationOptions.ExecuteSynchronously);
+                
+                return InvokeDecoratedExecution(activity, targetMethod, args, activityAttribute?.RecordExceptions);
+            }
 
-            return result;
+
         }
 
-        activity?.Dispose();
-        return result;
+        return InvokeDecoratedExecution(null, targetMethod, args, null);
     }
 
-    protected override object? Invoke(MethodInfo? targetMethod, object?[]? args)
+    private TAttribute? GetInstanceMethodAttribute<TAttribute>(MethodInfo targetMethod) where TAttribute : Attribute
     {
-        var noActivityAttribute = targetMethod.GetCustomAttribute<NoActivityAttribute>(false);
-        var activityAttribute = targetMethod.GetCustomAttribute<TraceActivityAttribute>(false);
+        return _decorated.GetType().GetMethod(targetMethod.Name)?.GetCustomAttribute<TAttribute>(inherit:true);
+    }
 
-        if (noActivityAttribute == null && (_decorateAllMethods || activityAttribute != null))
-        {
-            var classType = _decorated!.GetType();
-
-            var defaultSpanName = _namingSchema.GetSpanName(classType, targetMethod);
-            var activity = _activity.StartActivity(activityAttribute?.Name ?? defaultSpanName);
 
-            SpanUtils.AddCommonTags(classType, targetMethod, activity);
-            InjectAttributes(targetMethod, activity);
+    private TAttribute? GetFromOrStoreInAttributeCache<TAttribute>(
+        MethodInfo targetMethod,
+        IDictionary<string, TAttribute?> cache) where TAttribute : Attribute
+    {
+        TAttribute? attributeInfo;
 
-            return InvokeDecoratedExecution(activity, targetMethod, args, activityAttribute?.RecordExceptions);
-        }
+        if (!cache.TryGetValue(targetMethod.Name, out attributeInfo))
 
-        return InvokeDecoratedExecution(null, targetMethod, args, null);
+        {
+            attributeInfo = GetInstanceMethodAttribute<TAttribute>(targetMethod);
+            cache[targetMethod.Name] = attributeInfo;
+
+        }
+
+        return attributeInfo;
+    }
+    private ActivitiesAttributesAttribute? GetActivityTagsForInstanceMethod(MethodInfo targetMethod)
+    {
+        return GetFromOrStoreInAttributeCache(targetMethod, _methodActivitiesTagsCache);
+    }
+
+    private TraceActivityAttribute? GetMethodActivityAttribute(MethodInfo targetMethod)
+    {
+        return GetFromOrStoreInAttributeCache(targetMethod, _methodActivityAttributeCache);
+    }
+
+    private NoActivityAttribute? GetMethodNoActivityAttribute(MethodInfo targetMethod)
+    {
+        return GetFromOrStoreInAttributeCache(targetMethod, _methodNoActivityAttributeCache);
     }
 
     private void InjectAttributes(MethodInfo targetMethod, Activity? activity)
     {
-        var methodActivityAttributes = targetMethod.GetCustomAttribute<ActivitiesAttributesAttribute>(inherit: false);
+        var methodActivityAttributes = GetActivityTagsForInstanceMethod(targetMethod);
         var classActivityAttributes =
             _decorated.GetType().GetCustomAttribute<ActivitiesAttributesAttribute>(inherit: false);
 
         if (methodActivityAttributes != null)
         {
-            foreach (var key in classActivityAttributes.Attributes.Keys)
+            foreach (var key in methodActivityAttributes.Attributes.Keys)
             {
                 activity.AddTag(key, methodActivityAttributes.Attributes[key]);
             }
@@ -113,7 +188,7 @@ private void InjectAttributes(MethodInfo targetMethod, Activity? activity)
         {
             foreach (var key in classActivityAttributes.Attributes.Keys)
             {
-                activity.AddTag(key, methodActivityAttributes.Attributes[key]);
+                activity.AddTag(key, classActivityAttributes.Attributes[key]);
             }
         }
     }