Merge pull request #834 from hjgraca/feature/coldstart-provisioned-concurrency

chore: Cold start with provisioned concurrency support

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/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/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;
     }
 }

libraries/tests/AWS.Lambda.Powertools.Common.Tests/Core/LambdaLifecycleTrackerTests.cs

@@ -0,0 +1,105 @@
+using System;
+using AWS.Lambda.Powertools.Common.Core;
+using Xunit;
+
+namespace AWS.Lambda.Powertools.Common.Tests;
+
+public class LambdaLifecycleTrackerTests : IDisposable
+    {
+        public LambdaLifecycleTrackerTests()
+        {
+            // Reset before each test to ensure clean state
+            LambdaLifecycleTracker.Reset();
+            Environment.SetEnvironmentVariable(Constants.AWSInitializationTypeEnv, null);
+        }
+
+        public void Dispose()
+        {
+            // Reset after each test
+            LambdaLifecycleTracker.Reset();
+            Environment.SetEnvironmentVariable(Constants.AWSInitializationTypeEnv, null);
+        }
+
+        [Fact]
+        public void IsColdStart_FirstInvocation_ReturnsTrue()
+        {
+            // Act
+            var result = LambdaLifecycleTracker.IsColdStart;
+
+            // Assert
+            Assert.True(result);
+        }
+
+        [Fact]
+        public void IsColdStart_SecondInvocation_ReturnsFalse()
+        {
+            // Arrange - first access to trigger cold start
+            _ = LambdaLifecycleTracker.IsColdStart;
+    
+            // Clear just the AsyncLocal value to simulate new invocation in same container
+            LambdaLifecycleTracker.Reset(resetContainer: false);
+    
+            // Act - second invocation on same container
+            var result = LambdaLifecycleTracker.IsColdStart;
+    
+            // Assert
+            Assert.False(result);
+        }
+
+        [Fact]
+        public void IsColdStart_WithProvisionedConcurrency_ReturnsFalse()
+        {
+            // Arrange
+            Environment.SetEnvironmentVariable(Constants.AWSInitializationTypeEnv, "provisioned-concurrency");
+
+            // Act
+            var result = LambdaLifecycleTracker.IsColdStart;
+
+            // Assert
+            Assert.False(result);
+        }
+
+        [Fact]
+        public void IsColdStart_ReturnsSameValueWithinInvocation()
+        {
+            // Act - access multiple times in the same invocation
+            var firstAccess = LambdaLifecycleTracker.IsColdStart;
+            var secondAccess = LambdaLifecycleTracker.IsColdStart;
+            var thirdAccess = LambdaLifecycleTracker.IsColdStart;
+
+            // Assert
+            Assert.True(firstAccess);
+            Assert.Equal(firstAccess, secondAccess);
+            Assert.Equal(firstAccess, thirdAccess);
+        }
+
+        [Fact]
+        public void Reset_ResetsState()
+        {
+            // Arrange
+            _ = LambdaLifecycleTracker.IsColdStart; // First invocation
+            
+            // Act
+            LambdaLifecycleTracker.Reset();
+            var result = LambdaLifecycleTracker.IsColdStart;
+
+            // Assert
+            Assert.True(result); // Should be true again after reset
+        }
+
+        [Fact]
+        public void Reset_ClearsEnvironmentSetting()
+        {
+            // Arrange
+            Environment.SetEnvironmentVariable(Constants.AWSInitializationTypeEnv, "provisioned-concurrency");
+            _ = LambdaLifecycleTracker.IsColdStart; // Load the environment variable
+            
+            // Act
+            LambdaLifecycleTracker.Reset();
+            Environment.SetEnvironmentVariable(Constants.AWSInitializationTypeEnv, null); // Clear the environment
+            var result = LambdaLifecycleTracker.IsColdStart;
+
+            // Assert
+            Assert.True(result); // Should be true when env var is cleared
+        }
+    }

libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/TracingAttributeTest.cs

@@ -17,6 +17,7 @@
 using System.Linq;
 using System.Text;
 using Amazon.XRay.Recorder.Core;
+using AWS.Lambda.Powertools.Common.Core;
 using AWS.Lambda.Powertools.Tracing.Internal;
 using Xunit;
 
@@ -50,6 +51,8 @@ public void OnEntry_WhenFirstCall_CapturesColdStart()
             var subSegmentCold = segmentCold.Subsegments[0];
 
             // Warm Start Execution
+            // Clear just the AsyncLocal value to simulate new invocation in same container
+            LambdaLifecycleTracker.Reset(resetContainer: false);
             // Start segment
             var segmentWarm = AWSXRayRecorder.Instance.TraceContext.GetEntity();
             _handler.Handle();
@@ -87,6 +90,9 @@ public void OnEntry_WhenFirstCall_And_Service_Not_Set_CapturesColdStart()
             var subSegmentCold = segmentCold.Subsegments[0];
 
             // Warm Start Execution
+            // Clear just the AsyncLocal value to simulate new invocation in same container
+            LambdaLifecycleTracker.Reset(resetContainer: false);
+
             // Start segment
             var segmentWarm = AWSXRayRecorder.Instance.TraceContext.GetEntity();
             _handler.Handle();