chore: Tracing Fix async Task Handler exception

docs/core/tracing.md

@@ -244,7 +244,7 @@ under a subsegment, or you are doing multithreaded programming. Refer examples b
     }
     ```
 
-## Instrumenting SDK clients and HTTP calls
+## Instrumenting SDK clients
 
 You should make sure to instrument the SDK clients explicitly based on the function dependency. You can instrument all of your AWS SDK for .NET clients by calling RegisterForAllServices before you create them.
 
@@ -277,25 +277,98 @@ To instrument clients for some services and not others, call Register instead of
 Tracing.Register<IAmazonDynamoDB>()
 ```
 
-This functionality is a thin wrapper for AWS X-Ray .NET SDK. Refer details on [how to instrument SDK client with Xray](https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-dotnet-sdkclients.html) and [outgoing http calls](https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-dotnet-httpclients.html).
+This functionality is a thin wrapper for AWS X-Ray .NET SDK. Refer details on [how to instrument SDK client with Xray](https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-dotnet-sdkclients.html)
+
+## Instrumenting outgoing HTTP calls 
+
+=== "Function.cs"
+
+    ```c# hl_lines="7"
+    using Amazon.XRay.Recorder.Handlers.System.Net;
+
+    public class Function
+    {
+        public Function()
+        {
+            var httpClient = new HttpClient(new HttpClientXRayTracingHandler(new HttpClientHandler()));
+            var myIp = await httpClient.GetStringAsync("https://checkip.amazonaws.com/");
+        }
+    }
+    ```
+
+More information about instrumenting [outgoing http calls](https://docs.aws.amazon.com/xray/latest/devguide/xray-sdk-dotnet-httpclients.html).
 
 ## AOT Support
 
 Native AOT trims your application code as part of the compilation to ensure that the binary is as small as possible. .NET 8 for Lambda provides improved trimming support compared to previous versions of .NET.
 
-These improvements offer the potential to eliminate build-time trimming warnings, but .NET will never be completely trim safe. This means that parts of libraries that your function relies on may be trimmed out as part of the compilation step. You can manage this by defining TrimmerRootAssemblies as part of your `.csproj` file as shown in the following example.
 
-For the Tracing utility to work correctly and without trim warnings please add the following to your `.csproj` file
+### WithTracing()
 
-```xml
-<ItemGroup>
-    <TrimmerRootAssembly Include="AWSSDK.Core" />
-    <TrimmerRootAssembly Include="AWSXRayRecorder.Core" />
-    <TrimmerRootAssembly Include="AWSXRayRecorder.Handlers.AwsSdk" />
-    <TrimmerRootAssembly Include="Amazon.Lambda.APIGatewayEvents" />
-</ItemGroup>
-```
+To use Tracing utility with AOT support you first need to add `WithTracing()` to the source generator you are using either the default `SourceGeneratorLambdaJsonSerializer`
+or the Powertools Logging utility [source generator](logging.md#aot-support){:target="_blank"} `PowertoolsSourceGeneratorSerializer`.
+
+Examples:
+
+=== "Without Powertools Logging"
+
+    ```c# hl_lines="8"
+    using AWS.Lambda.Powertools.Tracing;
+    using AWS.Lambda.Powertools.Tracing.Serializers;
+
+    private static async Task Main()
+    {
+        Func<string, ILambdaContext, string> handler = FunctionHandler;
+        await LambdaBootstrapBuilder.Create(handler, new SourceGeneratorLambdaJsonSerializer<LambdaFunctionJsonSerializerContext>()
+        .WithTracing())
+            .Build()
+            .RunAsync();
+    }
+    ```
+
+=== "With Powertools Logging"
+
+    ```c# hl_lines="10 11"
+    using AWS.Lambda.Powertools.Logging;
+    using AWS.Lambda.Powertools.Logging.Serializers;
+    using AWS.Lambda.Powertools.Tracing;
+    using AWS.Lambda.Powertools.Tracing.Serializers;
+
+    private static async Task Main()
+    {
+        Func<string, ILambdaContext, string> handler = FunctionHandler;
+        await LambdaBootstrapBuilder.Create(handler, 
+            new PowertoolsSourceGeneratorSerializer<LambdaFunctionJsonSerializerContext>()
+            .WithTracing())
+                .Build()
+                .RunAsync();
+    }
+    ```
+
+### Publishing
+
+!!! warning "Publishing"
+    Make sure you are publishing your code with `--self-contained true` and that you have `<TrimMode>partial</TrimMode>` in your `.csproj` file 
+
+### Trimming
+
+!!! warning "Trim warnings"
+    ```xml
+    <ItemGroup>
+        <TrimmerRootAssembly Include="AWSSDK.Core" />
+        <TrimmerRootAssembly Include="AWSXRayRecorder.Core" />
+        <TrimmerRootAssembly Include="AWSXRayRecorder.Handlers.AwsSdk" />
+        <TrimmerRootAssembly Include="Amazon.Lambda.APIGatewayEvents" />
+        <TrimmerRootAssembly Include="bootstrap" />
+        <TrimmerRootAssembly Include="Shared" />
+    </ItemGroup>
+    ```
+
+    Note that when you receive a trim warning, adding the class that generates the warning to TrimmerRootAssembly might not resolve the issue. A trim warning indicates that the class is trying to access some other class that can't be determined until runtime. To avoid runtime errors, add this second class to TrimmerRootAssembly. 
+
+    To learn more about managing trim warnings, see [Introduction to trim warnings](https://learn.microsoft.com/en-us/dotnet/core/deploying/trimming/fixing-warnings) in the Microsoft .NET documentation.
 
-Note that when you receive a trim warning, adding the class that generates the warning to TrimmerRootAssembly might not resolve the issue. A trim warning indicates that the class is trying to access some other class that can't be determined until runtime. To avoid runtime errors, add this second class to TrimmerRootAssembly.
+### Not supported
 
-To learn more about managing trim warnings, see [Introduction to trim warnings](https://learn.microsoft.com/en-us/dotnet/core/deploying/trimming/fixing-warnings) in the Microsoft .NET documentation.
+!!! warning "Not supported"
+    Currently instrumenting SDK clients with `Tracing.RegisterForAllServices()` is not supported on AOT mode.

libraries/src/AWS.Lambda.Powertools.Tracing/AWS.Lambda.Powertools.Tracing.csproj

@@ -16,6 +16,7 @@
         <PackageReference Include="AWSSDK.XRay" />
         <PackageReference Include="AWSXRayRecorder.Core" />
         <PackageReference Include="AWSXRayRecorder.Handlers.AwsSdk" />
+        <PackageReference Include="Amazon.Lambda.Serialization.SystemTextJson" />
         <ProjectReference Include="..\AWS.Lambda.Powertools.Common\AWS.Lambda.Powertools.Common.csproj" Condition="'$(Configuration)'=='Debug'"/>
     </ItemGroup>
 

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

@@ -0,0 +1,259 @@
+/*
+ * 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 System.Threading.Tasks;
+using AspectInjector.Broker;
+using AWS.Lambda.Powertools.Common;
+using AWS.Lambda.Powertools.Common.Utils;
+
+namespace AWS.Lambda.Powertools.Tracing.Internal;
+
+/// <summary>
+///     Tracing Aspect
+///     Scope.Global is singleton
+/// </summary>
+[Aspect(Scope.Global, Factory = typeof(TracingAspectFactory))]
+public class TracingAspect
+{
+    /// <summary>
+    ///     The Powertools for AWS Lambda (.NET) configurations
+    /// </summary>
+    private readonly IPowertoolsConfigurations _powertoolsConfigurations;
+
+    /// <summary>
+    ///     X-Ray Recorder
+    /// </summary>
+    private readonly IXRayRecorder _xRayRecorder;
+
+    /// <summary>
+    ///     If true, then is cold start
+    /// </summary>
+    private static bool _isColdStart = true;
+
+    /// <summary>
+    ///     If true, capture annotations
+    /// </summary>
+    private static bool _captureAnnotations = true;
+
+    /// <summary>
+    ///     If true, annotations have been captured
+    /// </summary>
+    private bool _isAnnotationsCaptured;
+
+    /// <summary>
+    /// Aspect constructor
+    /// </summary>
+    /// <param name="powertoolsConfigurations"></param>
+    /// <param name="xRayRecorder"></param>
+    public TracingAspect(IPowertoolsConfigurations powertoolsConfigurations, IXRayRecorder xRayRecorder)
+    {
+        _powertoolsConfigurations = powertoolsConfigurations;
+        _xRayRecorder = xRayRecorder;
+    }
+
+    /// <summary>
+    /// Surrounds the specific method with Tracing aspect
+    /// </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)
+    {
+        var trigger = triggers.OfType<TracingAttribute>().First();
+
+        if (TracingDisabled())
+            return target(args);
+
+        var @namespace = !string.IsNullOrWhiteSpace(trigger.Namespace)
+            ? trigger.Namespace
+            : _powertoolsConfigurations.Service;
+
+        var (segmentName, metadataName) = string.IsNullOrWhiteSpace(trigger.SegmentName)
+            ? ($"## {name}", name)
+            : (trigger.SegmentName, trigger.SegmentName);
+
+        BeginSegment(segmentName, @namespace);
+
+        try
+        {
+            var result = target(args);
+
+            if (result is Task task)
+            {
+                if (task.IsFaulted)  
+                {  
+                    // Force the exception to be thrown  
+                    task.Exception?.Handle(ex => false);    
+                } 
+
+                // Only handle response if it's not a void Task
+                if (task.GetType().IsGenericType)
+                {
+                    var taskResult = task.GetType().GetProperty("Result")?.GetValue(task);
+                    HandleResponse(metadataName, taskResult, trigger.CaptureMode, @namespace);
+                }
+                _xRayRecorder.EndSubsegment();
+                return task;
+            }
+
+            HandleResponse(metadataName, result, trigger.CaptureMode, @namespace);
+
+            _xRayRecorder.EndSubsegment();
+            return result;
+        }
+        catch (Exception ex)
+        {
+            var actualException = ex is AggregateException ae ? ae.InnerException! : ex;  
+            HandleException(actualException, metadataName, trigger.CaptureMode, @namespace);  
+            _xRayRecorder.EndSubsegment();  
+
+            // Capture and rethrow the original exception preserving the stack trace
+            ExceptionDispatchInfo.Capture(actualException).Throw();  
+            throw;
+        }
+        finally
+        {
+            if (_isAnnotationsCaptured)
+                _captureAnnotations = true;
+        }
+    }
+
+    private void BeginSegment(string segmentName, string @namespace)
+    {
+        _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;
+    }
+
+    private void HandleResponse(string name, object result, TracingCaptureMode captureMode, string @namespace)
+    {
+        if (!CaptureResponse(captureMode)) return;
+        if (result == null) return;  // Don't try to serialize null results
+
+        // Skip if the result is VoidTaskResult
+        if (result.GetType().Name == "VoidTaskResult") return;
+
+#if NET8_0_OR_GREATER
+        if (!RuntimeFeatureWrapper.IsDynamicCodeSupported) // is AOT
+        {
+            _xRayRecorder.AddMetadata(
+                @namespace,
+                $"{name} response",
+                Serializers.PowertoolsTracingSerializer.Serialize(result)
+            );
+            return;
+        }
+#endif
+
+        _xRayRecorder.AddMetadata(
+            @namespace,
+            $"{name} response",
+            result
+        );
+    }
+
+    private void HandleException(Exception exception, string name, TracingCaptureMode captureMode, string @namespace)
+    {
+        if (!CaptureError(captureMode)) return;
+
+        var sb = new StringBuilder();
+        sb.AppendLine($"Exception type: {exception.GetType()}");
+        sb.AppendLine($"Exception message: {exception.Message}");
+        sb.AppendLine($"Stack trace: {exception.StackTrace}");
+
+        if (exception.InnerException != null)
+        {
+            sb.AppendLine("---BEGIN InnerException--- ");
+            sb.AppendLine($"Exception type {exception.InnerException.GetType()}");
+            sb.AppendLine($"Exception message: {exception.InnerException.Message}");
+            sb.AppendLine($"Stack trace: {exception.InnerException.StackTrace}");
+            sb.AppendLine("---END Inner Exception");
+        }
+
+        _xRayRecorder.AddMetadata(
+            @namespace,
+            $"{name} error",
+            sb.ToString()
+        );
+    }
+
+    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;
+    }
+
+    private bool CaptureResponse(TracingCaptureMode captureMode)
+    {
+        return captureMode switch
+        {
+            TracingCaptureMode.EnvironmentVariable => _powertoolsConfigurations.TracerCaptureResponse,
+            TracingCaptureMode.Response => true,
+            TracingCaptureMode.ResponseAndError => true,
+            _ => false
+        };
+    }
+
+    private bool CaptureError(TracingCaptureMode captureMode)
+    {
+        return captureMode switch
+        {
+            TracingCaptureMode.EnvironmentVariable => _powertoolsConfigurations.TracerCaptureError,
+            TracingCaptureMode.Error => true,
+            TracingCaptureMode.ResponseAndError => true,
+            _ => false
+        };
+    }
+
+    internal static void ResetForTest()
+    {
+        _isColdStart = true;
+        _captureAnnotations = true;
+    }
+}