chore: Fix Lambda timeout with Tracing 1.5.1 using async methods

libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs → libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspectHandler.cs

@@ -1,33 +1,31 @@
 /*
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
- *
+ * 
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
  * A copy of the License is located at
- *
+ * 
  *  http://aws.amazon.com/apache2.0
- *
+ * 
  * or in the "license" file accompanying this file. This file is distributed
  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
  * express or implied. See the License for the specific language governing
  * permissions and limitations under the License.
  */
 
 using System;
-using System.Linq;
 using System.Runtime.ExceptionServices;
 using System.Text;
-using AspectInjector.Broker;
 using AWS.Lambda.Powertools.Common;
 
 namespace AWS.Lambda.Powertools.Tracing.Internal;
 
 /// <summary>
-///     This aspect will automatically trace all function handlers.
-///     Scope.Global is singleton
+///     Class TracingAspectHandler.
+///     Implements the <see cref="IMethodAspectHandler" />
 /// </summary>
-[Aspect(Scope.Global)]
-public class TracingAspect
+/// <seealso cref="IMethodAspectHandler" />
+internal class TracingAspectHandler : IMethodAspectHandler
 {
     /// <summary>
     ///     The Powertools for AWS Lambda (.NET) configurations
@@ -48,130 +46,120 @@ public class TracingAspect
     ///     If true, capture annotations
     /// </summary>
     private static bool _captureAnnotations = true;
-
+    
     /// <summary>
     ///     If true, annotations have been captured
     /// </summary>
     private bool _isAnnotationsCaptured;
-
+    
     /// <summary>
     ///     Tracing namespace
     /// </summary>
-    private string _namespace;
-
+    private readonly string _namespace;
+    
     /// <summary>
     ///     The capture mode
     /// </summary>
-    private TracingCaptureMode _captureMode;
+    private readonly TracingCaptureMode _captureMode;
+
+    /// <summary>
+    ///     The segment name
+    /// </summary>
+    private readonly string _segmentName;
 
     /// <summary>
-    /// Initializes a new instance
+    ///     Initializes a new instance of the <see cref="TracingAspectHandler" /> class.
     /// </summary>
-    public TracingAspect()
+    /// <param name="segmentName">Name of the segment.</param>
+    /// <param name="nameSpace">The namespace.</param>
+    /// <param name="captureMode">The capture mode.</param>
+    /// <param name="powertoolsConfigurations">The Powertools for AWS Lambda (.NET) configurations.</param>
+    /// <param name="xRayRecorder">The X-Ray recorder.</param>
+    internal TracingAspectHandler
+    (
+        string segmentName,
+        string nameSpace,
+        TracingCaptureMode captureMode,
+        IPowertoolsConfigurations powertoolsConfigurations,
+        IXRayRecorder xRayRecorder
+    )
     {
-        _xRayRecorder = XRayRecorder.Instance;
-        _powertoolsConfigurations = PowertoolsConfigurations.Instance;
+        _segmentName = segmentName;
+        _namespace = nameSpace;
+        _captureMode = captureMode;
+        _powertoolsConfigurations = powertoolsConfigurations;
+        _xRayRecorder = xRayRecorder;
     }
 
     /// <summary>
-    /// the code is executed instead of the target method.
-    /// The call to original method is wrapped around the following code
-    /// the original code is called with var result = target(args);
+    ///     Handles the <see cref="E:Entry" /> event.
     /// </summary>
-    /// <param name="name"></param>
-    /// <param name="args"></param>
-    /// <param name="target"></param>
-    /// <param name="triggers"></param>
-    /// <returns></returns>
-    [Advice(Kind.Around)]
-    public object Around(
-        [Argument(Source.Name)] string name,
-        [Argument(Source.Arguments)] object[] args,
-        [Argument(Source.Target)] Func<object[], object> target,
-        [Argument(Source.Triggers)] Attribute[] triggers)
+    /// <param name="eventArgs">
+    ///     The <see cref="T:AWS.Lambda.Powertools.Aspects.AspectEventArgs" /> instance containing the
+    ///     event data.
+    /// </param>
+    public void OnEntry(AspectEventArgs eventArgs)
     {
-        // Before running Function
-
-        var trigger = triggers.OfType<TracingAttribute>().First();
-        try
-        {
-            if (TracingDisabled())
-                return target(args);
-
-            _namespace = trigger.Namespace;
-
-            var segmentName = !string.IsNullOrWhiteSpace(trigger.SegmentName) ? trigger.SegmentName : $"## {name}";
-            var nameSpace = GetNamespace();
-
-            _xRayRecorder.BeginSubsegment(segmentName);
-            _xRayRecorder.SetNamespace(nameSpace);
-
-            if (_captureAnnotations)
-            {
-                _xRayRecorder.AddAnnotation("ColdStart", _isColdStart);
-
-                _captureAnnotations = false;
-                _isAnnotationsCaptured = true;
-
-                if (_powertoolsConfigurations.IsServiceDefined)
-                    _xRayRecorder.AddAnnotation("Service", _powertoolsConfigurations.Service);
-            }
-
-            _isColdStart = false;
+        if(TracingDisabled())
+            return;
 
-            // return of the handler
-            var result = target(args);
+        var segmentName = !string.IsNullOrWhiteSpace(_segmentName) ? _segmentName : $"## {eventArgs.Name}";
+        var nameSpace = GetNamespace();
 
-            // must get capture after all subsegments run
-            _captureMode = trigger.CaptureMode;
+        _xRayRecorder.BeginSubsegment(segmentName);
+        _xRayRecorder.SetNamespace(nameSpace);
 
-            if (CaptureResponse())
-            {
-                _xRayRecorder.AddMetadata
-                (
-                    nameSpace,
-                    $"{name} response",
-                    result
-                );
-            }
-
-            // after 
-            return result;
-        }
-        catch (Exception e)
+        if (_captureAnnotations)
         {
-            _captureMode = trigger.CaptureMode;
-            HandleException(e, name);
-            throw;
+            _xRayRecorder.AddAnnotation("ColdStart", _isColdStart);
+
+            _captureAnnotations = false;
+            _isAnnotationsCaptured = true;
+
+            if (_powertoolsConfigurations.IsServiceDefined)
+                _xRayRecorder.AddAnnotation("Service", _powertoolsConfigurations.Service);
         }
+
+        _isColdStart = false;
     }
 
     /// <summary>
-    /// the code is injected after the method ends.
+    ///     Called when [success].
     /// </summary>
-    [Advice(Kind.After)]
-    public void OnExit()
+    /// <param name="eventArgs">
+    ///     The <see cref="T:AWS.Lambda.Powertools.Aspects.AspectEventArgs" /> instance containing the
+    ///     event data.
+    /// </param>
+    /// <param name="result">The result.</param>
+    public void OnSuccess(AspectEventArgs eventArgs, object result)
     {
-        if (TracingDisabled())
-            return;
-
-        if (_isAnnotationsCaptured)
-            _captureAnnotations = true;
-
-        _xRayRecorder.EndSubsegment();
+        if (CaptureResponse())
+        {
+            var nameSpace = GetNamespace();
+
+            _xRayRecorder.AddMetadata
+            (
+                nameSpace,
+                $"{eventArgs.Name} response",
+                result
+            );
+        }
     }
 
     /// <summary>
-    /// Code that handles when exceptions occur in the client method
+    ///     Called when [exception].
     /// </summary>
-    /// <param name="exception"></param>
-    /// <param name="name"></param>
-    private void HandleException(Exception exception, string name)
+    /// <param name="eventArgs">
+    ///     The <see cref="T:AWS.Lambda.Powertools.Aspects.AspectEventArgs" /> instance containing the
+    ///     event data.
+    /// </param>
+    /// <param name="exception">The exception.</param>
+    public void OnException(AspectEventArgs eventArgs, Exception exception)
     {
         if (CaptureError())
         {
             var nameSpace = GetNamespace();
-
+            
             var sb = new StringBuilder();
             sb.AppendLine($"Exception type: {exception.GetType()}");
             sb.AppendLine($"Exception message: {exception.Message}");
@@ -185,48 +173,45 @@ private void HandleException(Exception exception, string name)
                 sb.AppendLine($"Stack trace: {exception.InnerException.StackTrace}");
                 sb.AppendLine("---END Inner Exception");
             }
-
+            
             _xRayRecorder.AddMetadata
             (
                 nameSpace,
-                $"{name} error",
+                $"{eventArgs.Name} error",
                 sb.ToString()
             );
         }
 
-        // // The purpose of ExceptionDispatchInfo.Capture is to capture a potentially mutating exception's StackTrace at a point in time:
-        // // https://learn.microsoft.com/en-us/dotnet/standard/exceptions/best-practices-for-exceptions#capture-exceptions-to-rethrow-later
+        // The purpose of ExceptionDispatchInfo.Capture is to capture a potentially mutating exception's StackTrace at a point in time:
+        // https://learn.microsoft.com/en-us/dotnet/standard/exceptions/best-practices-for-exceptions#capture-exceptions-to-rethrow-later
         ExceptionDispatchInfo.Capture(exception).Throw();
     }
 
     /// <summary>
-    ///     Gets the namespace.
+    ///     Handles the <see cref="E:Exit" /> event.
     /// </summary>
-    /// <returns>System.String.</returns>
-    private string GetNamespace()
+    /// <param name="eventArgs">
+    ///     The <see cref="T:AWS.Lambda.Powertools.Aspects.AspectEventArgs" /> instance containing the
+    ///     event data.
+    /// </param>
+    public void OnExit(AspectEventArgs eventArgs)
     {
-        return !string.IsNullOrWhiteSpace(_namespace) ? _namespace : _powertoolsConfigurations.Service;
+        if(TracingDisabled())
+            return;
+
+        if (_isAnnotationsCaptured)
+            _captureAnnotations = true;
+
+        _xRayRecorder.EndSubsegment();
     }
 
     /// <summary>
-    /// Method that checks if tracing is disabled
+    ///     Gets the namespace.
     /// </summary>
-    /// <returns></returns>
-    private bool TracingDisabled()
+    /// <returns>System.String.</returns>
+    private string GetNamespace()
     {
-        if (_powertoolsConfigurations.TracingDisabled)
-        {
-            Console.WriteLine("Tracing has been disabled via env var POWERTOOLS_TRACE_DISABLED");
-            return true;
-        }
-
-        if (!_powertoolsConfigurations.IsLambdaEnvironment)
-        {
-            Console.WriteLine("Running outside Lambda environment; disabling Tracing");
-            return true;
-        }
-
-        return false;
+        return !string.IsNullOrWhiteSpace(_namespace) ? _namespace : _powertoolsConfigurations.Service;
     }
 
     /// <summary>
@@ -235,6 +220,9 @@ private bool TracingDisabled()
     /// <returns><c>true</c> if tracing should capture responses, <c>false</c> otherwise.</returns>
     private bool CaptureResponse()
     {
+        if(TracingDisabled())
+            return false;
+
         switch (_captureMode)
         {
             case TracingCaptureMode.EnvironmentVariable:
@@ -255,9 +243,9 @@ private bool CaptureResponse()
     /// <returns><c>true</c> if tracing should capture errors, <c>false</c> otherwise.</returns>
     private bool CaptureError()
     {
-        if (TracingDisabled())
+        if(TracingDisabled())
             return false;
-
+        
         switch (_captureMode)
         {
             case TracingCaptureMode.EnvironmentVariable:
@@ -271,7 +259,28 @@ private bool CaptureError()
                 return false;
         }
     }
+
+    /// <summary>
+    ///     Tracing disabled.
+    /// </summary>
+    /// <returns><c>true</c> if tracing is disabled, <c>false</c> otherwise.</returns>
+    private bool TracingDisabled()
+    {
+        if (_powertoolsConfigurations.TracingDisabled)
+        {
+            Console.WriteLine("Tracing has been disabled via env var POWERTOOLS_TRACE_DISABLED");
+            return true;
+        }
+
+        if (!_powertoolsConfigurations.IsLambdaEnvironment)
+        {
+            Console.WriteLine("Running outside Lambda environment; disabling Tracing");
+            return true;
+        }
 
+        return false;
+    }
+
     /// <summary>
     ///     Resets static variables for test.
     /// </summary>
@@ -280,4 +289,4 @@ internal static void ResetForTest()
         _isColdStart = true;
         _captureAnnotations = true;
     }
-}
+}