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 000000000..1cc91de31
Binary files /dev/null and b/buildtools/snapshotrestore.snk differ