feat(logging): add ExtraKeys to logger (#1128)

Author: hjgraca

docs/core/logging.md

@@ -604,10 +604,125 @@ You can remove any additional key from entry using `Logger.RemoveKeys()`.
     }
     ```
 
-## Extra Keys
+### Temporary keys with ExtraKeys
 
-Extra keys allow you to append additional keys to a log entry. Unlike `AppendKey`, extra keys will only apply to the
-current log entry.
+The `ExtraKeys` method allows temporary modification of the Logger's context without manual cleanup. It's useful for adding context keys to specific workflows while maintaining the logger's overall state.
+
+Keys are automatically removed when the scope ends, eliminating the need to manually call `AppendKey` and `RemoveKeys`.
+
+=== "Using Dictionary"
+
+    ```c# hl_lines="12-16"
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class Function
+    {
+        [Logging]
+        public async Task<APIGatewayProxyResponse> FunctionHandler
+            (APIGatewayProxyRequest apigProxyEvent, ILambdaContext context)
+        {
+            var orderId = apigProxyEvent.PathParameters["orderId"];
+            
+            using (Logger.ExtraKeys(new Dictionary<string, object> { { "orderId", orderId } }))
+            {
+                Logger.LogInformation("Processing order");
+                await ProcessOrderAsync(orderId);
+                Logger.LogInformation("Order processed"); // orderId included
+            }
+            // orderId is automatically removed
+            
+            Logger.LogInformation("Continuing without orderId");
+        }
+    }
+    ```
+
+=== "Using Tuples"
+
+    ```c# hl_lines="12-16"
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class Function
+    {
+        [Logging]
+        public async Task<APIGatewayProxyResponse> FunctionHandler
+            (APIGatewayProxyRequest apigProxyEvent, ILambdaContext context)
+        {
+            var orderId = apigProxyEvent.PathParameters["orderId"];
+            
+            using (Logger.ExtraKeys(("orderId", orderId), ("customerId", "customer-123")))
+            {
+                Logger.LogInformation("Processing order");
+                await ProcessOrderAsync(orderId);
+                Logger.LogInformation("Order processed"); // orderId and customerId included
+            }
+            // Both keys are automatically removed
+        }
+    }
+    ```
+
+=== "Nested Scopes"
+
+    ```c# hl_lines="10-19"
+    /**
+     * Handler for requests to Lambda function.
+     */
+    public class Function
+    {
+        [Logging]
+        public async Task<APIGatewayProxyResponse> FunctionHandler
+            (APIGatewayProxyRequest apigProxyEvent, ILambdaContext context)
+        {
+            using (Logger.ExtraKeys(("requestId", context.AwsRequestId)))
+            {
+                Logger.LogInformation("Starting request"); // requestId included
+                
+                using (Logger.ExtraKeys(("step", "validation")))
+                {
+                    Logger.LogInformation("Validating"); // requestId AND step included
+                }
+                // step removed, requestId still present
+                
+                Logger.LogInformation("Request complete"); // only requestId
+            }
+        }
+    }
+    ```
+
+=== "Example CloudWatch Logs excerpt"
+
+    ```json hl_lines="14 15"
+    {
+        "level": "Information",
+        "message": "Processing order",
+        "timestamp": "2024-01-15T10:30:00.0000000Z",
+        "service": "order-service",
+        "cold_start": true,
+        "function_name": "OrderProcessor",
+        "function_memory_size": 256,
+        "function_arn": "arn:aws:lambda:eu-west-1:123456789:function:OrderProcessor",
+        "function_request_id": "abc-123-def",
+        "function_version": "$LATEST",
+        "xray_trace_id": "1-abc-123",
+        "name": "AWS.Lambda.Powertools.Logging.Logger",
+        "order_id": "order-456",
+        "customer_id": "customer-123"
+    }
+    ```
+
+!!! tip "When to use ExtraKeys vs AppendKey"
+    Use `ExtraKeys` when you need keys for a specific operation or code block. Use `AppendKey` when keys should persist for the entire Lambda invocation.
+
+!!! warning "Key overwrite behavior"
+    If a key already exists when entering an `ExtraKeys` scope, it will be overwritten and then **removed** when the scope ends. The original value is not restored. Use unique key names within `ExtraKeys` scopes to avoid unexpected behavior.
+
+!!! info "Async safe"
+    `ExtraKeys` is safe to use across `async/await` boundaries. Keys will correctly flow through asynchronous operations within the same execution context.
+
+## Extra Keys (Single Log Entry)
+
+Extra keys can also be added to a single log entry using message templates. Unlike `AppendKey` or `ExtraKeys()`, these keys will only apply to the current log entry.
 
 Extra keys argument is available for all log levels' methods, as implemented in the standard logging library - e.g.
 Logger.Information, Logger.Warning.

libraries/src/AWS.Lambda.Powertools.Logging/Logger.Scope.cs

@@ -11,24 +11,63 @@ namespace AWS.Lambda.Powertools.Logging;
 public static partial class Logger
 {
     /// <summary>
-    ///     Thread-safe dictionary for per-thread scope storage.
-    ///     Uses ManagedThreadId as key to ensure isolation when Lambda processes
-    ///     multiple concurrent requests (AWS_LAMBDA_MAX_CONCURRENCY > 1).
-    ///     Inner dictionary is ConcurrentDictionary for thread-safe operations.
+    ///     AsyncLocal storage for per-invocation scope.
+    ///     Uses AsyncLocal to ensure keys flow correctly across async/await boundaries.
+    ///     
+    ///     In Lambda's multi-threaded mode, each invocation starts with a fresh execution
+    ///     context, providing natural isolation between concurrent invocations.
+    ///     Keys added via AppendKey flow across async/await within the same invocation.
     /// </summary>
-    private static readonly ConcurrentDictionary<int, ConcurrentDictionary<string, object>> _threadScopes = new();
+    private static readonly AsyncLocal<ConcurrentDictionary<string, object>> _asyncScope = new();
 
     /// <summary>
-    ///     Gets the scope for the current thread.
-    ///     Creates a new dictionary if one doesn't exist for this thread.
+    ///     Gets the scope for the current async execution context.
+    ///     Creates a new dictionary if one doesn't exist for this context.
     /// </summary>
     /// <value>The scope.</value>
     private static ConcurrentDictionary<string, object> Scope
     {
         get
         {
-            var threadId = Environment.CurrentManagedThreadId;
-            return _threadScopes.GetOrAdd(threadId, _ => new ConcurrentDictionary<string, object>(StringComparer.Ordinal));
+            var scope = _asyncScope.Value;
+            if (scope == null)
+            {
+                scope = new ConcurrentDictionary<string, object>(StringComparer.Ordinal);
+                _asyncScope.Value = scope;
+            }
+            return scope;
+        }
+    }
+    
+    /// <summary>
+    ///     Creates a new isolated scope for the current execution context.
+    ///     Used internally to simulate Lambda invocation isolation in tests.
+    /// </summary>
+    /// <returns>An IDisposable that restores the previous scope when disposed.</returns>
+    internal static IDisposable UseScope()
+    {
+        return new LoggerScopeContext();
+    }
+
+    /// <summary>
+    ///     Manages a new isolated scope context.
+    /// </summary>
+    private sealed class LoggerScopeContext : IDisposable
+    {
+        private readonly ConcurrentDictionary<string, object> _previousScope;
+        private bool _disposed;
+
+        public LoggerScopeContext()
+        {
+            _previousScope = _asyncScope.Value;
+            _asyncScope.Value = new ConcurrentDictionary<string, object>(StringComparer.Ordinal);
+        }
+
+        public void Dispose()
+        {
+            if (_disposed) return;
+            _disposed = true;
+            _asyncScope.Value = _previousScope;
         }
     }
 
@@ -49,7 +88,8 @@ public static string CorrelationId
     }
 
     /// <summary>
-    ///     Appending additional key to the log context.
+    ///     Appends a key-value pair to the log context.
+    ///     Keys persist across async/await boundaries within the same execution context.
     /// </summary>
     /// <param name="key">The key.</param>
     /// <param name="value">The value.</param>
@@ -66,7 +106,7 @@ public static void AppendKey(string key, object value)
     }
 
     /// <summary>
-    ///     Appending additional key to the log context.
+    ///     Appends multiple keys to the log context.
     /// </summary>
     /// <param name="keys">The list of keys.</param>
     public static void AppendKeys(IEnumerable<KeyValuePair<string, object>> keys)
@@ -76,7 +116,7 @@ public static void AppendKeys(IEnumerable<KeyValuePair<string, object>> keys)
     }
 
     /// <summary>
-    ///     Appending additional key to the log context.
+    ///     Appends multiple keys to the log context.
     /// </summary>
     /// <param name="keys">The list of keys.</param>
     public static void AppendKeys(IEnumerable<KeyValuePair<string, string>> keys)
@@ -86,7 +126,7 @@ public static void AppendKeys(IEnumerable<KeyValuePair<string, string>> keys)
     }
 
     /// <summary>
-    ///     Remove additional keys from the log context.
+    ///     Removes keys from the log context.
     /// </summary>
     /// <param name="keys">The list of keys.</param>
     public static void RemoveKeys(params string[] keys)
@@ -97,7 +137,7 @@ public static void RemoveKeys(params string[] keys)
     }
 
     /// <summary>
-    ///     Returns all additional keys added to the log context.
+    ///     Returns all keys added to the log context.
     ///     Returns a snapshot to ensure thread-safety during enumeration.
     /// </summary>
     /// <returns>IEnumerable&lt;KeyValuePair&lt;System.String, System.Object&gt;&gt;.</returns>
@@ -108,15 +148,11 @@ public static IEnumerable<KeyValuePair<string, object>> GetAllKeys()
     }
 
     /// <summary>
-    ///     Removes all additional keys from the log context for the current thread.
+    ///     Removes all keys from the log context for the current execution context.
     /// </summary>
     internal static void RemoveAllKeys()
     {
-        var threadId = Environment.CurrentManagedThreadId;
-        if (_threadScopes.TryGetValue(threadId, out var scope))
-        {
-            scope.Clear();
-        }
+        _asyncScope.Value?.Clear();
     }
 
     /// <summary>
@@ -126,4 +162,95 @@ public static void RemoveKey(string key)
     {
         Scope.TryRemove(key, out _);
     }
+
+    /// <summary>
+    ///     Adds temporary keys to the log context that are automatically removed when disposed.
+    ///     Safe to use across async/await boundaries.
+    /// </summary>
+    /// <param name="keys">The keys to add temporarily.</param>
+    /// <returns>An IDisposable that removes the keys when disposed.</returns>
+    /// <remarks>
+    ///     <para>
+    ///         <b>Important:</b> If a key already exists in the context, it will be overwritten
+    ///         and then removed when the scope is disposed. The original value is NOT restored.
+    ///     </para>
+    ///     <para>
+    ///         For example, if "orderId" = "A" exists and you create a scope with "orderId" = "B",
+    ///         disposing the scope will remove "orderId" entirely, not restore it to "A".
+    ///     </para>
+    /// </remarks>
+    /// <example>
+    /// <code>
+    /// using (Logger.ExtraKeys(new Dictionary&lt;string, object&gt; { {"orderId", "123"} }))
+    /// {
+    ///     await ProcessOrderAsync();
+    ///     Logger.LogInformation("Order processed"); // includes orderId
+    /// }
+    /// // orderId is automatically removed
+    /// </code>
+    /// </example>
+    public static IDisposable ExtraKeys(IEnumerable<KeyValuePair<string, object>> keys)
+    {
+        if (keys == null) throw new ArgumentNullException(nameof(keys));
+        return new LoggerExtraKeysScope(keys);
+    }
+
+    /// <summary>
+    ///     Adds temporary keys to the log context that are automatically removed when disposed.
+    ///     Safe to use across async/await boundaries.
+    /// </summary>
+    /// <param name="keys">The keys to add temporarily as tuples.</param>
+    /// <returns>An IDisposable that removes the keys when disposed.</returns>
+    /// <remarks>
+    ///     <para>
+    ///         <b>Important:</b> If a key already exists in the context, it will be overwritten
+    ///         and then removed when the scope is disposed. The original value is NOT restored.
+    ///     </para>
+    /// </remarks>
+    /// <example>
+    /// <code>
+    /// using (Logger.ExtraKeys(("orderId", "123"), ("customerId", "456")))
+    /// {
+    ///     Logger.LogInformation("Processing"); // includes orderId and customerId
+    /// }
+    /// </code>
+    /// </example>
+    public static IDisposable ExtraKeys(params (string Key, object Value)[] keys)
+    {
+        if (keys == null) throw new ArgumentNullException(nameof(keys));
+        return new LoggerExtraKeysScope(keys.Select(k => new KeyValuePair<string, object>(k.Key, k.Value)));
+    }
+
+    /// <summary>
+    ///     Scope that manages temporary logging keys.
+    ///     Keys are added on construction and removed on disposal.
+    /// </summary>
+    private sealed class LoggerExtraKeysScope : IDisposable
+    {
+        private readonly string[] _keysToRemove;
+        private bool _disposed;
+
+        public LoggerExtraKeysScope(IEnumerable<KeyValuePair<string, object>> keys)
+        {
+            var keyList = new List<string>();
+            
+            foreach (var (key, value) in keys)
+            {
+                if (string.IsNullOrWhiteSpace(key)) continue;
+                
+                AppendKey(key, value);
+                keyList.Add(key);
+            }
+            
+            _keysToRemove = keyList.ToArray();
+        }
+
+        public void Dispose()
+        {
+            if (_disposed) return;
+            _disposed = true;
+            
+            RemoveKeys(_keysToRemove);
+        }
+    }
 }