From 943862ff5c96e2508f4278c1f869379202fec9a9 Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Fri, 15 Nov 2024 15:56:50 -0800 Subject: [PATCH] Add SnapStart support to Amazon.Lambda.Core and Amazon.Lambda.RuntimeSupport The SnapStart support includes adding a new SnapStart.Registry package for registering and invoking hooks during the SnapStart lifecycle. --------- Co-authored-by: Phil Asmar Co-authored-by: Philippe El Asmar <53088140+philasmar@users.noreply.github.com> Co-authored-by: Saksham Bhalla Co-authored-by: Saksham Bhalla Co-authored-by: Philip Pittle --- .gitignore | 1 + CHANGELOG.md | 11 +++ .../Images/net8/amd64/Dockerfile | 1 + .../Images/net8/arm64/Dockerfile | 1 + Libraries/Amazon.Lambda.RuntimeSupport.slnf | 2 + Libraries/Libraries.sln | 14 +++ .../Amazon.Lambda.Core.csproj | 10 +- .../src/Amazon.Lambda.Core/SnapshotRestore.cs | 59 ++++++++++++ .../Amazon.Lambda.RuntimeSupport.csproj | 5 +- .../Bootstrap/Constants.cs | 5 + .../Bootstrap/LambdaBootstrap.cs | 63 ++++++++++-- .../Client/IRuntimeApiClient.cs | 22 ++++- .../Client/InternalClientAdapted.cs | 84 +++++++++++----- .../Client/RuntimeApiClient.cs | 36 ++++++- .../Context/LambdaBootstrapConfiguration.cs | 35 +++++++ ...tartHelperCopySnapshotCallbacksIsolated.cs | 21 ++++ ...perInitializeWithSnapstartIsolatedAsync.cs | 54 +++++++++++ .../Amazon.Lambda.RuntimeSupport/Program.cs | 22 ++++- .../src/SnapshotRestore.Registry/README.md | 39 ++++++++ .../RestoreHooksRegistry.cs | 77 +++++++++++++++ .../SnapshotRestore.Registry.csproj | 25 +++++ .../Helpers/LambdaToolsHelper.cs | 3 + .../IntegrationTestFixture.cs | 1 + ...zon.Lambda.RuntimeSupport.UnitTests.csproj | 5 +- .../LambdaBootstrapTests.cs | 5 +- .../SnapstartTests.cs | 80 ++++++++++++++++ .../TestHelpers/TestRuntimeApiClient.cs | 19 +++- .../RestoreHooksRegistryTests.cs | 90 ++++++++++++++++++ .../SnapshotRestore.Registry.Tests.csproj | 19 ++++ buildtools/build.proj | 1 + buildtools/snapshotrestore.snk | Bin 0 -> 596 bytes 31 files changed, 766 insertions(+), 44 deletions(-) create mode 100644 Libraries/src/Amazon.Lambda.Core/SnapshotRestore.cs create mode 100644 Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaBootstrapConfiguration.cs create mode 100644 Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperCopySnapshotCallbacksIsolated.cs create mode 100644 Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperInitializeWithSnapstartIsolatedAsync.cs create mode 100644 Libraries/src/SnapshotRestore.Registry/README.md create mode 100644 Libraries/src/SnapshotRestore.Registry/RestoreHooksRegistry.cs create mode 100644 Libraries/src/SnapshotRestore.Registry/SnapshotRestore.Registry.csproj create mode 100644 Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/SnapstartTests.cs create mode 100644 Libraries/test/SnapshotRestore.Registry.Tests/RestoreHooksRegistryTests.cs create mode 100644 Libraries/test/SnapshotRestore.Registry.Tests/SnapshotRestore.Registry.Tests.csproj create mode 100644 buildtools/snapshotrestore.snk diff --git a/.gitignore b/.gitignore index 31dc52b60..f91715274 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ **/packages **/launchSettings.json **/Debug/ +**/build/ **/project.lock.json diff --git a/CHANGELOG.md b/CHANGELOG.md index b37790cd9..c1145ebd3 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/LambdaRuntimeDockerfiles/Images/net8/amd64/Dockerfile b/LambdaRuntimeDockerfiles/Images/net8/amd64/Dockerfile index 3ccc3b915..9c498135b 100644 --- a/LambdaRuntimeDockerfiles/Images/net8/amd64/Dockerfile +++ b/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" diff --git a/LambdaRuntimeDockerfiles/Images/net8/arm64/Dockerfile b/LambdaRuntimeDockerfiles/Images/net8/arm64/Dockerfile index 604cc489b..60c2309f8 100644 --- a/LambdaRuntimeDockerfiles/Images/net8/arm64/Dockerfile +++ b/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" diff --git a/Libraries/Amazon.Lambda.RuntimeSupport.slnf b/Libraries/Amazon.Lambda.RuntimeSupport.slnf index 1f1b383ce..fb03ebc05 100644 --- a/Libraries/Amazon.Lambda.RuntimeSupport.slnf +++ b/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" ] diff --git a/Libraries/Libraries.sln b/Libraries/Libraries.sln index 4f0e88170..3df1056d9 100644 --- a/Libraries/Libraries.sln +++ b/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} diff --git a/Libraries/src/Amazon.Lambda.Core/Amazon.Lambda.Core.csproj b/Libraries/src/Amazon.Lambda.Core/Amazon.Lambda.Core.csproj index 429bfb07f..c6a034fe0 100644 --- a/Libraries/src/Amazon.Lambda.Core/Amazon.Lambda.Core.csproj +++ b/Libraries/src/Amazon.Lambda.Core/Amazon.Lambda.Core.csproj @@ -6,7 +6,7 @@ netstandard2.0;net6.0;net8.0 Amazon Lambda .NET Core support - Core package. Amazon.Lambda.Core - 2.4.0 + 2.5.0 Amazon.Lambda.Core Amazon.Lambda.Core AWS;Amazon;Lambda @@ -15,7 +15,13 @@ - + + <_Parameter1>Amazon.Lambda.RuntimeSupport, PublicKey="0024000004800000940000000602000000240000525341310004000001000100db5f59f098d27276c7833875a6263a3cc74ab17ba9a9df0b52aedbe7252745db7274d5271fd79c1f08f668ecfa8eaab5626fa76adc811d3c8fc55859b0d09d3bc0a84eecd0ba891f2b8a2fc55141cdcc37c2053d53491e650a479967c3622762977900eddbf1252ed08a2413f00a28f3a0752a81203f03ccb7f684db373518b4" + + + <_Parameter1>Amazon.Lambda.RuntimeSupport.UnitTests, PublicKey="0024000004800000940000000602000000240000525341310004000001000100db5f59f098d27276c7833875a6263a3cc74ab17ba9a9df0b52aedbe7252745db7274d5271fd79c1f08f668ecfa8eaab5626fa76adc811d3c8fc55859b0d09d3bc0a84eecd0ba891f2b8a2fc55141cdcc37c2053d53491e650a479967c3622762977900eddbf1252ed08a2413f00a28f3a0752a81203f03ccb7f684db373518b4" + + IL2026,IL2067,IL2075 diff --git a/Libraries/src/Amazon.Lambda.Core/SnapshotRestore.cs b/Libraries/src/Amazon.Lambda.Core/SnapshotRestore.cs new file mode 100644 index 000000000..87854bf01 --- /dev/null +++ b/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 + /// + /// 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. + /// + 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> BeforeSnapshotRegistry = new(); + private static readonly ConcurrentQueue> AfterRestoreRegistry = new(); + + internal static void CopyBeforeSnapshotCallbacksToRegistry(Action> 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> restoreHooksRegistryMethod) + { + while (AfterRestoreRegistry.TryDequeue(out var registeredAction)) + { + restoreHooksRegistryMethod?.Invoke(registeredAction); + } + } + + /// + /// 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. + /// + /// + public static void RegisterBeforeSnapshot(Func beforeSnapshotAction) + { + BeforeSnapshotRegistry.Enqueue(beforeSnapshotAction); + } + + /// + /// 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. + /// + /// + public static void RegisterAfterRestore(Func afterRestoreAction) + { + AfterRestoreRegistry.Enqueue(afterRestoreAction); + } + } +#endif +} \ No newline at end of file diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj b/Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj index bf1a5cf34..7cf4f87c9 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj @@ -4,7 +4,7 @@ netstandard2.0;net5.0;net6.0;net8.0 - 1.11.0 + 1.12.0 Provides a bootstrap and Lambda Runtime API Client to help you to develop custom .NET Core Lambda Runtimes. Amazon.Lambda.RuntimeSupport Amazon.Lambda.RuntimeSupport @@ -41,6 +41,9 @@ + + + Always diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/Constants.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/Constants.cs index e3ef794f3..3b01339f3 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/Constants.cs +++ b/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, diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs index 019d79ea8..fa7996a84 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs +++ b/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; } /// @@ -65,7 +66,7 @@ public LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, La /// Delegate called to initialize the Lambda function. If not provided the initialization step is skipped. /// public LambdaBootstrap(LambdaBootstrapHandler handler, LambdaBootstrapInitializer initializer = null) - : this(ConstructHttpClient(), handler, initializer, ownsHttpClient: true) + : this(ConstructHttpClient(), handler, initializer, ownsHttpClient: true ) { } /// @@ -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) { } + + /// + /// Create a LambdaBootstrap that will call the given initializer and handler with custom configuration. + /// + /// Delegate called for each invocation of the Lambda function. + /// Delegate called to initialize the Lambda function. If not provided the initialization step is skipped. + /// Get configuration to check if Invoke is with Pre JIT or SnapStart enabled + /// + internal LambdaBootstrap(LambdaBootstrapHandler handler, + LambdaBootstrapInitializer initializer, + LambdaBootstrapConfiguration configuration) : this(ConstructHttpClient(), handler, initializer, false, configuration) + { } /// /// Create a LambdaBootstrap that will call the given initializer and handler. @@ -97,7 +110,7 @@ public LambdaBootstrap(HttpClient httpClient, HandlerWrapper handlerWrapper, Lam /// Delegate called to initialize the Lambda function. If not provided the initialization step is skipped. /// Whether the instance owns the HTTP client and should dispose of it. /// - 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)); @@ -105,6 +118,7 @@ private LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, L _initializer = initializer; _httpClient.Timeout = RuntimeApiHttpTimeout; Client = new RuntimeApiClient(new SystemEnvironmentVariables(), _httpClient); + _configuration = configuration ?? LambdaBootstrapConfiguration.GetDefaultConfiguration(); } /// @@ -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 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) diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/IRuntimeApiClient.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/IRuntimeApiClient.cs index 34b0162ec..8dbb34257 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/IRuntimeApiClient.cs +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/IRuntimeApiClient.cs @@ -35,9 +35,10 @@ public interface IRuntimeApiClient /// Report an initialization error as an asynchronous operation. /// /// The exception to report. + /// An optional errorType string that can be used to log higher-context error to customer instead of generic Runtime.Unknown by the Lambda Sandbox. /// The optional cancellation token to use. /// A Task representing the asynchronous operation. - Task ReportInitializationErrorAsync(Exception exception, CancellationToken cancellationToken = default); + Task ReportInitializationErrorAsync(Exception exception, String errorType = null, CancellationToken cancellationToken = default); /// /// Send an initialization error with a type string but no other information as an asynchronous operation. @@ -64,7 +65,26 @@ public interface IRuntimeApiClient /// The optional cancellation token to use. /// A Task representing the asynchronous operation. Task ReportInvocationErrorAsync(string awsRequestId, Exception exception, CancellationToken cancellationToken = default); + +#if NET8_0_OR_GREATER + /// + /// Triggers the snapshot to be taken, and then after resume, restores the lambda + /// context from the Runtime API as an asynchronous operation when SnapStart is enabled. + /// + /// The optional cancellation token to use. + /// A Task representing the asynchronous operation. + Task RestoreNextInvocationAsync(CancellationToken cancellationToken = default); + /// + /// Report a restore error as an asynchronous operation when SnapStart is enabled. + /// + /// The exception to report. + /// An optional errorType string that can be used to log higher-context error to customer instead of generic Runtime.Unknown by the Lambda Sandbox. + /// The optional cancellation token to use. + /// A Task representing the asynchronous operation. + Task ReportRestoreErrorAsync(Exception exception, String errorType = null, CancellationToken cancellationToken = default); +#endif + /// /// Send a response to a function invocation to the Runtime API as an asynchronous operation. /// diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs index 448f955c0..a3eeff854 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs @@ -14,9 +14,13 @@ */ +using System; +using System.IO; using System.Text.Json; using System.Net; using System.Text.Json.Serialization; +using System.Threading; +using System.Threading.Tasks; using Amazon.Lambda.RuntimeSupport.Helpers; namespace Amazon.Lambda.RuntimeSupport @@ -27,24 +31,27 @@ internal partial interface IInternalRuntimeApiClient /// Non-recoverable initialization error. Runtime should exit after reporting the error. Error will be served in response to the first invoke. /// Accepted /// A server side error occurred. - System.Threading.Tasks.Task> ErrorAsync(string lambda_Runtime_Function_Error_Type, string errorJson); + Task> ErrorAsync(string lambda_Runtime_Function_Error_Type, string errorJson, CancellationToken cancellationToken); + - /// Non-recoverable initialization error. Runtime should exit after reporting the error. Error will be served in response to the first invoke. - /// Accepted - /// A server side error occurred. - /// A cancellation token that can be used by other objects or threads to receive notice of cancellation. - System.Threading.Tasks.Task> ErrorAsync(string lambda_Runtime_Function_Error_Type, string errorJson, System.Threading.CancellationToken cancellationToken); - - /// Runtime makes this HTTP request when it is ready to receive and process a new invoke. - /// This is an iterator-style blocking API call. Response contains event JSON document, specific to the invoking service. +#if NET8_0_OR_GREATER + /// + /// Triggers the snapshot to be taken, and then after resume, restores the lambda + /// context from the Runtime API as an asynchronous operation when SnapStart is enabled. + /// + /// /// A Task representing the asynchronous operation. /// A server side error occurred. - System.Threading.Tasks.Task> NextAsync(); + System.Threading.Tasks.Task> RestoreNextAsync(CancellationToken cancellationToken); + Task> RestoreErrorAsync(string lambda_Runtime_Function_Error_Type, + string errorJson, CancellationToken cancellationToken); +#endif + /// Runtime makes this HTTP request when it is ready to receive and process a new invoke. /// This is an iterator-style blocking API call. Response contains event JSON document, specific to the invoking service. /// A server side error occurred. - /// A cancellation token that can be used by other objects or threads to receive notice of cancellation. - System.Threading.Tasks.Task> NextAsync(System.Threading.CancellationToken cancellationToken); + System.Threading.Tasks.Task> NextAsync(CancellationToken cancellationToken); + /// Runtime makes this request in order to submit a response. /// Accepted @@ -54,6 +61,8 @@ internal partial interface IInternalRuntimeApiClient /// Runtime makes this request in order to submit a response. /// Accepted /// A server side error occurred. + /// + /// /// A cancellation token that can be used by other objects or threads to receive notice of cancellation. System.Threading.Tasks.Task> ResponseAsync(string awsRequestId, System.IO.Stream outputStream, System.Threading.CancellationToken cancellationToken); @@ -106,19 +115,17 @@ public string BaseUrl /// Non-recoverable initialization error. Runtime should exit after reporting the error. Error will be served in response to the first invoke. /// Accepted /// A server side error occurred. - public System.Threading.Tasks.Task> ErrorAsync(string lambda_Runtime_Function_Error_Type, string errorJson) + public Task> ErrorAsync(string lambda_Runtime_Function_Error_Type, string errorJson, CancellationToken cancellationToken) { - return ErrorAsync(lambda_Runtime_Function_Error_Type, errorJson, System.Threading.CancellationToken.None); + return ErrorAsync(lambda_Runtime_Function_Error_Type, errorJson, "/runtime/init/error", cancellationToken ); } - /// Non-recoverable initialization error. Runtime should exit after reporting the error. Error will be served in response to the first invoke. - /// Accepted - /// A server side error occurred. - /// A cancellation token that can be used by other objects or threads to receive notice of cancellation. - public async System.Threading.Tasks.Task> ErrorAsync(string lambda_Runtime_Function_Error_Type, string errorJson, System.Threading.CancellationToken cancellationToken) + private async System.Threading.Tasks.Task> ErrorAsync( + string lambda_Runtime_Function_Error_Type, string errorJson, string url, + System.Threading.CancellationToken cancellationToken) { - var urlBuilder_ = new System.Text.StringBuilder(); - urlBuilder_.Append(BaseUrl != null ? BaseUrl.TrimEnd('/') : "").Append("/runtime/init/error"); + var urlBuilder_ = new System.Text.StringBuilder(); + urlBuilder_.Append(BaseUrl != null ? BaseUrl.TrimEnd('/') : "").Append(url); var client_ = _httpClient; try @@ -215,16 +222,40 @@ public async System.Threading.Tasks.Task> ErrorA /// Runtime makes this HTTP request when it is ready to receive and process a new invoke. /// This is an iterator-style blocking API call. Response contains event JSON document, specific to the invoking service. /// A server side error occurred. - public System.Threading.Tasks.Task> NextAsync() + public System.Threading.Tasks.Task> NextAsync(CancellationToken cancellationToken) { - return NextAsync(System.Threading.CancellationToken.None); + return NextAsync("/runtime/invocation/next", cancellationToken); + } + +#if NET8_0_OR_GREATER + /// + /// Restores the lambda context from the Runtime API as an asynchronous operation when SnapStart is enabled + /// + /// A Task representing the asynchronous operation. + public Task> RestoreNextAsync(CancellationToken cancellationToken) + { + return NextAsync("/runtime/restore/next", cancellationToken); } + + /// Non-recoverable restore error when SnapStart is enabled. Runtime should exit after reporting the error. + /// A Task representing the asynchronous operation. + /// A server side error occurred. + public async Task> RestoreErrorAsync(string lambda_Runtime_Function_Error_Type, + string errorJson, CancellationToken cancellationToken) + { + return await ErrorAsync(lambda_Runtime_Function_Error_Type, errorJson, "/runtime/restore/error", cancellationToken); + + } +#endif + + /// Runtime makes this HTTP request when it is ready to receive and process a new invoke. /// This is an iterator-style blocking API call. Response contains event JSON document, specific to the invoking service. /// A server side error occurred. + /// RAPID API endpointUrl that is invoked to process the request /// A cancellation token that can be used by other objects or threads to receive notice of cancellation. - public async System.Threading.Tasks.Task> NextAsync(System.Threading.CancellationToken cancellationToken) + public async System.Threading.Tasks.Task> NextAsync(String endpointUrl, CancellationToken cancellationToken) { this._logger.LogInformation("Starting InternalClient.NextAsync"); @@ -234,9 +265,10 @@ public async System.Threading.Tasks.Task> ErrorA using (var request_ = new System.Net.Http.HttpRequestMessage()) { request_.Method = new System.Net.Http.HttpMethod("GET"); + request_.Headers.Accept.Add(System.Net.Http.Headers.MediaTypeWithQualityHeaderValue.Parse("application/json")); - var url_ = BaseUrl.TrimEnd('/') + "/runtime/invocation/next"; + var url_ = BaseUrl.TrimEnd('/') + endpointUrl; request_.RequestUri = new System.Uri(url_, System.UriKind.Absolute); var response_ = await client_.SendAsync(request_, System.Net.Http.HttpCompletionOption.ResponseHeadersRead, cancellationToken).ConfigureAwait(false); @@ -306,6 +338,8 @@ public System.Threading.Tasks.Task> ResponseAsyn /// Runtime makes this request in order to submit a response. /// Accepted /// A server side error occurred. + /// + /// /// A cancellation token that can be used by other objects or threads to receive notice of cancellation. public async System.Threading.Tasks.Task> ResponseAsync(string awsRequestId, System.IO.Stream outputStream, System.Threading.CancellationToken cancellationToken) { diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiClient.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiClient.cs index 2cbc4cb83..f8e619f12 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiClient.cs +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiClient.cs @@ -72,21 +72,22 @@ internal RuntimeApiClient(IEnvironmentVariables environmentVariables, IInternalR /// Report an initialization error as an asynchronous operation. /// /// The exception to report. + /// An optional errorType string that can be used to log higher-context error to customer instead of generic Runtime.Unknown by the Lambda Sandbox. /// The optional cancellation token to use. /// A Task representing the asynchronous operation. - public Task ReportInitializationErrorAsync(Exception exception, CancellationToken cancellationToken = default) + public Task ReportInitializationErrorAsync(Exception exception, String errorType = null, CancellationToken cancellationToken = default) { if (exception == null) throw new ArgumentNullException(nameof(exception)); - return _internalClient.ErrorAsync(null, LambdaJsonExceptionWriter.WriteJson(ExceptionInfo.GetExceptionInfo(exception)), cancellationToken); + return _internalClient.ErrorAsync(errorType, LambdaJsonExceptionWriter.WriteJson(ExceptionInfo.GetExceptionInfo(exception)), cancellationToken); } /// /// Send an initialization error with a type string but no other information as an asynchronous operation. /// This can be used to directly control flow in Step Functions without creating an Exception class and throwing it. /// - /// The type of the error to report to Lambda. This does not need to be a .NET type name. + /// The type of the error to report to Lambda. This does not need to be a .NET type name. /// The optional cancellation token to use. /// A Task representing the asynchronous operation. public Task ReportInitializationErrorAsync(string errorType, CancellationToken cancellationToken = default) @@ -140,6 +141,35 @@ public Task ReportInvocationErrorAsync(string awsRequestId, Exception exception, return _internalClient.ErrorWithXRayCauseAsync(awsRequestId, exceptionInfo.ErrorType, exceptionInfoJson, exceptionInfoXRayJson, cancellationToken); } + +#if NET8_0_OR_GREATER + + /// + /// Triggers the snapshot to be taken, and then after resume, restores the lambda + /// context from the Runtime API as an asynchronous operation when SnapStart is enabled. + /// + /// The optional cancellation token to use. + /// A Task representing the asynchronous operation. + public async Task RestoreNextInvocationAsync(CancellationToken cancellationToken = default) + { + await _internalClient.RestoreNextAsync(cancellationToken); + } + + /// + /// Report a restore error as an asynchronous operation when SnapStart is enabled. + /// + /// The exception to report. + /// An optional errorType string that can be used to log higher-context error to customer instead of generic Runtime.Unknown by the Lambda Sandbox. + /// The optional cancellation token to use. + /// A Task representing the asynchronous operation. + public Task ReportRestoreErrorAsync(Exception exception, String errorType = null, CancellationToken cancellationToken = default) + { + if (exception == null) + throw new ArgumentNullException(nameof(exception)); + + return _internalClient.RestoreErrorAsync(errorType, LambdaJsonExceptionWriter.WriteJson(ExceptionInfo.GetExceptionInfo(exception)), cancellationToken); + } +#endif /// diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaBootstrapConfiguration.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaBootstrapConfiguration.cs new file mode 100644 index 000000000..84ea76615 --- /dev/null +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaBootstrapConfiguration.cs @@ -0,0 +1,35 @@ +using System; +using Amazon.Lambda.RuntimeSupport.Bootstrap; +using Amazon.Lambda.RuntimeSupport.Helpers; + +namespace Amazon.Lambda.RuntimeSupport +{ + internal class LambdaBootstrapConfiguration + { + internal bool IsCallPreJit { get; set; } + internal bool IsInitTypeSnapstart { get; set; } + + internal LambdaBootstrapConfiguration(bool isCallPreJit, bool isInitTypeSnapstart) + { + if (IsInitTypeSnapstart) + InternalLogger.GetDefaultLogger().LogInformation("Setting Init type to SnapStart"); + + IsCallPreJit = isCallPreJit; + IsInitTypeSnapstart = isInitTypeSnapstart; + } + + internal static LambdaBootstrapConfiguration GetDefaultConfiguration() + { + bool isCallPreJit = UserCodeInit.IsCallPreJit(); +#if NET8_0_OR_GREATER + bool isInitTypeSnapstart = + string.Equals( + Environment.GetEnvironmentVariable(Constants.ENVIRONMENT_VARIABLE_AWS_LAMBDA_INITIALIZATION_TYPE), + Constants.AWS_LAMBDA_INITIALIZATION_TYPE_SNAP_START); + + return new LambdaBootstrapConfiguration(isCallPreJit, isInitTypeSnapstart); +#endif + return new LambdaBootstrapConfiguration(isCallPreJit, false); + } + } +} \ No newline at end of file diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperCopySnapshotCallbacksIsolated.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperCopySnapshotCallbacksIsolated.cs new file mode 100644 index 000000000..740b29b0d --- /dev/null +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperCopySnapshotCallbacksIsolated.cs @@ -0,0 +1,21 @@ +using System; +using System.Collections.Generic; +using System.Text; + +namespace Amazon.Lambda.RuntimeSupport.Helpers +{ +#if NET8_0_OR_GREATER + internal static class SnapstartHelperCopySnapshotCallbacksIsolated + { + internal static object CopySnapshotCallbacks() + { + var logger = InternalLogger.GetDefaultLogger(); + var restoreHooksRegistry = new SnapshotRestore.Registry.RestoreHooksRegistry(logger.LogInformation); + Core.SnapshotRestore.CopyBeforeSnapshotCallbacksToRegistry(restoreHooksRegistry.RegisterBeforeSnapshot); + Core.SnapshotRestore.CopyAfterRestoreCallbacksToRegistry(restoreHooksRegistry.RegisterAfterRestore); + + return restoreHooksRegistry; + } + } +#endif +} diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperInitializeWithSnapstartIsolatedAsync.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperInitializeWithSnapstartIsolatedAsync.cs new file mode 100644 index 000000000..e0874f50b --- /dev/null +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperInitializeWithSnapstartIsolatedAsync.cs @@ -0,0 +1,54 @@ +using System; +using System.Threading.Tasks; +using Amazon.Lambda.RuntimeSupport.Bootstrap; + +namespace Amazon.Lambda.RuntimeSupport.Helpers +{ +#if NET8_0_OR_GREATER + /// + /// Anywhere this class is used in RuntimeSupport it should be wrapped around a try/catch block catching TypeLoadException. + /// If the version of Amazon.Lambda.Core in the deployment bundle is out of date the type that is accessing SnapshotRestore + /// will throw a TypeLoadException when the type is loaded. This extra layer for accessing SnapshotRestore is used so + /// classes like LambdaBootstrap can attempt accessing SnapshotRestore and catch the TypeLoadException if the type does not exist. + /// If LambdaBootstrap was to directly access SnapshotRestore from Amazon.Lambda.Core a TypeLoadException would be thrown + /// when LambdaBootstrap is loaded. + /// + internal static class SnapstartHelperInitializeWithSnapstartIsolatedAsync + { + /// + /// This function will invoke the beforeSnapshot hooks, restore lambda context and run the afterRestore hooks. + /// This will be used when SnapStart is enabled + /// + internal static async Task InitializeWithSnapstartAsync(IRuntimeApiClient client, object restoreHooksRegistry) + { + restoreHooksRegistry = restoreHooksRegistry == null ? new SnapshotRestore.Registry.RestoreHooksRegistry() : restoreHooksRegistry; + var logger = InternalLogger.GetDefaultLogger(); + try + { + await ((SnapshotRestore.Registry.RestoreHooksRegistry)restoreHooksRegistry).InvokeBeforeSnapshotCallbacks(); + await client.RestoreNextInvocationAsync(); + } + catch (Exception ex) + { + client.ConsoleLogger.FormattedWriteLine(LogLevelLoggerWriter.LogLevel.Error.ToString(), ex, + $"Failed to invoke before snapshot hooks: {ex}"); + await client.ReportInitializationErrorAsync(ex, Constants.LAMBDA_ERROR_TYPE_BEFORE_SNAPSHOT); + return false; + } + try + { + await ((SnapshotRestore.Registry.RestoreHooksRegistry)restoreHooksRegistry).InvokeAfterRestoreCallbacks(); + } + catch (Exception ex) + { + client.ConsoleLogger.FormattedWriteLine(LogLevelLoggerWriter.LogLevel.Error.ToString(), ex, + $"Failed to invoke after restore callables: {ex}"); + await client.ReportRestoreErrorAsync(ex, Constants.LAMBDA_ERROR_TYPE_AFTER_RESTORE); + return false; + } + + return true; + } + } +#endif +} diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Program.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Program.cs index fd63562c8..b3c7f8d91 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Program.cs +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Program.cs @@ -13,7 +13,10 @@ * permissions and limitations under the License. */ +using Amazon.Lambda.RuntimeSupport.Helpers; using System; +using System.IO; +using System.Runtime.Loader; using System.Threading.Tasks; namespace Amazon.Lambda.RuntimeSupport @@ -27,15 +30,32 @@ class Program #endif private static async Task Main(string[] args) { +#if NET8_0_OR_GREATER + AssemblyLoadContext.Default.Resolving += ResolveSnapshotRestoreAssembly; if (args.Length == 0) { throw new ArgumentException("The function handler was not provided via command line arguments.", nameof(args)); } - +#endif var handler = args[0]; RuntimeSupportInitializer runtimeSupportInitializer = new RuntimeSupportInitializer(handler); await runtimeSupportInitializer.RunLambdaBootstrap(); } + +#if NET8_0_OR_GREATER + [System.Diagnostics.CodeAnalysis.RequiresUnreferencedCode("This code is only exercised in the class library programming model. Native AOT will not use this code path.")] + private static System.Reflection.Assembly ResolveSnapshotRestoreAssembly(AssemblyLoadContext assemblyContext, System.Reflection.AssemblyName assemblyName) + { + const string assemblyPath = "/var/runtime/SnapshotRestore.Registry.dll"; + InternalLogger.GetDefaultLogger().LogInformation("Resolving assembly: " + assemblyName.Name); + if (string.Equals(assemblyName.Name, "SnapshotRestore.Registry", StringComparison.InvariantCultureIgnoreCase) && File.Exists(assemblyPath)) + { + return assemblyContext.LoadFromAssemblyPath(assemblyPath); + } + + return null; + } +#endif } } diff --git a/Libraries/src/SnapshotRestore.Registry/README.md b/Libraries/src/SnapshotRestore.Registry/README.md new file mode 100644 index 000000000..73b0ef128 --- /dev/null +++ b/Libraries/src/SnapshotRestore.Registry/README.md @@ -0,0 +1,39 @@ +### Overview +The primary aim of this project is to develop a new API to register and retrieve tasks of type `ValueTask`. + +The class uses a `ConcurrentStack` and a `ConcurrentQueue` to store the registered hooks, which are `Func` objects. + +The `RegisterBeforeSnapshot` and `RegisterAfterRestore` methods allow users to register their own hooks, while the `InvokeBeforeSnapshotCallbacks` and `InvokeAfterRestoreCallbacks` methods allow the caller to invoke these snapstart hooks. + +This implementation is used for `Snapstart`, a feature that allows for quick restoration of application state. + +### Sample Usage + +``` +/// +/// Example class to demonstrate usage of SnapshotRestore.Registry library +/// +public class SnapstartExample +{ + private Guid _myExecutionEnvironmentGuid; + public SnapstartExample() + { + // This GUID is set for non-restore use-cases such as testing or if SnapStart is turned off + _myExecutionEnvironmentGuid = new Guid(); + // Register the method which will run after each restore. You may need to update Amazon.Lambda.Core to see this + Amazon.Lambda.Core.SnapshotRestore.RegisterAfterRestore(MyAfterRestore); + } + + private ValueTask MyAfterRestore() + { + // After we restore this snapshot to a new execution environment, update the GUID + _myExecutionEnvironmentGuid = new Guid(); + return ValueTask.CompletedTask; + } + + public string Handler() + { + return $"Hello World! My Execution Environment GUID is {_myExecutionEnvironmentGuid}"; + } +} +``` \ No newline at end of file diff --git a/Libraries/src/SnapshotRestore.Registry/RestoreHooksRegistry.cs b/Libraries/src/SnapshotRestore.Registry/RestoreHooksRegistry.cs new file mode 100644 index 000000000..e150643f8 --- /dev/null +++ b/Libraries/src/SnapshotRestore.Registry/RestoreHooksRegistry.cs @@ -0,0 +1,77 @@ +using System; +using System.Collections.Concurrent; +using System.Threading.Tasks; + +namespace SnapshotRestore.Registry; + +/// +/// .NET Implementation for Registering BeforeSnapshot and AfterRestore hooks +/// for Snapstart +/// +public class RestoreHooksRegistry +{ + private ConcurrentStack> _beforeSnapshotRegistry = new(); + private ConcurrentQueue> _afterRestoreRegistry = new(); + + private Action _logger; + + /// + /// Creates an instance of RestoreHooksRegistry. + /// + /// An optional callback logger. + public RestoreHooksRegistry(Action logger = null) + { + _logger = logger ?? (x => { }); + } + + /// + /// Register a ValueTask by adding it into the Before Snapshot Registry + /// + /// + public void RegisterBeforeSnapshot(Func func) + { + _beforeSnapshotRegistry.Push(func); + } + /// + /// Register a ValueTask by adding it into the After Restore Registry + /// + /// + public void RegisterAfterRestore(Func func) + { + _afterRestoreRegistry.Enqueue(func); + } + + /// + /// Invoke all the registered before snapshot callbacks. + /// + /// + public async Task InvokeBeforeSnapshotCallbacks() + { + if (_beforeSnapshotRegistry != null) + { + _logger($"Invoking {_beforeSnapshotRegistry.Count} beforeSnapshotCallables"); + while (_beforeSnapshotRegistry.TryPop(out var beforeSnapshotCallable)) + { + _logger($"Calling beforeSnapshotCallable: {beforeSnapshotCallable.Method.Name}"); + await beforeSnapshotCallable(); + } + } + } + + /// + /// Invoke all the registered after restore callbacks. + /// + /// + public async Task InvokeAfterRestoreCallbacks() + { + if (_afterRestoreRegistry != null) + { + _logger($"Invoking {_afterRestoreRegistry.Count} afterRestoreCallables"); + while (_afterRestoreRegistry.TryDequeue(out var afterRestoreCallable)) + { + _logger($"Calling afterRestoreCallable: {afterRestoreCallable.Method.Name}"); + await afterRestoreCallable(); + } + } + } +} \ No newline at end of file diff --git a/Libraries/src/SnapshotRestore.Registry/SnapshotRestore.Registry.csproj b/Libraries/src/SnapshotRestore.Registry/SnapshotRestore.Registry.csproj new file mode 100644 index 000000000..049fcf90e --- /dev/null +++ b/Libraries/src/SnapshotRestore.Registry/SnapshotRestore.Registry.csproj @@ -0,0 +1,25 @@ + + + + net8.0 + 1.0.0 + Provides a Restore Hooks library to help you register before snapshot and after restore hooks. + SnapshotRestore.Registry + SnapshotRestore.Registry + AWS;Amazon;Lambda + README.md + true + true + latest + IL2026,IL2067,IL2075 + true + true + Amazon Web Services + ..\..\..\buildtools\snapshotrestore.snk + true + + + + + + diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/Helpers/LambdaToolsHelper.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/Helpers/LambdaToolsHelper.cs index 067d6f689..42a02aac6 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/Helpers/LambdaToolsHelper.cs +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/Helpers/LambdaToolsHelper.cs @@ -70,6 +70,9 @@ private static void CopyDirectory(DirectoryInfo dir, string destDirName) foreach (var subdir in dirs) { + if (string.Equals(subdir.Name, ".vs", System.StringComparison.OrdinalIgnoreCase)) + continue; + var tempPath = Path.Combine(destDirName, subdir.Name); var subDir = new DirectoryInfo(subdir.FullName); CopyDirectory(subDir, tempPath); diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/IntegrationTestFixture.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/IntegrationTestFixture.cs index eba66b47b..bc6ddad99 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/IntegrationTestFixture.cs +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/IntegrationTestFixture.cs @@ -34,6 +34,7 @@ public async Task InitializeAsync() await LambdaToolsHelper.LambdaPackage(toolPath, "net6.0", testAppPath); } + public Task DisposeAsync() { foreach (var tempPath in _tempPaths) diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/Amazon.Lambda.RuntimeSupport.UnitTests.csproj b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/Amazon.Lambda.RuntimeSupport.UnitTests.csproj index 943c29ada..51b2e9fba 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/Amazon.Lambda.RuntimeSupport.UnitTests.csproj +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/Amazon.Lambda.RuntimeSupport.UnitTests.csproj @@ -1,7 +1,9 @@  - + net8.0 + ..\..\..\..\buildtools\public.snk + true @@ -15,6 +17,7 @@ + diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs index aeabdc108..44f20aa7e 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs @@ -83,13 +83,12 @@ public async Task NoInitializer() } [Fact] - public async Task InitializerThrowsException() + public async Task InitializerHandlesExceptionsGracefully() { using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerAsync, _testInitializer.InitializeThrowAsync)) { bootstrap.Client = _testRuntimeApiClient; - var exception = await Assert.ThrowsAsync(async () => { await bootstrap.RunAsync(); }); - Assert.Equal(TestInitializer.InitializeExceptionMessage, exception.Message); + await bootstrap.RunAsync(); } Assert.True(_testRuntimeApiClient.ReportInitializationErrorAsyncExceptionCalled); diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/SnapstartTests.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/SnapstartTests.cs new file mode 100644 index 000000000..aaedf943a --- /dev/null +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/SnapstartTests.cs @@ -0,0 +1,80 @@ +using System; +using System.Collections.Generic; +using System.Threading.Tasks; +using Xunit; +using static Amazon.Lambda.RuntimeSupport.Bootstrap.Constants; + +namespace Amazon.Lambda.RuntimeSupport.UnitTests; +public class SnapstartTests +{ + TestHandler _testFunction; + TestInitializer _testInitializer; + TestRuntimeApiClient _testRuntimeApiClient; + TestEnvironmentVariables _environmentVariables; + + public SnapstartTests() + { + _environmentVariables = new TestEnvironmentVariables(); + var headers = new Dictionary> + { + { + RuntimeApiHeaders.HeaderAwsRequestId, new List { "request_id" } + }, + { + RuntimeApiHeaders.HeaderInvokedFunctionArn, new List { "invoked_function_arn" } + } + }; + _testRuntimeApiClient = new TestRuntimeApiClient(_environmentVariables, headers); + _testInitializer = new TestInitializer(); + _testFunction = new TestHandler(); + } + + [Fact] + public async void VerifyRestoreNextIsCalledWhenSnapstartIsEnabled() + { + using var bootstrap = + new LambdaBootstrap(_testFunction.BaseHandlerAsync, _testInitializer.InitializeTrueAsync, configuration: new LambdaBootstrapConfiguration(false, true)); + bootstrap.Client = _testRuntimeApiClient; + await bootstrap.RunAsync(_testFunction.CancellationSource.Token); + Assert.True(_testRuntimeApiClient.RestoreNextInvocationAsyncCalled); + } + + [Fact] + public async void VerifyRestoreNextIsNotCalledWhenSnapstartIsDisabled() + { + using var bootstrap = + new LambdaBootstrap(_testFunction.BaseHandlerAsync, _testInitializer.InitializeTrueAsync, configuration: new LambdaBootstrapConfiguration(false, false)); + bootstrap.Client = _testRuntimeApiClient; + Environment.SetEnvironmentVariable(ENVIRONMENT_VARIABLE_AWS_LAMBDA_INITIALIZATION_TYPE, AWS_LAMBDA_INITIALIZATION_TYPE_ON_DEMAND); + await bootstrap.RunAsync(_testFunction.CancellationSource.Token); + Assert.False(_testRuntimeApiClient.RestoreNextInvocationAsyncCalled); + } + + + [Fact] + public async void VerifyInitializeErrorIsCalledWhenExceptionInBeforeSnapshotCallables() + { + using var bootstrap = + new LambdaBootstrap(_testFunction.BaseHandlerAsync, _testInitializer.InitializeTrueAsync, configuration: new LambdaBootstrapConfiguration(false, true)); + bootstrap.Client = _testRuntimeApiClient; + Core.SnapshotRestore.RegisterBeforeSnapshot( + () => throw new Exception("Error in Before snapshot callable 1")); + Core.SnapshotRestore.RegisterBeforeSnapshot(() => ValueTask.CompletedTask); + await bootstrap.RunAsync(_testFunction.CancellationSource.Token); + Assert.True(_testRuntimeApiClient.ReportInitializationErrorAsyncExceptionCalled); + } + + [Fact] + public async void VerifyRestoreErrorIsCalledWhenExceptionInAfterRestoreCallables() + { + using (var bootstrap = + new LambdaBootstrap(_testFunction.BaseHandlerAsync, _testInitializer.InitializeTrueAsync, new LambdaBootstrapConfiguration(false, true))) + { + bootstrap.Client = _testRuntimeApiClient; + Core.SnapshotRestore.RegisterAfterRestore(() => ValueTask.CompletedTask); + Core.SnapshotRestore.RegisterAfterRestore(() => throw new Exception("Error in After restore callable 1")); + await bootstrap.RunAsync(_testFunction.CancellationSource.Token); + Assert.True(_testRuntimeApiClient.ReportRestoreErrorAsyncCalled); + } + } +} \ No newline at end of file diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/TestHelpers/TestRuntimeApiClient.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/TestHelpers/TestRuntimeApiClient.cs index bf449e803..ef500e746 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/TestHelpers/TestRuntimeApiClient.cs +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/TestHelpers/TestRuntimeApiClient.cs @@ -38,6 +38,10 @@ public TestRuntimeApiClient(IEnvironmentVariables environmentVariables, Dictiona } public bool GetNextInvocationAsyncCalled { get; private set; } + public bool RestoreNextInvocationAsyncCalled { get; private set; } + public bool ReportRestoreErrorAsyncCalled { get; private set; } + + public bool ReportInitializationErrorAsyncExceptionCalled { get; private set; } public bool ReportInitializationErrorAsyncTypeCalled { get; private set; } public bool ReportInvocationErrorAsyncExceptionCalled { get; private set; } @@ -98,8 +102,14 @@ public Task GetNextInvocationAsync(CancellationToken cancella new TestDateTimeHelper(), new Helpers.SimpleLoggerWriter()) }); } + + public Task RestoreNextInvocationAsync(CancellationToken cancellationToken = default) + { + RestoreNextInvocationAsyncCalled = true; + return Task.Run(() => { }); + } - public Task ReportInitializationErrorAsync(Exception exception, CancellationToken cancellationToken = default) + public Task ReportInitializationErrorAsync(Exception exception, String errorType = null, CancellationToken cancellationToken = default) { LastRecordedException = exception; ReportInitializationErrorAsyncExceptionCalled = true; @@ -124,6 +134,13 @@ public Task ReportInvocationErrorAsync(string awsRequestId, string errorType, Ca ReportInvocationErrorAsyncTypeCalled = true; return Task.Run(() => { }); } + + public Task ReportRestoreErrorAsync(Exception exception, String errorType = null, CancellationToken cancellationToken = default) + { + ReportRestoreErrorAsyncCalled = true; + + + return Task.Run(() => { }); } public Task SendResponseAsync(string awsRequestId, Stream outputStream, CancellationToken cancellationToken = default) { diff --git a/Libraries/test/SnapshotRestore.Registry.Tests/RestoreHooksRegistryTests.cs b/Libraries/test/SnapshotRestore.Registry.Tests/RestoreHooksRegistryTests.cs new file mode 100644 index 000000000..37ef4267f --- /dev/null +++ b/Libraries/test/SnapshotRestore.Registry.Tests/RestoreHooksRegistryTests.cs @@ -0,0 +1,90 @@ +using System; +using Xunit; +namespace SnapshotRestore.Registry.Tests; + +public class RestoreHooksRegistryTests +{ + private DateTimeOffset? _func1InvokeTime = null; + private DateTimeOffset? _func2InvokeTime = null; + + [Fact] + public async Task RegisterBeforeSnapshotAsyncShouldAddValueTaskToRegistryAsync() + { + // Arrange + _func1InvokeTime = null; + _func2InvokeTime = null; + RestoreHooksRegistry registry = new(Console.WriteLine); + registry.RegisterBeforeSnapshot(TestFunc1); + registry.RegisterBeforeSnapshot(TestFunc2); + + // Act + await registry.InvokeBeforeSnapshotCallbacks(); + + // Assert + Assert.NotNull(_func1InvokeTime); + Assert.NotNull(_func2InvokeTime); + Assert.True(_func2InvokeTime < _func1InvokeTime, "func2InvokeTime should be less than func1InvokeTime, " + + "since func2InvokeTime was registered second, and BeforeSnapshot " + + "tasks are called in the reverse order they were registered."); + } + + [Fact] + public async Task RegisterAfterRestoreAsync_ShouldAddValueTaskToRegistryAsync() + { + // Arrange + _func1InvokeTime = null; + _func2InvokeTime = null; + RestoreHooksRegistry registry = new(Console.WriteLine); + registry.RegisterAfterRestore(TestFunc1); + registry.RegisterAfterRestore(TestFunc2); + + // Act + await registry.InvokeAfterRestoreCallbacks(); + + // Assert + Assert.NotNull(_func1InvokeTime); + Assert.NotNull(_func2InvokeTime); + Assert.True(_func1InvokeTime < _func2InvokeTime, "func1InvokeTime should be less than or equal to " + + "func2InvokeTime, since it was registered first, and AfterRestore " + + "tasks are called in the order they were registered."); + } + + [Fact] + public async Task LoggerIsNotRequired() + { + // Arrange + RestoreHooksRegistry registry = new(logger: null); + registry.RegisterAfterRestore(TestFunc1); + registry.RegisterAfterRestore(TestFunc2); + + Exception? exception = null; + + // Act + try + { + await registry.InvokeAfterRestoreCallbacks(); + } + catch (Exception e) + { + exception = e; + } + + // Assert + Assert.Null(exception); + } + + + private ValueTask TestFunc1() + { + _func1InvokeTime = DateTimeOffset.UtcNow; + Thread.Sleep(10); // So the times of func1 and func2 aren't ever exactly equal + return ValueTask.CompletedTask; + } + + private ValueTask TestFunc2() + { + _func2InvokeTime = DateTimeOffset.UtcNow; + Thread.Sleep(10); // So the times of func1 and func2 aren't ever exactly equal + return ValueTask.CompletedTask; + } +} \ No newline at end of file diff --git a/Libraries/test/SnapshotRestore.Registry.Tests/SnapshotRestore.Registry.Tests.csproj b/Libraries/test/SnapshotRestore.Registry.Tests/SnapshotRestore.Registry.Tests.csproj new file mode 100644 index 000000000..dd41e38a4 --- /dev/null +++ b/Libraries/test/SnapshotRestore.Registry.Tests/SnapshotRestore.Registry.Tests.csproj @@ -0,0 +1,19 @@ + + + + net8.0 + enable + enable + + + + + + all + runtime; build; native; contentfiles; analyzers; buildtransitive + + + + + + diff --git a/buildtools/build.proj b/buildtools/build.proj index 90d792318..bb9ac9a43 100644 --- a/buildtools/build.proj +++ b/buildtools/build.proj @@ -194,6 +194,7 @@ + diff --git a/buildtools/snapshotrestore.snk b/buildtools/snapshotrestore.snk new file mode 100644 index 0000000000000000000000000000000000000000..1cc91de316b19d356f28f672dc0c43db843baa9f GIT binary patch literal 596 zcmV-a0;~N80ssI2Bme+XQ$aES1ONa50096s2%7=D;yhm!FQ;4y*rn2U7#QVbb`+sR zAFZ6StdSF4xjq#!)UL{}t3N=A%tw0w+j}5pX?w`r87%2KxvzNH>($>b)_en}|*l#UM`Vg<|6EoUHSzBAA(>Zi=ln+<2N~{pMHw=@lOfYp~eEy+B zF9BdXzl+sU>*P5El7}jKVvU*`SJTx|qD%!}9-JW;d(JW2!40&A5Rd3HjC*~lXl|O2 zdKmV7ajc?RUnkBd2R4hiY^?zD2*>jxMvG-1YkRW7&wv=CQ}hF z&5@lefB-PTP`IWi4CB}3dYTVvaV6^%EC0tt54-3;qO*P1S+p(scAjA198V93>Tl&@ zwC57`=i|qbLApUftlEF3ImC=iT;5v`cM;$>$#_XM%Km4+cYB!UA!Jqp!_JUqd`pzs zk3?%?!>$iLvN@a{klo0CVmZfMQZQr>CTR4@JVIAbPOn4QgK@OM9V4{@FpETe?pNVl z`Dw54o~k5i3ykQL&Oh#9