Sync Dev from Master

.gitignore

@@ -18,6 +18,7 @@
 **/packages
 **/launchSettings.json
 **/Debug/
+**/build/
 
 **/project.lock.json
 

CHANGELOG.md

@@ -1,3 +1,14 @@
+## Release 2024-11-18
+
+### Amazon.Lambda.Core (2.5.0)
+* Added the new `SnapshotRestore` static class for registering SnapStart hooks for before snapshot and after restore.
+
+### Amazon.Lambda.RuntimeSupport (1.12.0)
+* Added support for handling Lambda SnapStart events.
+
+### SnapshotRestore.Registry (1.0.0)
+* New package used by Amazon.Lambda.RuntimeSupport for registering and executing SnapStart hooks.
+
 ## Release 2024-11-14
 
 ### Amazon.Lambda.TestTool.BlazorTester (0.16.0)

LambdaRuntimeDockerfiles/Images/net8/amd64/Dockerfile

@@ -31,6 +31,7 @@ FROM mcr.microsoft.com/dotnet/sdk:8.0-bookworm-slim AS builder
 WORKDIR /src
 COPY ["Libraries/src/Amazon.Lambda.RuntimeSupport", "Repo/Libraries/src/Amazon.Lambda.RuntimeSupport/"]
 COPY ["Libraries/src/Amazon.Lambda.Core", "Repo/Libraries/src/Amazon.Lambda.Core/"]
+COPY ["Libraries/src/SnapshotRestore.Registry", "Repo/Libraries/src/SnapshotRestore.Registry/"]
 COPY ["buildtools/", "Repo/buildtools/"]
 RUN dotnet restore "Repo/Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj" /p:TargetFrameworks=net8.0
 WORKDIR "Repo/Libraries/src/Amazon.Lambda.RuntimeSupport"

LambdaRuntimeDockerfiles/Images/net8/arm64/Dockerfile

@@ -31,6 +31,7 @@ FROM mcr.microsoft.com/dotnet/sdk:8.0-bookworm-slim AS builder
 WORKDIR /src
 COPY ["Libraries/src/Amazon.Lambda.RuntimeSupport", "Repo/Libraries/src/Amazon.Lambda.RuntimeSupport/"]
 COPY ["Libraries/src/Amazon.Lambda.Core", "Repo/Libraries/src/Amazon.Lambda.Core/"]
+COPY ["Libraries/src/SnapshotRestore.Registry", "Repo/Libraries/src/SnapshotRestore.Registry/"]
 COPY ["buildtools/", "Repo/buildtools/"]
 RUN dotnet restore "Repo/Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj" /p:TargetFrameworks=net8.0
 WORKDIR "Repo/Libraries/src/Amazon.Lambda.RuntimeSupport"

Libraries/Amazon.Lambda.RuntimeSupport.slnf

@@ -11,11 +11,13 @@
       "src\\Amazon.Lambda.RuntimeSupport\\Amazon.Lambda.RuntimeSupport.csproj",
       "src\\Amazon.Lambda.Serialization.Json\\Amazon.Lambda.Serialization.Json.csproj",
       "src\\Amazon.Lambda.Serialization.SystemTextJson\\Amazon.Lambda.Serialization.SystemTextJson.csproj",
+      "src\\SnapshotRestore.Registry\\SnapshotRestore.Registry.csproj",
       "test\\Amazon.Lambda.RuntimeSupport.Tests\\Amazon.Lambda.RuntimeSupport.IntegrationTests\\Amazon.Lambda.RuntimeSupport.IntegrationTests.csproj",
       "test\\Amazon.Lambda.RuntimeSupport.Tests\\Amazon.Lambda.RuntimeSupport.UnitTests\\Amazon.Lambda.RuntimeSupport.UnitTests.csproj",
       "test\\Amazon.Lambda.RuntimeSupport.Tests\\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest\\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest.csproj",
       "test\\Amazon.Lambda.RuntimeSupport.Tests\\CustomRuntimeAspNetCoreMinimalApiTest\\CustomRuntimeAspNetCoreMinimalApiTest.csproj",
       "test\\Amazon.Lambda.RuntimeSupport.Tests\\CustomRuntimeFunctionTest\\CustomRuntimeFunctionTest.csproj",
+      "test\\SnapshotRestore.Registry.Tests\\SnapshotRestore.Registry.Tests.csproj",
       "test\\HandlerTestNoSerializer\\HandlerTestNoSerializer.csproj",
       "test\\HandlerTest\\HandlerTest.csproj"
     ]

Libraries/Libraries.sln

@@ -131,6 +131,10 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "TestExecutableServerlessApp
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "TestServerlessApp.NET8", "test\TestServerlessApp.NET8\TestServerlessApp.NET8.csproj", "{7300983D-8FCE-42EA-9B9E-B1C5347D15D8}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "SnapshotRestore.Registry", "src\SnapshotRestore.Registry\SnapshotRestore.Registry.csproj", "{7261A438-8C1D-47AD-98B0-7678F72E4382}"
+EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "SnapshotRestore.Registry.Tests", "test\SnapshotRestore.Registry.Tests\SnapshotRestore.Registry.Tests.csproj", "{A699E183-D0D4-4F26-A0A7-88DA5607F455}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -357,6 +361,14 @@ Global
 		{7300983D-8FCE-42EA-9B9E-B1C5347D15D8}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{7300983D-8FCE-42EA-9B9E-B1C5347D15D8}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{7300983D-8FCE-42EA-9B9E-B1C5347D15D8}.Release|Any CPU.Build.0 = Release|Any CPU
+		{7261A438-8C1D-47AD-98B0-7678F72E4382}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{7261A438-8C1D-47AD-98B0-7678F72E4382}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{7261A438-8C1D-47AD-98B0-7678F72E4382}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{7261A438-8C1D-47AD-98B0-7678F72E4382}.Release|Any CPU.Build.0 = Release|Any CPU
+		{A699E183-D0D4-4F26-A0A7-88DA5607F455}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{A699E183-D0D4-4F26-A0A7-88DA5607F455}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{A699E183-D0D4-4F26-A0A7-88DA5607F455}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{A699E183-D0D4-4F26-A0A7-88DA5607F455}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -421,6 +433,8 @@ Global
 		{0BD83939-458C-4EF5-8663-7098AD1200F2} = {B5BD0336-7D08-492C-8489-42C987E29B39}
 		{DD378063-C54A-44C7-9A6F-32A6A1AE94B3} = {1DE4EE60-45BA-4EF7-BE00-B9EB861E4C69}
 		{7300983D-8FCE-42EA-9B9E-B1C5347D15D8} = {1DE4EE60-45BA-4EF7-BE00-B9EB861E4C69}
+		{7261A438-8C1D-47AD-98B0-7678F72E4382} = {AAB54E74-20B1-42ED-BC3D-CE9F7BC7FD12}
+		{A699E183-D0D4-4F26-A0A7-88DA5607F455} = {1DE4EE60-45BA-4EF7-BE00-B9EB861E4C69}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {503678A4-B8D1-4486-8915-405A3E9CF0EB}

Libraries/src/Amazon.Lambda.Core/Amazon.Lambda.Core.csproj

@@ -6,7 +6,7 @@
     <TargetFrameworks>netstandard2.0;net6.0;net8.0</TargetFrameworks>
     <Description>Amazon Lambda .NET Core support - Core package.</Description>
     <AssemblyTitle>Amazon.Lambda.Core</AssemblyTitle>
-    <Version>2.4.0</Version>
+    <Version>2.5.0</Version>
     <AssemblyName>Amazon.Lambda.Core</AssemblyName>
     <PackageId>Amazon.Lambda.Core</PackageId>
     <PackageTags>AWS;Amazon;Lambda</PackageTags>
@@ -15,7 +15,13 @@
 
   <ItemGroup>
 	<None Include="README.md" Pack="true" PackagePath="\" />
-  </ItemGroup>	
+    <AssemblyAttribute Include="System.Runtime.CompilerServices.InternalsVisibleToAttribute">
+      <_Parameter1>Amazon.Lambda.RuntimeSupport, PublicKey="0024000004800000940000000602000000240000525341310004000001000100db5f59f098d27276c7833875a6263a3cc74ab17ba9a9df0b52aedbe7252745db7274d5271fd79c1f08f668ecfa8eaab5626fa76adc811d3c8fc55859b0d09d3bc0a84eecd0ba891f2b8a2fc55141cdcc37c2053d53491e650a479967c3622762977900eddbf1252ed08a2413f00a28f3a0752a81203f03ccb7f684db373518b4"</_Parameter1>
+    </AssemblyAttribute>
+    <AssemblyAttribute Include="System.Runtime.CompilerServices.InternalsVisibleToAttribute">
+      <_Parameter1>Amazon.Lambda.RuntimeSupport.UnitTests, PublicKey="0024000004800000940000000602000000240000525341310004000001000100db5f59f098d27276c7833875a6263a3cc74ab17ba9a9df0b52aedbe7252745db7274d5271fd79c1f08f668ecfa8eaab5626fa76adc811d3c8fc55859b0d09d3bc0a84eecd0ba891f2b8a2fc55141cdcc37c2053d53491e650a479967c3622762977900eddbf1252ed08a2413f00a28f3a0752a81203f03ccb7f684db373518b4"</_Parameter1>
+    </AssemblyAttribute>
+  </ItemGroup>
 
   <PropertyGroup Condition="'$(TargetFramework)' == 'net8.0'">
     <WarningsAsErrors>IL2026,IL2067,IL2075</WarningsAsErrors>

Libraries/src/Amazon.Lambda.Core/SnapshotRestore.cs

@@ -0,0 +1,59 @@
+using System;
+using System.Collections.Concurrent;
+using System.Threading.Tasks;
+namespace Amazon.Lambda.Core
+{
+#if NET8_0_OR_GREATER
+    /// <summary>
+    /// Static class to register callback hooks to during the snapshot and restore phases of Lambda SnapStart. Hooks 
+    /// should be registered as part of the constructor of the type containing the function handler or before the 
+    /// `LambdaBootstrap` is started in executable assembly Lambda functions.    
+    /// </summary>
+    public static class SnapshotRestore
+    {
+        // We don't want Amazon.Lambda.Core to have any dependencies because the packaged handler code
+        // that gets uploaded to AWS Lambda could have a version mismatch with the version that is already
+        // included in the managed runtime. This class allows us to define a simple API that both the
+        // RuntimeClient and handler code can use to register and then call these actions without
+        // depending on a specific version of SnapshotRestore.Registry.
+        private static readonly ConcurrentQueue<Func<ValueTask>> BeforeSnapshotRegistry = new();
+        private static readonly ConcurrentQueue<Func<ValueTask>> AfterRestoreRegistry = new();
+
+        internal static void CopyBeforeSnapshotCallbacksToRegistry(Action<Func<ValueTask>> restoreHooksRegistryMethod)
+        {
+            // To preserve the order of registry, BeforeSnapshotRegistry in Core needs to be a Queue
+            // These callbacks will be added to the Stack that SnapshotRestore.Registry maintains
+            while (BeforeSnapshotRegistry.TryDequeue(out var registeredAction))
+            {
+                restoreHooksRegistryMethod?.Invoke(registeredAction);
+            }
+        }
+
+        internal static void CopyAfterRestoreCallbacksToRegistry(Action<Func<ValueTask>> restoreHooksRegistryMethod)
+        {
+            while (AfterRestoreRegistry.TryDequeue(out var registeredAction))
+            {
+                restoreHooksRegistryMethod?.Invoke(registeredAction);
+            }
+        }
+
+        /// <summary>
+        /// Register callback hook to be called before Lambda creates a snapshot of the running process. This can be used to warm code in the .NET process or close connections before the snapshot is taken.
+        /// </summary>
+        /// <param name="beforeSnapshotAction"></param>
+        public static void RegisterBeforeSnapshot(Func<ValueTask> beforeSnapshotAction)
+        {
+            BeforeSnapshotRegistry.Enqueue(beforeSnapshotAction);
+        }
+
+        /// <summary>
+        /// Register callback hook to be called after Lambda restores a snapshot of the running process. This can be used to ensure uniqueness after restoration. For example reseeding random number generators.
+        /// </summary>
+        /// <param name="afterRestoreAction"></param>
+        public static void RegisterAfterRestore(Func<ValueTask> afterRestoreAction)
+        {
+            AfterRestoreRegistry.Enqueue(afterRestoreAction);
+        }       
+    }
+#endif
+}

Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj

@@ -4,7 +4,7 @@
 
   <PropertyGroup>
     <TargetFrameworks>netstandard2.0;net5.0;net6.0;net8.0</TargetFrameworks>
-    <Version>1.11.0</Version>
+    <Version>1.12.0</Version>
     <Description>Provides a bootstrap and Lambda Runtime API Client to help you to develop custom .NET Core Lambda Runtimes.</Description>
     <AssemblyTitle>Amazon.Lambda.RuntimeSupport</AssemblyTitle>
     <AssemblyName>Amazon.Lambda.RuntimeSupport</AssemblyName>
@@ -41,6 +41,9 @@
   <ItemGroup>
     <ProjectReference Include="..\Amazon.Lambda.Core\Amazon.Lambda.Core.csproj" />
   </ItemGroup>
+  <ItemGroup Condition="'$(TargetFramework)' == 'net8.0'">
+    <ProjectReference Include="..\SnapshotRestore.Registry\SnapshotRestore.Registry.csproj" />
+  </ItemGroup>
   <ItemGroup>
     <None Update="bootstrap.sh">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/Constants.cs

@@ -32,6 +32,8 @@ internal class Constants
         internal const string ENVIRONMENT_VARIABLE_TELEMETRY_LOG_FD = "_LAMBDA_TELEMETRY_LOG_FD";
         internal const string AWS_LAMBDA_INITIALIZATION_TYPE_PC = "provisioned-concurrency";
         internal const string AWS_LAMBDA_INITIALIZATION_TYPE_ON_DEMAND = "on-demand";
+        internal const string AWS_LAMBDA_INITIALIZATION_TYPE_SNAP_START = "snap-start";
+
 
         internal const string NET_RIC_LOG_LEVEL_ENVIRONMENT_VARIABLE = "AWS_LAMBDA_HANDLER_LOG_LEVEL";
         internal const string NET_RIC_LOG_FORMAT_ENVIRONMENT_VARIABLE = "AWS_LAMBDA_HANDLER_LOG_FORMAT";
@@ -41,6 +43,9 @@ internal class Constants
 
         internal const string LAMBDA_LOG_FORMAT_JSON = "Json";
 
+        internal const string LAMBDA_ERROR_TYPE_BEFORE_SNAPSHOT = "Runtime.BeforeSnapshotError";
+        internal const string LAMBDA_ERROR_TYPE_AFTER_RESTORE = "Runtime.AfterRestoreError";
+
         internal enum AwsLambdaDotNetPreJit
         {
             Never,

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs

@@ -45,6 +45,7 @@ public class LambdaBootstrap : IDisposable
         private InternalLogger _logger = InternalLogger.GetDefaultLogger();
 
         private HttpClient _httpClient;
+        private LambdaBootstrapConfiguration _configuration;
         internal IRuntimeApiClient Client { get; set; }
 
         /// <summary>
@@ -65,7 +66,7 @@ public LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, La
         /// <param name="initializer">Delegate called to initialize the Lambda function.  If not provided the initialization step is skipped.</param>
         /// <returns></returns>
         public LambdaBootstrap(LambdaBootstrapHandler handler, LambdaBootstrapInitializer initializer = null)
-            : this(ConstructHttpClient(), handler, initializer, ownsHttpClient: true)
+            : this(ConstructHttpClient(), handler, initializer, ownsHttpClient: true )
         { }
 
         /// <summary>
@@ -88,6 +89,18 @@ public LambdaBootstrap(HandlerWrapper handlerWrapper, LambdaBootstrapInitializer
         public LambdaBootstrap(HttpClient httpClient, HandlerWrapper handlerWrapper, LambdaBootstrapInitializer initializer = null)
             : this(httpClient, handlerWrapper.Handler, initializer, ownsHttpClient: false)
         { }
+
+        /// <summary>
+        /// Create a LambdaBootstrap that will call the given initializer and handler with custom configuration.
+        /// </summary>
+        /// <param name="handler">Delegate called for each invocation of the Lambda function.</param>
+        /// <param name="initializer">Delegate called to initialize the Lambda function.  If not provided the initialization step is skipped.</param>
+        /// <param name="configuration"> Get configuration to check if Invoke is with Pre JIT or SnapStart enabled </param>
+        /// <returns></returns>
+        internal LambdaBootstrap(LambdaBootstrapHandler handler,
+            LambdaBootstrapInitializer initializer,
+            LambdaBootstrapConfiguration configuration) : this(ConstructHttpClient(), handler, initializer, false, configuration)
+        { }
 
         /// <summary>
         /// Create a LambdaBootstrap that will call the given initializer and handler.
@@ -97,14 +110,15 @@ public LambdaBootstrap(HttpClient httpClient, HandlerWrapper handlerWrapper, Lam
         /// <param name="initializer">Delegate called to initialize the Lambda function.  If not provided the initialization step is skipped.</param>
         /// <param name="ownsHttpClient">Whether the instance owns the HTTP client and should dispose of it.</param>
         /// <returns></returns>
-        private LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, LambdaBootstrapInitializer initializer, bool ownsHttpClient)
+        private LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, LambdaBootstrapInitializer initializer, bool ownsHttpClient, LambdaBootstrapConfiguration configuration = null)
         {
             _httpClient = httpClient ?? throw new ArgumentNullException(nameof(httpClient));
             _handler = handler ?? throw new ArgumentNullException(nameof(handler));
             _ownsHttpClient = ownsHttpClient;
             _initializer = initializer;
             _httpClient.Timeout = RuntimeApiHttpTimeout;
             Client = new RuntimeApiClient(new SystemEnvironmentVariables(), _httpClient);
+            _configuration = configuration ?? LambdaBootstrapConfiguration.GetDefaultConfiguration();
         }
 
         /// <summary>
@@ -124,7 +138,7 @@ private LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, L
             AdjustMemorySettings();
 #endif
 
-            if (UserCodeInit.IsCallPreJit())
+            if (_configuration.IsCallPreJit)
             {
                 this._logger.LogInformation("PreJit: CultureInfo");
                 UserCodeInit.LoadStringCultureInfo();
@@ -137,10 +151,41 @@ private LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, L
             // and then shut down cleanly. Useful for profiling or running local tests with the .NET Lambda Test Tool. This environment
             // variable should never be set when function is deployed to Lambda.
             var runOnce = string.Equals(Environment.GetEnvironmentVariable(Constants.ENVIRONMENT_VARIABLE_AWS_LAMBDA_DOTNET_DEBUG_RUN_ONCE), "true", StringComparison.OrdinalIgnoreCase);
+
+
+            if (_initializer != null && !(await InitializeAsync()))
+            {
+                return;
+            }
+#if NET8_0_OR_GREATER
+            // Check if Initialization type is SnapStart, and invoke the snapshot restore logic.
+            if (_configuration.IsInitTypeSnapstart)
+            {
+                InternalLogger.GetDefaultLogger().LogInformation($"In LambdaBootstrap, Initializing with SnapStart.");
 
-            bool doStartInvokeLoop = _initializer == null || await InitializeAsync();
+                object registry = null;
+                try
+                {
+                    registry = SnapstartHelperCopySnapshotCallbacksIsolated.CopySnapshotCallbacks();
+                }
+                catch (TypeLoadException ex)
+                {
+                    Client.ConsoleLogger.FormattedWriteLine(
+                        Amazon.Lambda.RuntimeSupport.Helpers.LogLevelLoggerWriter.LogLevel.Error.ToString(),
+                        $"Failed to retrieve snapshot hooks from Amazon.Lambda.Core.SnapshotRestore, " +
+                        $"this can be fixed by updating the version of Amazon.Lambda.Core: {ex}",
+                        null);
+                }
+                // no exceptions in calling SnapStart hooks or /restore/next RAPID endpoint
+                if (!(await SnapstartHelperInitializeWithSnapstartIsolatedAsync.InitializeWithSnapstartAsync(Client,
+                        registry)))
+                {
+                    return;
+                };
+            }
+#endif
 
-            while (doStartInvokeLoop && !cancellationToken.IsCancellationRequested)
+            while (!cancellationToken.IsCancellationRequested)
             {
                 try
                 {
@@ -168,8 +213,14 @@ internal async Task<bool> InitializeAsync()
             {
                 WriteUnhandledExceptionToLog(exception);
                 await Client.ReportInitializationErrorAsync(exception);
-                throw;
+#if NET8_0_OR_GREATER
+                if (_configuration.IsInitTypeSnapstart)
+                {
+                    System.Environment.Exit(1); // This needs to be non-zero for Lambda Sandbox to know that Runtime client encountered an exception
+                }
+#endif
             }
+            return false;
         }
 
         internal async Task InvokeOnceAsync(CancellationToken cancellationToken = default)