Snapstart Minimal API Performance Improvements

.autover/changes/38c5bace-4ca5-4f83-8094-ae6d912ca20a.json

@@ -0,0 +1,11 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.AspNetCoreServer.Hosting",
+      "Type": "Patch",
+      "ChangelogMessages": [
+        "Add overrideable method GetBeforeSnapshotRequests() and AddAWSLambdaBeforeSnapshotRequest() extension method to support warming up the asp.net/lambda pipelines automatically during BeforeSnapshot callback."
+      ]
+    }
+  ]
+}

Libraries/Libraries.sln

@@ -139,6 +139,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Amazon.Lambda.DynamoDBEvent
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Amazon.Lambda.DynamoDBEvents.SDK.Convertor.Tests", "test\Amazon.Lambda.DynamoDBEvents.SDK.Convertor.Tests\Amazon.Lambda.DynamoDBEvents.SDK.Convertor.Tests.csproj", "{074DB940-82BA-47D4-B888-C213D4220A82}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Amazon.Lambda.AspNetCoreServer.Hosting.Tests", "test\Amazon.Lambda.AspNetCoreServer.Hosting.Tests\Amazon.Lambda.AspNetCoreServer.Hosting.Tests.csproj", "{D61CBB71-17AB-4EC2-8C6A-70E9D7C60526}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -381,6 +383,10 @@ Global
 		{074DB940-82BA-47D4-B888-C213D4220A82}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{074DB940-82BA-47D4-B888-C213D4220A82}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{074DB940-82BA-47D4-B888-C213D4220A82}.Release|Any CPU.Build.0 = Release|Any CPU
+		{D61CBB71-17AB-4EC2-8C6A-70E9D7C60526}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{D61CBB71-17AB-4EC2-8C6A-70E9D7C60526}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{D61CBB71-17AB-4EC2-8C6A-70E9D7C60526}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{D61CBB71-17AB-4EC2-8C6A-70E9D7C60526}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -449,6 +455,7 @@ Global
 		{A699E183-D0D4-4F26-A0A7-88DA5607F455} = {1DE4EE60-45BA-4EF7-BE00-B9EB861E4C69}
 		{3400F4E9-BA12-4D3D-9BA1-2798AA8D0AFC} = {AAB54E74-20B1-42ED-BC3D-CE9F7BC7FD12}
 		{074DB940-82BA-47D4-B888-C213D4220A82} = {1DE4EE60-45BA-4EF7-BE00-B9EB861E4C69}
+		{D61CBB71-17AB-4EC2-8C6A-70E9D7C60526} = {1DE4EE60-45BA-4EF7-BE00-B9EB861E4C69}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {503678A4-B8D1-4486-8915-405A3E9CF0EB}

Libraries/src/Amazon.Lambda.AspNetCoreServer.Hosting/Internal/GetBeforeSnapshotRequestsCollector.cs

@@ -0,0 +1,16 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Amazon.Lambda.AspNetCoreServer.Hosting.Internal;
+
+/// <summary>
+/// Helper class for storing Requests for
+/// <see cref="ServiceCollectionExtensions.AddAWSLambdaBeforeSnapshotRequest(IServiceCollection, Func{IEnumerable{HttpRequestMessage}})"/>
+/// </summary>
+internal class GetBeforeSnapshotRequestsCollector
+{
+    public Func<IEnumerable<HttpRequestMessage>> Requests { get; set; }
+}
+

Libraries/src/Amazon.Lambda.AspNetCoreServer.Hosting/Internal/LambdaRuntimeSupportServer.cs

@@ -1,9 +1,8 @@
-using Amazon.Lambda.AspNetCoreServer.Internal;
+using System.Diagnostics.CodeAnalysis;
+using Amazon.Lambda.AspNetCoreServer.Internal;
 using Amazon.Lambda.Core;
 using Amazon.Lambda.RuntimeSupport;
-using Amazon.Lambda.Serialization.SystemTextJson;
 using Microsoft.AspNetCore.Hosting.Server;
-using Microsoft.AspNetCore.Mvc.ApplicationParts;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Amazon.Lambda.AspNetCoreServer.Hosting.Internal
@@ -16,7 +15,8 @@ namespace Amazon.Lambda.AspNetCoreServer.Hosting.Internal
     /// </summary>
     public abstract class LambdaRuntimeSupportServer : LambdaServer
     {
-        IServiceProvider _serviceProvider;
+        private readonly IServiceProvider _serviceProvider;
+
         internal ILambdaSerializer Serializer;
 
         /// <summary>
@@ -26,6 +26,7 @@ public abstract class LambdaRuntimeSupportServer : LambdaServer
         public LambdaRuntimeSupportServer(IServiceProvider serviceProvider)
         {
             _serviceProvider = serviceProvider;
+
             Serializer = serviceProvider.GetRequiredService<ILambdaSerializer>();
         }
 
@@ -41,6 +42,7 @@ public override Task StartAsync<TContext>(IHttpApplication<TContext> application
             base.StartAsync(application, cancellationToken);
 
             var handlerWrapper = CreateHandlerWrapper(_serviceProvider);
+
             var bootStrap = new LambdaBootstrap(handlerWrapper);
             return bootStrap.RunAsync();
         }
@@ -83,14 +85,30 @@ protected override HandlerWrapper CreateHandlerWrapper(IServiceProvider serviceP
         /// </summary>
         public class APIGatewayHttpApiV2MinimalApi : APIGatewayHttpApiV2ProxyFunction
         {
+            private readonly IEnumerable<GetBeforeSnapshotRequestsCollector> _beforeSnapshotRequestsCollectors;
+
             /// <summary>
             /// Create instances
             /// </summary>
             /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
             public APIGatewayHttpApiV2MinimalApi(IServiceProvider serviceProvider)
                 : base(serviceProvider)
             {
+                _beforeSnapshotRequestsCollectors = serviceProvider.GetServices<GetBeforeSnapshotRequestsCollector>();
             }
+
+            #if NET8_0_OR_GREATER
+            protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests()
+            {
+                foreach (var collector in _beforeSnapshotRequestsCollectors)
+                {
+                    foreach (var req in collector.Requests())
+                    {
+                        yield return req;
+                    }
+                }
+            }
+            #endif
         }
     }
 
@@ -124,14 +142,30 @@ protected override HandlerWrapper CreateHandlerWrapper(IServiceProvider serviceP
         /// </summary>
         public class APIGatewayRestApiMinimalApi : APIGatewayProxyFunction
         {
+            private readonly IEnumerable<GetBeforeSnapshotRequestsCollector> _beforeSnapshotRequestsCollectors;
+
             /// <summary>
             /// Create instances
             /// </summary>
             /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
             public APIGatewayRestApiMinimalApi(IServiceProvider serviceProvider)
                 : base(serviceProvider)
             {
+                _beforeSnapshotRequestsCollectors = serviceProvider.GetServices<GetBeforeSnapshotRequestsCollector>();
             }
+
+            #if NET8_0_OR_GREATER
+            protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests()
+            {
+                foreach (var collector in _beforeSnapshotRequestsCollectors)
+                {
+                    foreach (var req in collector.Requests())
+                    {
+                        yield return req;
+                    }
+                }
+            }
+            #endif
         }
     }
 
@@ -165,14 +199,30 @@ protected override HandlerWrapper CreateHandlerWrapper(IServiceProvider serviceP
         /// </summary>
         public class ApplicationLoadBalancerMinimalApi : ApplicationLoadBalancerFunction
         {
+            private readonly IEnumerable<GetBeforeSnapshotRequestsCollector> _beforeSnapshotRequestsCollectors;
+
             /// <summary>
             /// Create instances
             /// </summary>
             /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
             public ApplicationLoadBalancerMinimalApi(IServiceProvider serviceProvider)
                 : base(serviceProvider)
             {
+                _beforeSnapshotRequestsCollectors = serviceProvider.GetServices<GetBeforeSnapshotRequestsCollector>();
+            }
+
+#if NET8_0_OR_GREATER
+            protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests()
+            {
+                foreach (var collector in _beforeSnapshotRequestsCollectors)
+                {
+                    foreach (var req in collector.Requests())
+                    {
+                        yield return req;
+                    }
+                }
             }
+#endif
         }
     }
-}
+}

Libraries/src/Amazon.Lambda.AspNetCoreServer.Hosting/ServiceCollectionExtensions.cs

@@ -1,4 +1,4 @@
-using Amazon.Lambda.AspNetCoreServer.Hosting;
+using Amazon.Lambda.AspNetCoreServer.Hosting;
 using Amazon.Lambda.AspNetCoreServer.Internal;
 using Amazon.Lambda.AspNetCoreServer.Hosting.Internal;
 using Amazon.Lambda.Core;
@@ -88,6 +88,73 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
             return services;
         }
 
+        #if NET8_0_OR_GREATER
+        /// <summary>
+        /// Adds a function meant to initialize the asp.net and lambda pipelines during <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>
+        /// improving the performance gains offered by SnapStart.
+        /// <para />
+        /// Pass a function with one or more <see cref="HttpRequestMessage"/>s that will be used to invoke
+        /// Routes in your lambda function.  The returned <see cref="HttpRequestMessage"/> must have a relative
+        /// <see cref="HttpRequestMessage.RequestUri"/>.
+        /// <para />.
+        /// Be aware that this will invoke your applications function handler code
+        /// multiple times.  Additionally, it uses a mock <see cref="ILambdaContext"/>
+        /// which may not be fully populated.
+        /// <para />
+        /// This method automatically registers with <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>.
+        /// <para />
+        /// If SnapStart is not enabled, then this method is ignored and <paramref name="beforeSnapStartRequests"/> is never invoked.
+        /// <para />
+        /// Example:
+        /// <para />
+        /// <code>
+        /// <![CDATA[
+        /// // Example Minimal Api
+        /// var builder = WebApplication.CreateSlimBuilder(args);
+        ///
+        /// builder.Services.AddAWSLambdaHosting(LambdaEventSource.HttpApi);
+        /// 
+        /// // Initialize asp.net pipeline before Snapshot
+        /// builder.Services.AddAWSLambdaBeforeSnapshotRequest(() => [
+        ///     new HttpRequestMessage(HttpMethod.Get, "/test")
+        /// });
+        /// 
+        /// var app = builder.Build();
+        /// 
+        /// app.MapGet("/test", () => "Success");
+        /// 
+        /// app.Run(); 
+        /// ]]>
+        /// </code>
+        /// </summary>
+        #else
+        /// <summary>
+        /// Snapstart requires your application to target .NET 8 or above.
+        /// </summary>
+        #endif
+        public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, Func<IEnumerable<HttpRequestMessage>> beforeSnapStartRequests)
+        {
+            #if NET8_0_OR_GREATER
+            services.AddSingleton(new GetBeforeSnapshotRequestsCollector
+            {
+                Requests = beforeSnapStartRequests
+            });
+
+            #endif
+
+            return services;
+        }
+
+        /// <inheritdoc cref="AddAWSLambdaBeforeSnapshotRequest(IServiceCollection,Func{IEnumerable{HttpRequestMessage}})"/>
+        public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, Func<HttpRequestMessage> beforeSnapStartRequest)
+        {
+            return AddAWSLambdaBeforeSnapshotRequest(services,
+                () => new List<HttpRequestMessage>
+                {
+                    beforeSnapStartRequest()
+                });
+        }
+
         private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSource eventSource, Action<HostingOptions>? configure, out HostingOptions? hostingOptions)
         {
             hostingOptions = null;