Merge pull request #1011 from aws-powertools/feature/tracing-sanitize-metadata

chore: Tracing sanitize metadata

Author: hjgraca

docs/core/tracing.md

@@ -196,9 +196,76 @@ context for an operation using any native object.
 
 ## Utilities
 
-Tracing modules comes with certain utility method when you don't want to use attribute for capturing a code block
+Tracing modules comes with certain utility methods when you don't want to use attribute for capturing a code block
 under a subsegment, or you are doing multithreaded programming. Refer examples below.
 
+### Using Statement Pattern
+
+You can create subsegments using the familiar `using` statement pattern for automatic cleanup and exception safety.
+
+=== "Basic Using Statement"
+
+    ```c# hl_lines="8 9 10 11 12 13"
+    using AWS.Lambda.Powertools.Tracing;
+
+    public class Function
+    {
+        public async Task<APIGatewayProxyResponse> FunctionHandler
+            (APIGatewayProxyRequest apigProxyEvent, ILambdaContext context)
+        {
+            using var gatewaySegment = Tracing.BeginSubsegment("PaymentGatewayIntegration");
+            gatewaySegment.AddAnnotation("Operation", "ProcessPayment");
+            gatewaySegment.AddAnnotation("PaymentMethod", "CreditCard");
+
+            var result = await ProcessPaymentAsync();
+            gatewaySegment.AddAnnotation("ProcessingTimeMs", result.ProcessingTimeMs);
+            // Subsegment automatically ends when disposed
+        }
+    }
+    ```
+
+=== "With Custom Namespace"
+
+    ```c# hl_lines="8 9 10"
+    using AWS.Lambda.Powertools.Tracing;
+
+    public class Function
+    {
+        public async Task<APIGatewayProxyResponse> FunctionHandler
+            (APIGatewayProxyRequest apigProxyEvent, ILambdaContext context)
+        {
+            using var segment = Tracing.BeginSubsegment("MyCustomNamespace", "DatabaseOperation");
+            segment.AddAnnotation("TableName", "Users");
+            segment.AddMetadata("query", "SELECT * FROM Users WHERE Active = 1");
+        }
+    }
+    ```
+
+=== "Nested Subsegments"
+
+    ```c# hl_lines="8 9 10 11 12 13 14 15 16"
+    using AWS.Lambda.Powertools.Tracing;
+
+    public class Function
+    {
+        public async Task<APIGatewayProxyResponse> FunctionHandler
+            (APIGatewayProxyRequest apigProxyEvent, ILambdaContext context)
+        {
+            using var outerSegment = Tracing.BeginSubsegment("PaymentProcessing");
+            outerSegment.AddAnnotation("Operation", "ProcessPayment");
+
+            var result = await ProcessPaymentAsync();
+
+            using var postProcessingSegment = Tracing.BeginSubsegment("PaymentPostProcessing");
+            postProcessingSegment.AddAnnotation("PaymentId", result.PaymentId);
+
+            await PostProcessPaymentAsync(result);
+        }
+    }
+    ```
+
+### Callback Pattern
+
 === "Functional Api"
 
     ```c# hl_lines="8 9 10 12 13 14"
@@ -244,6 +311,31 @@ under a subsegment, or you are doing multithreaded programming. Refer examples b
     }
     ```
 
+### Subsegment Methods
+
+When using the `using` statement pattern, the returned `TracingSubsegment` object provides direct access to tracing methods:
+
+=== "Available Methods"
+
+    ```c# hl_lines="8 9 10 11 12 13 14 15 16"
+    using var segment = Tracing.BeginSubsegment("PaymentProcessing");
+
+    // Add annotations (indexed by X-Ray)
+    segment.AddAnnotation("PaymentMethod", "CreditCard");
+    segment.AddAnnotation("Amount", 99.99);
+
+    // Add metadata (not indexed, for additional context)
+    segment.AddMetadata("PaymentDetails", paymentObject);
+    segment.AddMetadata("CustomNamespace", "RequestId", requestId);
+
+    // Add exception information
+    segment.AddException(exception);
+
+    // Add HTTP information
+    segment.AddHttpInformation("response_code", 200);
+    segment.AddHttpInformation("url", "https://api.payment.com/process");
+    ```
+
 ## 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.

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

@@ -1,4 +1,6 @@
+using System;
 using Amazon.XRay.Recorder.Core.Internal.Entities;
+using AWS.Lambda.Powertools.Common;
 
 namespace AWS.Lambda.Powertools.Tracing.Internal;
 
@@ -7,11 +9,106 @@ namespace AWS.Lambda.Powertools.Tracing.Internal;
 ///     It's a wrapper for Subsegment from Amazon.XRay.Recorder.Core.Internal.
 /// </summary>
 /// <seealso cref="Subsegment" />
-public class TracingSubsegment : Subsegment
+public class TracingSubsegment : Subsegment, IDisposable
 {
+    private bool _disposed = false;
+    private readonly bool _shouldAutoEnd;
+
     /// <summary>
     /// Wrapper constructor
     /// </summary>
     /// <param name="name"></param> 
-    public TracingSubsegment(string name) : base(name) { }
+    public TracingSubsegment(string name) : base(name) 
+    { 
+        _shouldAutoEnd = false;
+    }
+
+    /// <summary>
+    /// Constructor for disposable subsegments
+    /// </summary>
+    /// <param name="name">The name of the subsegment</param>
+    /// <param name="shouldAutoEnd">Whether this subsegment should auto-end when disposed</param>
+    internal TracingSubsegment(string name, bool shouldAutoEnd) : base(name) 
+    { 
+        _shouldAutoEnd = shouldAutoEnd;
+    }
+
+    /// <summary>
+    /// Adds an annotation to the subsegment
+    /// </summary>
+    /// <param name="key">The annotation key</param>
+    /// <param name="value">The annotation value</param>
+    public new void AddAnnotation(string key, object value)
+    {
+        XRayRecorder.Instance.AddAnnotation(key, value);
+    }
+
+    /// <summary>
+    /// Adds metadata to the subsegment
+    /// </summary>
+    /// <param name="key">The metadata key</param>
+    /// <param name="value">The metadata value</param>
+    public new void AddMetadata(string key, object value)
+    {
+        XRayRecorder.Instance.AddMetadata(Namespace ?? PowertoolsConfigurations.Instance.Service, key, value);
+    }
+
+    /// <summary>
+    /// Adds metadata to the subsegment with a specific namespace
+    /// </summary>
+    /// <param name="nameSpace">The namespace</param>
+    /// <param name="key">The metadata key</param>
+    /// <param name="value">The metadata value</param>
+    public new void AddMetadata(string nameSpace, string key, object value)
+    {
+        XRayRecorder.Instance.AddMetadata(nameSpace, key, value);
+    }
+
+    /// <summary>
+    /// Adds an exception to the subsegment
+    /// </summary>
+    /// <param name="exception">The exception to add</param>
+    public new void AddException(Exception exception)
+    {
+        XRayRecorder.Instance.AddException(exception);
+    }
+
+    /// <summary>
+    /// Adds HTTP information to the subsegment
+    /// </summary>
+    /// <param name="key">The HTTP information key</param>
+    /// <param name="value">The HTTP information value</param>
+    public void AddHttpInformation(string key, object value)
+    {
+        XRayRecorder.Instance.AddHttpInformation(key, value);
+    }
+
+    /// <summary>
+    /// Disposes the subsegment and ends it if configured to do so
+    /// </summary>
+    public void Dispose()
+    {
+        Dispose(true);
+        GC.SuppressFinalize(this);
+    }
+
+    /// <summary>
+    /// Protected dispose method
+    /// </summary>
+    /// <param name="disposing">Whether we're disposing</param>
+    protected virtual void Dispose(bool disposing)
+    {
+        if (!_disposed && disposing && _shouldAutoEnd)
+        {
+            try
+            {
+                XRayRecorder.Instance.EndSubsegment();
+            }
+            catch
+            {
+                // Swallow exceptions during disposal to prevent issues in using blocks
+            }
+            _disposed = true;
+        }
+    }
 }