Merge pull request #632 from hjgraca/fix(tracing)-batch-handler-result-null-reference

hjgraca · web-flow · commit 871f18be6d4e · 2024-08-26T10:31:07.000+01:00
chore: Fix Max depth serialization bug when using Batch RecordHanderResult and Tracing
diff --git a/libraries/src/AWS.Lambda.Powertools.BatchProcessing/RecordHandlerResult.cs b/libraries/src/AWS.Lambda.Powertools.BatchProcessing/RecordHandlerResult.cs
@@ -23,7 +23,7 @@ public class RecordHandlerResult
     /// <summary>
     /// Returns an empty <see cref="RecordHandlerResult"/> value.
     /// </summary>
-    public static readonly RecordHandlerResult None = new();
+    public static RecordHandlerResult None { get; } = null!;
 
     /// <summary>
     /// Convenience method for the creation of a <see cref="RecordHandlerResult"/>.
diff --git a/libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs b/libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs
@@ -15,7 +15,6 @@
 
 using System;
 using System.Linq;
-using System.Reflection;
 using System.Runtime.ExceptionServices;
 using System.Text;
 using AspectInjector.Broker;
@@ -39,7 +38,7 @@ public class TracingAspect
     ///     X-Ray Recorder
     /// </summary>
     private readonly IXRayRecorder _xRayRecorder;
-    
+
     /// <summary>
     ///     If true, then is cold start
     /// </summary>
@@ -49,17 +48,17 @@ 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;
-    
+
     /// <summary>
     ///     The capture mode
     /// </summary>
@@ -73,18 +72,14 @@ public TracingAspect()
         _xRayRecorder = XRayRecorder.Instance;
         _powertoolsConfigurations = PowertoolsConfigurations.Instance;
     }
-    
+
     /// <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);
     /// </summary>
-    /// <param name="instance"></param>
     /// <param name="name"></param>
     /// <param name="args"></param>
-    /// <param name="hostType"></param>
-    /// <param name="method"></param>
-    /// <param name="returnType"></param>
     /// <param name="target"></param>
     /// <param name="triggers"></param>
     /// <returns></returns>
@@ -97,38 +92,39 @@ public object Around(
     {
         // Before running Function
 
+        var trigger = triggers.OfType<TracingAttribute>().First();
         try
         {
-            var trigger = triggers.OfType<TracingAttribute>().First();
-        
-            if(TracingDisabled())
+            if (TracingDisabled())
                 return target(args);
 
             _namespace = trigger.Namespace;
-            _captureMode = trigger.CaptureMode;
-        
+
             var segmentName = !string.IsNullOrWhiteSpace(trigger.SegmentName) ? trigger.SegmentName : $"## {name}";
             var nameSpace = GetNamespace();
 
             _xRayRecorder.BeginSubsegment(segmentName);
             _xRayRecorder.SetNamespace(nameSpace);
 
-            if (_captureAnnotations)    
+            if (_captureAnnotations)
             {
                 _xRayRecorder.AddAnnotation("ColdStart", _isColdStart);
-            
+
                 _captureAnnotations = false;
                 _isAnnotationsCaptured = true;
 
                 if (_powertoolsConfigurations.IsServiceDefined)
                     _xRayRecorder.AddAnnotation("Service", _powertoolsConfigurations.Service);
             }
-        
+
             _isColdStart = false;
-        
+
             // return of the handler
             var result = target(args);
-        
+
+            // must get capture after all subsegments run
+            _captureMode = trigger.CaptureMode;
+
             if (CaptureResponse())
             {
                 _xRayRecorder.AddMetadata
@@ -138,12 +134,13 @@ public object Around(
                     result
                 );
             }
-        
+
             // after 
             return result;
         }
         catch (Exception e)
         {
+            _captureMode = trigger.CaptureMode;
             HandleException(e, name);
             throw;
         }
@@ -153,16 +150,17 @@ public object Around(
     /// the code is injected after the method ends.
     /// </summary>
     [Advice(Kind.After)]
-    public void OnExit() {
-        if(TracingDisabled())
+    public void OnExit()
+    {
+        if (TracingDisabled())
             return;
 
         if (_isAnnotationsCaptured)
             _captureAnnotations = true;
 
         _xRayRecorder.EndSubsegment();
     }
-    
+
     /// <summary>
     /// Code that handles when exceptions occur in the client method
     /// </summary>
@@ -173,7 +171,7 @@ private void HandleException(Exception exception, string name)
         if (CaptureError())
         {
             var nameSpace = GetNamespace();
-            
+
             var sb = new StringBuilder();
             sb.AppendLine($"Exception type: {exception.GetType()}");
             sb.AppendLine($"Exception message: {exception.Message}");
@@ -187,7 +185,7 @@ private void HandleException(Exception exception, string name)
                 sb.AppendLine($"Stack trace: {exception.InnerException.StackTrace}");
                 sb.AppendLine("---END Inner Exception");
             }
-            
+
             _xRayRecorder.AddMetadata
             (
                 nameSpace,
@@ -209,7 +207,7 @@ private string GetNamespace()
     {
         return !string.IsNullOrWhiteSpace(_namespace) ? _namespace : _powertoolsConfigurations.Service;
     }
-    
+
     /// <summary>
     /// Method that checks if tracing is disabled
     /// </summary>
@@ -230,7 +228,7 @@ private bool TracingDisabled()
 
         return false;
     }
-    
+
     /// <summary>
     ///     Captures the response.
     /// </summary>
@@ -250,16 +248,16 @@ private bool CaptureResponse()
                 return false;
         }
     }
-    
+
     /// <summary>
     ///     Captures the error.
     /// </summary>
     /// <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:
@@ -273,7 +271,7 @@ private bool CaptureError()
                 return false;
         }
     }
-    
+
     /// <summary>
     ///     Resets static variables for test.
     /// </summary>
diff --git a/libraries/src/AWS.Lambda.Powertools.Tracing/Internal/XRayRecorder.cs b/libraries/src/AWS.Lambda.Powertools.Tracing/Internal/XRayRecorder.cs
@@ -1,12 +1,12 @@
 /*
  * 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
@@ -41,7 +41,8 @@ internal class XRayRecorder : IXRayRecorder
     ///     Gets the instance.
     /// </summary>
     /// <value>The instance.</value>
-    public static IXRayRecorder Instance => _instance ??= new XRayRecorder(AWSXRayRecorder.Instance, PowertoolsConfigurations.Instance);
+    public static IXRayRecorder Instance =>
+        _instance ??= new XRayRecorder(AWSXRayRecorder.Instance, PowertoolsConfigurations.Instance);
 
     public XRayRecorder(IAWSXRayRecorder awsxRayRecorder, IPowertoolsConfigurations powertoolsConfigurations)
     {
@@ -56,7 +57,7 @@ public XRayRecorder(IAWSXRayRecorder awsxRayRecorder, IPowertoolsConfigurations
     ///     Checks whether current execution is in AWS Lambda.
     /// </summary>
     /// <returns>Returns true if current execution is in AWS Lambda.</returns>
-    private static bool _isLambda; 
+    private static bool _isLambda;
 
     /// <summary>
     ///     Gets the emitter.
@@ -112,14 +113,30 @@ public void AddMetadata(string nameSpace, string key, object value)
         if (_isLambda)
             _awsxRayRecorder.AddMetadata(nameSpace, key, value);
     }
-
+    
     /// <summary>
     ///     Ends the subsegment.
     /// </summary>
     public void EndSubsegment()
     {
-        if (_isLambda)
+        if (!_isLambda) return;
+        try
+        {
+            _awsxRayRecorder.EndSubsegment();
+        }
+        catch (Exception e)
+        {
+            // if it fails at this stage the data is lost
+            // so lets create a new subsegment with the error
+
+            Console.WriteLine("Error in Tracing utility - see Exceptions tab in Cloudwatch Traces");
+
+            _awsxRayRecorder.TraceContext.ClearEntity();
+            _awsxRayRecorder.BeginSubsegment("Error in Tracing utility - see Exceptions tab");
+            _awsxRayRecorder.AddException(e);
+            _awsxRayRecorder.MarkError();
             _awsxRayRecorder.EndSubsegment();
+        }
     }
 
     /// <summary>
diff --git a/libraries/src/Directory.Packages.props b/libraries/src/Directory.Packages.props
@@ -9,13 +9,13 @@
     <PackageVersion Include="AWSSDK.AppConfig" Version="3.7.301" />
     <PackageVersion Include="AWSSDK.AppConfigData" Version="3.7.301.15" />
     <PackageVersion Include="AWSSDK.DynamoDBv2" Version="3.7.301.18" />
-    <PackageVersion Include="AWSXRayRecorder.Handlers.AwsSdk" Version="2.12.0" />
+    <PackageVersion Include="AWSXRayRecorder.Handlers.AwsSdk" Version="2.13.0" />
     <PackageVersion Include="Microsoft.Extensions.Logging" Version="8.0.0" />
     <PackageVersion Include="AWSSDK.SecretsManager" Version="3.7.302.29" />
     <PackageVersion Include="AWSSDK.SimpleSystemsManagement" Version="3.7.303.2" />
     <PackageVersion Include="Microsoft.Extensions.Configuration" Version="8.0.0" />
-    <PackageVersion Include="AWSSDK.XRay" Version="3.7.300.54" />
-    <PackageVersion Include="AWSXRayRecorder.Core" Version="2.14.0" />
+    <PackageVersion Include="AWSSDK.XRay" Version="3.7.400.8" />
+    <PackageVersion Include="AWSXRayRecorder.Core" Version="2.15.0" />
     <PackageVersion Include="Amazon.Lambda.DynamoDBEvents" Version="3.1.0" />
     <PackageVersion Include="Amazon.Lambda.KinesisEvents" Version="2.2.0" />
     <PackageVersion Include="Amazon.Lambda.SQSEvents" Version="2.2.0" />
diff --git a/libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/Handlers/Handlers.cs b/libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/Handlers/Handlers.cs
@@ -83,4 +83,24 @@ public string[] HandleWithCaptureModeDisabled(bool exception = false)
             throw new Exception("Failed");
         return new[] { "A", "B" };
     }
+    
+    [Tracing(CaptureMode = TracingCaptureMode.ResponseAndError)]
+    public string DecoratedHandlerCaptureResponse()
+    {
+        DecoratedMethodCaptureDisabled();
+        return "Hello World";
+    }
+
+    [Tracing(CaptureMode = TracingCaptureMode.Disabled)]
+    public string DecoratedMethodCaptureDisabled()
+    {
+        DecoratedMethodCaptureEnabled();
+        return "DecoratedMethod Disabled";
+    }
+    
+    [Tracing(CaptureMode = TracingCaptureMode.ResponseAndError)]
+    private string DecoratedMethodCaptureEnabled()
+    {
+        return "DecoratedMethod Enabled";
+    }
 }
diff --git a/libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/TracingAttributeTest.cs b/libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/TracingAttributeTest.cs
@@ -429,6 +429,69 @@ public void OnSuccess_WhenTracerCaptureModeIsDisabled_DoesNotCaptureResponse()
             Assert.Single(segment.Subsegments);
             Assert.False(subSegment.IsMetadataAdded); // does not add metadata
         }
+        
+        [Fact]
+        public void OnSuccess_WhenTracerCaptureResponseEnvironmentVariableIsFalse_ButDecoratorCapturesResponse()
+        {
+            // Arrange
+            Environment.SetEnvironmentVariable("LAMBDA_TASK_ROOT", "AWS");
+            Environment.SetEnvironmentVariable("POWERTOOLS_SERVICE_NAME", "POWERTOOLS");
+            Environment.SetEnvironmentVariable("POWERTOOLS_TRACER_CAPTURE_RESPONSE", "false");
+            
+            // Act
+            var segment = AWSXRayRecorder.Instance.TraceContext.GetEntity();
+            _handler.DecoratedHandlerCaptureResponse();
+            var subSegment = segment.Subsegments[0];
+
+            // Assert
+            Assert.True(segment.IsSubsegmentsAdded);
+            Assert.Single(segment.Subsegments);
+            Assert.True(subSegment.IsMetadataAdded);
+            Assert.True(subSegment.Metadata.ContainsKey("POWERTOOLS"));
+
+            var metadata = subSegment.Metadata["POWERTOOLS"];
+            Assert.Equal("DecoratedHandlerCaptureResponse response", metadata.Keys.Cast<string>().First());
+            var handlerResponse = metadata.Values.Cast<string>().First();
+            Assert.Equal("Hello World", handlerResponse);
+
+            var decoratedMethodSegmentDisabled = subSegment.Subsegments[0];
+            Assert.False(decoratedMethodSegmentDisabled.IsMetadataAdded);
+            Assert.Equal("## DecoratedMethodCaptureDisabled", decoratedMethodSegmentDisabled.Name);
+            
+            var decoratedMethodSegmentEnabled = decoratedMethodSegmentDisabled.Subsegments[0];
+            Assert.True(decoratedMethodSegmentEnabled.IsMetadataAdded);
+            
+            var decoratedMethodSegmentEnabledMetadata = decoratedMethodSegmentEnabled.Metadata["POWERTOOLS"];
+            var decoratedMethodSegmentEnabledResponse = decoratedMethodSegmentEnabledMetadata.Values.Cast<string>().First();
+            Assert.Equal("DecoratedMethod Enabled", decoratedMethodSegmentEnabledResponse);
+            Assert.Equal("## DecoratedMethodCaptureEnabled", decoratedMethodSegmentEnabled.Name);
+        }
+        
+        [Fact]
+        public void OnSuccess_WhenTracerCaptureResponseEnvironmentVariableIsTrue_ButDecoratorCapturesResponseDisabled()
+        {
+            // Arrange
+            Environment.SetEnvironmentVariable("LAMBDA_TASK_ROOT", "AWS");
+            Environment.SetEnvironmentVariable("POWERTOOLS_SERVICE_NAME", "POWERTOOLS");
+            Environment.SetEnvironmentVariable("POWERTOOLS_TRACER_CAPTURE_RESPONSE", "true");
+            
+            // Act
+            var segment = AWSXRayRecorder.Instance.TraceContext.GetEntity();
+            _handler.DecoratedMethodCaptureDisabled();
+            var subSegment = segment.Subsegments[0];
+
+            // Assert
+            Assert.True(segment.IsSubsegmentsAdded);
+            Assert.Single(segment.Subsegments);
+            Assert.False(subSegment.IsMetadataAdded);
+
+            var decoratedMethodSegmentEnabled = subSegment.Subsegments[0];
+            var metadata = decoratedMethodSegmentEnabled.Metadata["POWERTOOLS"];
+            Assert.True(decoratedMethodSegmentEnabled.IsMetadataAdded);
+            var decoratedMethodSegmentEnabledResponse = metadata.Values.Cast<string>().First();
+            Assert.Equal("DecoratedMethod Enabled", decoratedMethodSegmentEnabledResponse);
+            Assert.Equal("## DecoratedMethodCaptureEnabled", decoratedMethodSegmentEnabled.Name);
+        }
 
         #endregion
 
diff --git a/libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/XRayRecorderTests.cs b/libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/XRayRecorderTests.cs
