Merge branch 'develop' into dependabot/github_actions/actions/setup-dotnet-4.3.1

Author: hjgraca

libraries/src/AWS.Lambda.Powertools.Common/Core/Constants.cs

@@ -20,6 +20,14 @@ namespace AWS.Lambda.Powertools.Common;
 /// </summary>
 internal static class Constants
 {
+    /// <summary>
+    ///     Constant for AWS_LAMBDA_INITIALIZATION_TYPE environment variable
+    ///     This is used to determine if the Lambda function is running in provisioned concurrency mode
+    ///     or not. If the value is "provisioned-concurrency", it indicates that the function is running in provisioned
+    ///     concurrency mode. Otherwise, it is running in standard mode.
+    /// </summary>
+    internal const string AWSInitializationTypeEnv = "AWS_LAMBDA_INITIALIZATION_TYPE";
+
     /// <summary>
     ///     Constant for POWERTOOLS_SERVICE_NAME environment variable
     /// </summary>

libraries/src/AWS.Lambda.Powertools.Common/Core/IPowertoolsConfigurations.cs

@@ -167,4 +167,15 @@ public interface IPowertoolsConfigurations
     /// Gets a value indicating whether Metrics are disabled.
     /// </summary>
     bool MetricsDisabled { get; }
+    
+    /// <summary>
+    /// Indicates if the current execution is a cold start.
+    /// </summary>
+    bool IsColdStart { get; }
+
+    /// <summary>
+    /// AWS Lambda initialization type.
+    /// This is set to "on-demand" for on-demand Lambda functions and "provisioned-concurrency" for provisioned concurrency.
+    /// </summary>
+    string AwsInitializationType { get; }
 }

libraries/src/AWS.Lambda.Powertools.Common/Core/ISystemWrapper.cs

@@ -59,15 +59,4 @@ public interface ISystemWrapper
     /// </summary>
     /// <param name="type"></param>
     void SetExecutionEnvironment<T>(T type);
-
-    /// <summary>
-    /// Sets console output
-    /// Useful for testing and checking the console output
-    /// <code>
-    /// var consoleOut = new StringWriter();
-    /// SystemWrapper.Instance.SetOut(consoleOut);
-    /// </code>
-    /// </summary>
-    /// <param name="writeTo">The TextWriter instance where to write to</param>
-    void SetOut(TextWriter writeTo);
 }

libraries/src/AWS.Lambda.Powertools.Common/Core/LambdaLifecycleTracker.cs

@@ -0,0 +1,66 @@
+using System;
+using System.Threading;
+
+namespace AWS.Lambda.Powertools.Common.Core;
+
+/// <summary>
+/// Tracks Lambda lifecycle state including cold starts
+/// </summary>
+internal static class LambdaLifecycleTracker
+{
+    // Static flag that's true only for the first Lambda container initialization
+    private static bool _isFirstContainer = true;
+        
+    // Store the cold start state for the current invocation
+    private static readonly AsyncLocal<bool?> CurrentInvocationColdStart = new AsyncLocal<bool?>();
+
+    private static string _lambdaInitType;
+    private static string LambdaInitType => _lambdaInitType ?? Environment.GetEnvironmentVariable(Constants.AWSInitializationTypeEnv);
+    
+    /// <summary>
+    /// Returns true if the current Lambda invocation is a cold start
+    /// </summary>
+    public static bool IsColdStart
+    {
+        get
+        {
+            if(LambdaInitType == "provisioned-concurrency")
+            {
+                // If the Lambda is provisioned concurrency, it is not a cold start
+                return false;
+            }
+            
+            // Initialize the cold start state for this invocation if not already set
+            if (!CurrentInvocationColdStart.Value.HasValue)
+            {
+                // Capture the container's cold start state for this entire invocation
+                CurrentInvocationColdStart.Value = _isFirstContainer;
+                    
+                // After detecting the first invocation, mark future ones as warm
+                if (_isFirstContainer)
+                {
+                    _isFirstContainer = false;
+                }
+            }
+                
+            // Return the cold start state for this invocation (cannot change during the invocation)
+            return CurrentInvocationColdStart.Value ?? false;
+        }
+    }
+    
+    
+
+    /// <summary>
+    /// Resets the cold start state for testing
+    /// </summary>
+    /// <param name="resetContainer">Whether to reset the container state (defaults to true)</param>
+    internal static void Reset(bool resetContainer = true)
+    {
+        if (resetContainer)
+        {
+            _isFirstContainer = true;
+        }
+        CurrentInvocationColdStart.Value = null;
+        _lambdaInitType = null;
+    }
+}

libraries/src/AWS.Lambda.Powertools.Common/Core/PowertoolsConfigurations.cs

@@ -14,6 +14,7 @@
  */
 
 using System.Globalization;
+using AWS.Lambda.Powertools.Common.Core;
 
 namespace AWS.Lambda.Powertools.Common;
 
@@ -222,4 +223,11 @@ public void SetExecutionEnvironment<T>(T type)
 
     /// <inheritdoc />
     public bool MetricsDisabled => GetEnvironmentVariableOrDefault(Constants.PowertoolsMetricsDisabledEnv, false);
+    
+    /// <inheritdoc />
+    public bool IsColdStart => LambdaLifecycleTracker.IsColdStart;
+    
+    /// <inheritdoc />
+    public string AwsInitializationType =>
+        GetEnvironmentVariable(Constants.AWSInitializationTypeEnv);
 }

libraries/src/AWS.Lambda.Powertools.Common/Core/SystemWrapper.cs

@@ -27,6 +27,9 @@ namespace AWS.Lambda.Powertools.Common;
 public class SystemWrapper : ISystemWrapper
 {
     private static IPowertoolsEnvironment _powertoolsEnvironment;
+    private static bool _inTestMode = false;
+    private static TextWriter _testOutputStream;
+    private static bool _outputResetPerformed = false;
 
     /// <summary>
     ///     The instance
@@ -41,13 +44,11 @@ public SystemWrapper(IPowertoolsEnvironment powertoolsEnvironment)
         _powertoolsEnvironment = powertoolsEnvironment;
         _instance ??= this;
 
-        // Clear AWS SDK Console injected parameters StdOut and StdErr
-        var standardOutput = new StreamWriter(Console.OpenStandardOutput());
-        standardOutput.AutoFlush = true;
-        Console.SetOut(standardOutput);
-        var errordOutput = new StreamWriter(Console.OpenStandardError());
-        errordOutput.AutoFlush = true;
-        Console.SetError(errordOutput);
+        if (!_inTestMode)
+        {
+            // Clear AWS SDK Console injected parameters in production only
+            ResetConsoleOutput();
+        }
     }
 
     /// <summary>
@@ -72,7 +73,15 @@ public string GetEnvironmentVariable(string variable)
     /// <param name="value">The value.</param>
     public void Log(string value)
     {
-        Console.Write(value);
+        if (_inTestMode && _testOutputStream != null)
+        {
+            _testOutputStream.Write(value);
+        }
+        else
+        {
+            EnsureConsoleOutputOnce();
+            Console.Write(value);
+        }
     }
 
     /// <summary>
@@ -81,7 +90,15 @@ public void Log(string value)
     /// <param name="value">The value.</param>
     public void LogLine(string value)
     {
-        Console.WriteLine(value);
+        if (_inTestMode && _testOutputStream != null)
+        {
+            _testOutputStream.WriteLine(value);
+        }
+        else
+        {
+            EnsureConsoleOutputOnce();
+            Console.WriteLine(value);
+        }
     }
 
     /// <summary>
@@ -126,9 +143,20 @@ public void SetExecutionEnvironment<T>(T type)
         SetEnvironmentVariable(envName, envValue.ToString());
     }
 
-    /// <inheritdoc />
-    public void SetOut(TextWriter writeTo)
+    /// <summary>
+    /// Sets console output
+    /// Useful for testing and checking the console output
+    /// <code>
+    /// var consoleOut = new StringWriter();
+    /// SystemWrapper.Instance.SetOut(consoleOut);
+    /// </code>
+    /// </summary>
+    /// <param name="writeTo">The TextWriter instance where to write to</param>
+
+    public static void SetOut(TextWriter writeTo)
     {
+        _testOutputStream = writeTo;
+        _inTestMode = true;
         Console.SetOut(writeTo);
     }
 
@@ -152,4 +180,33 @@ private string ParseAssemblyName(string assemblyName)
 
         return $"{Constants.FeatureContextIdentifier}/{assemblyName}";
     }
+    
+    private static void EnsureConsoleOutputOnce()
+    {
+        if (_outputResetPerformed) return;
+        ResetConsoleOutput();
+        _outputResetPerformed = true;
+    }
+    
+    private static void ResetConsoleOutput()
+    {
+        var standardOutput = new StreamWriter(Console.OpenStandardOutput());
+        standardOutput.AutoFlush = true;
+        Console.SetOut(standardOutput);
+        var errorOutput = new StreamWriter(Console.OpenStandardError());
+        errorOutput.AutoFlush = true;
+        Console.SetError(errorOutput);
+    }
+    
+    public static void ClearOutputResetFlag()
+    {
+        _outputResetPerformed = false;
+    }
+    
+    // For test cleanup
+    internal static void ResetTestMode()
+    {
+        _inTestMode = false;
+        _testOutputStream = null;
+    }
 }

libraries/src/AWS.Lambda.Powertools.Common/InternalsVisibleTo.cs

@@ -17,6 +17,7 @@
 
 [assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Logging")]
 [assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Metrics")]
+[assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Tracing")]
 [assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Idempotency")]
 [assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Common.Tests")]
 [assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Tracing.Tests")]

libraries/src/AWS.Lambda.Powertools.Logging/Internal/LoggingAspect.cs

@@ -21,6 +21,7 @@
 using System.Text.Json;
 using AspectInjector.Broker;
 using AWS.Lambda.Powertools.Common;
+using AWS.Lambda.Powertools.Common.Core;
 using AWS.Lambda.Powertools.Logging.Serializers;
 using Microsoft.Extensions.Logging;
 
@@ -34,11 +35,6 @@ namespace AWS.Lambda.Powertools.Logging.Internal;
 [Aspect(Scope.Global, Factory = typeof(LoggingAspectFactory))]
 public class LoggingAspect
 {
-    /// <summary>
-    ///     The is cold start
-    /// </summary>
-    private bool _isColdStart = true;
-
     /// <summary>
     ///     The initialize context
     /// </summary>
@@ -143,9 +139,8 @@ public void OnEntry(
             if (!_initializeContext)
                 return;
 
-            Logger.AppendKey(LoggingConstants.KeyColdStart, _isColdStart);
+            Logger.AppendKey(LoggingConstants.KeyColdStart, LambdaLifecycleTracker.IsColdStart);
 
-            _isColdStart = false;
             _initializeContext = false;
             _isContextInitialized = true;
 

libraries/src/AWS.Lambda.Powertools.Metrics/Internal/MetricsAspect.cs

@@ -20,6 +20,7 @@
 using Amazon.Lambda.Core;
 using AspectInjector.Broker;
 using AWS.Lambda.Powertools.Common;
+using AWS.Lambda.Powertools.Common.Core;
 
 namespace AWS.Lambda.Powertools.Metrics;
 
@@ -30,22 +31,12 @@ namespace AWS.Lambda.Powertools.Metrics;
 [Aspect(Scope.Global)]
 public class MetricsAspect
 {
-    /// <summary>
-    ///     The is cold start
-    /// </summary>
-    private static bool _isColdStart;
-
     /// <summary>
     ///     Gets the metrics instance.
     /// </summary>
     /// <value>The metrics instance.</value>
     private static IMetrics _metricsInstance;
 
-    static MetricsAspect()
-    {
-        _isColdStart = true;
-    }
-
     /// <summary>
     /// Runs before the execution of the method marked with the Metrics Attribute
     /// </summary>
@@ -89,10 +80,9 @@ public void Before(
             Triggers = triggers
         };
 
-        if (_isColdStart)
+        if (LambdaLifecycleTracker.IsColdStart)
         {
             _metricsInstance.CaptureColdStartMetric(GetContext(eventArgs));
-            _isColdStart = false;
         }
     }
 
@@ -112,7 +102,7 @@ public void Exit()
     internal static void ResetForTest()
     {
         _metricsInstance = null;
-        _isColdStart = true;
+        LambdaLifecycleTracker.Reset();
         Metrics.ResetForTest();
     }
 

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

@@ -20,6 +20,7 @@
 using System.Threading.Tasks;
 using AspectInjector.Broker;
 using AWS.Lambda.Powertools.Common;
+using AWS.Lambda.Powertools.Common.Core;
 using AWS.Lambda.Powertools.Common.Utils;
 
 namespace AWS.Lambda.Powertools.Tracing.Internal;
@@ -41,11 +42,6 @@ public class TracingAspect
     /// </summary>
     private readonly IXRayRecorder _xRayRecorder;
 
-    /// <summary>
-    ///     If true, then is cold start
-    /// </summary>
-    private static bool _isColdStart = true;
-
     /// <summary>
     ///     If true, capture annotations
     /// </summary>
@@ -148,16 +144,14 @@ private void BeginSegment(string segmentName, string @namespace)
 
         if (_captureAnnotations)
         {
-            _xRayRecorder.AddAnnotation("ColdStart", _isColdStart);
+            _xRayRecorder.AddAnnotation("ColdStart", LambdaLifecycleTracker.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)
@@ -253,7 +247,7 @@ private bool CaptureError(TracingCaptureMode captureMode)
 
     internal static void ResetForTest()
     {
-        _isColdStart = true;
+        LambdaLifecycleTracker.Reset();
         _captureAnnotations = true;
     }
 }