Snapstart Minimal API Performance Improvements (#2010)

* Add AddAWSLambdaBeforeSnapshotRequest to support warming up the asp.net/lambda pipelines automatically during BeforeSnapshot callback.

* inline SnapstartHelperLambdaRequest

* add auto version change

* add support for additional LambdaEventSource

* update documentation

* Fix bugs in LambdaSnapstartExecuteRequestsBeforeSnapshotHelper and add unit tests

* removed duplicate reference to snk

* update AddAWSLambdaBeforeSnapshotRequest() to support AOT by optimizing json serialization, adjust #if blocks and fix simple warnings

* remove requirement that Amazon.Lambda.RuntimeSupport make InternalsVisibileTo Amazon.Lambda.AspNetCorServer.Hosting

* migrate extension method to use HttpRequestMethods instead of collecting calls to HttpClient

* Add Snapstart helper to AbstractAspNetCoreFunction

* rename RegisterBeforeSnapshotRequest to GetBeforeSnapshotRequests

* optimize HttpRequestMessageSerializer and remove AspNetCoreServers's dependency on RuntimeSupport

* Update the Minimal API extension method to use GetBeforeSnapshotRequests.

* Add additional test for asp.net workloads

* improve documentation

* impvoe unit tests

Author: ppittle

.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,17 @@
+// 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;
+
+#if NET8_0_OR_GREATER
+/// <summary>
+/// Helper class for storing Requests for
+/// <see cref="ServiceCollectionExtensions.AddAWSLambdaBeforeSnapshotRequest"/>
+/// </summary>
+internal class GetBeforeSnapshotRequestsCollector
+{
+    public HttpRequestMessage? Request { get; set; }
+}
+#endif

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
         {
+            #if NET8_0_OR_GREATER
+            private readonly IEnumerable<GetBeforeSnapshotRequestsCollector> _beforeSnapshotRequestsCollectors;
+            #endif
+
             /// <summary>
             /// Create instances
             /// </summary>
             /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
             public APIGatewayHttpApiV2MinimalApi(IServiceProvider serviceProvider)
                 : base(serviceProvider)
             {
+                #if NET8_0_OR_GREATER
+                _beforeSnapshotRequestsCollectors = serviceProvider.GetServices<GetBeforeSnapshotRequestsCollector>();
+                #endif
             }
+
+            #if NET8_0_OR_GREATER
+            protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests()
+            {
+                foreach (var collector in _beforeSnapshotRequestsCollectors)
+                    if (collector.Request != null)
+                        yield return collector.Request;
+            }
+            #endif
         }
     }
 
@@ -124,14 +142,30 @@ protected override HandlerWrapper CreateHandlerWrapper(IServiceProvider serviceP
         /// </summary>
         public class APIGatewayRestApiMinimalApi : APIGatewayProxyFunction
         {
+            #if NET8_0_OR_GREATER
+            private readonly IEnumerable<GetBeforeSnapshotRequestsCollector> _beforeSnapshotRequestsCollectors;
+            #endif
+
             /// <summary>
             /// Create instances
             /// </summary>
             /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
             public APIGatewayRestApiMinimalApi(IServiceProvider serviceProvider)
                 : base(serviceProvider)
             {
+                #if NET8_0_OR_GREATER
+                _beforeSnapshotRequestsCollectors = serviceProvider.GetServices<GetBeforeSnapshotRequestsCollector>();
+                #endif
             }
+
+            #if NET8_0_OR_GREATER
+            protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests()
+            {
+                foreach (var collector in _beforeSnapshotRequestsCollectors)
+                    if (collector.Request != null)
+                        yield return collector.Request;
+            }
+            #endif
         }
     }
 
@@ -165,14 +199,30 @@ protected override HandlerWrapper CreateHandlerWrapper(IServiceProvider serviceP
         /// </summary>
         public class ApplicationLoadBalancerMinimalApi : ApplicationLoadBalancerFunction
         {
+            #if NET8_0_OR_GREATER
+            private readonly IEnumerable<GetBeforeSnapshotRequestsCollector> _beforeSnapshotRequestsCollectors;
+            #endif
+
             /// <summary>
             /// Create instances
             /// </summary>
             /// <param name="serviceProvider">The IServiceProvider created for the ASP.NET Core application</param>
             public ApplicationLoadBalancerMinimalApi(IServiceProvider serviceProvider)
                 : base(serviceProvider)
             {
+                #if NET8_0_OR_GREATER
+                _beforeSnapshotRequestsCollectors = serviceProvider.GetServices<GetBeforeSnapshotRequestsCollector>();
+                #endif
+            }
+
+            #if NET8_0_OR_GREATER
+            protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests()
+            {
+                foreach (var collector in _beforeSnapshotRequestsCollectors)
+                    if (collector.Request != null)
+                        yield return collector.Request;
             }
+            #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,62 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
             return services;
         }
 
+        #if NET8_0_OR_GREATER
+        /// <summary>
+        /// Adds a <see cref="HttpRequestMessage"/>> that will be used to invoke
+        /// Routes in your lambda function in order to initialize the ASP.NET Core and Lambda pipelines
+        /// during <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>.  This improves the performance gains
+        /// offered by SnapStart.
+        /// <para />
+        /// <paramref name="beforeSnapStartRequest"/> must have a relative
+        /// <see cref="HttpRequestMessage.RequestUri"/> and the <see cref="HttpRequestMessage.Content"/> only supports
+        /// text based payload.
+        /// <para />.
+        /// Be aware that this will invoke your applications function handler code
+        /// multiple times so that .NET runtime sees this code is a hot path and should be optimized.
+        /// <para />
+        /// When the function handler is called as part of SnapStart warm up, the instance will use a
+        /// mock <see cref="ILambdaContext"/>, which will not be fully populated.
+        /// <para />
+        /// This method automatically registers with <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>.
+        /// <para />
+        /// This method can be called multiple times to register additional urls.
+        /// <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>
+        /// <param name="services"></param>
+        /// <param name="beforeSnapStartRequest"></param>
+        public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, HttpRequestMessage beforeSnapStartRequest)
+        {
+            services.AddSingleton(new GetBeforeSnapshotRequestsCollector
+            {
+                Request = beforeSnapStartRequest
+            });
+
+            return services;
+        }
+        #endif
+
         private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSource eventSource, Action<HostingOptions>? configure, out HostingOptions? hostingOptions)
         {
             hostingOptions = null;

Libraries/src/Amazon.Lambda.AspNetCoreServer/AbstractAspNetCoreFunction.cs

@@ -8,6 +8,8 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
+using System.Net.Http;
 using System.Reflection;
 using System.Text;
 using System.Threading.Tasks;
@@ -120,6 +122,8 @@ protected AbstractAspNetCoreFunction(IServiceProvider hostedServices)
             _hostServices = hostedServices;
             _server = this._hostServices.GetService(typeof(Microsoft.AspNetCore.Hosting.Server.IServer)) as LambdaServer;
             _logger = ActivatorUtilities.CreateInstance<Logger<AbstractAspNetCoreFunction<TREQUEST, TRESPONSE>>>(this._hostServices);
+
+            AddRegisterBeforeSnapshot();
         }
 
         /// <summary>
@@ -251,6 +255,47 @@ protected virtual IHostBuilder CreateHostBuilder()
             return builder;
         }
 
+        #if NET8_0_OR_GREATER
+        /// <summary>
+        /// Return one or more <see cref="HttpRequestMessage"/>s that will be used to invoke
+        /// Routes in your lambda function in order to initialize the ASP.NET Core and Lambda pipelines
+        /// during <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>,
+        /// improving the performance gains offered by SnapStart.
+        /// <para />
+        /// The returned <see cref="HttpRequestMessage"/>s must have a relative
+        /// <see cref="HttpRequestMessage.RequestUri"/> and the <see cref="HttpRequestMessage.Content"/> only supports
+        /// text based payload.
+        /// <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 never invoked.
+        /// <para />
+        /// Example:
+        /// <para />
+        /// <code>
+        /// <![CDATA[
+        /// public class HttpV2LambdaFunction : APIGatewayHttpApiV2ProxyFunction<Startup>
+        /// {
+        ///     protected override IEnumerable<HttpRequestMessage> RegisterBeforeSnapshotRequest() =>
+        ///     [
+        ///         new HttpRequestMessage
+        ///         {
+        ///             RequestUri = new Uri("/api/ExampleSnapstartInit"),
+        ///             Method = HttpMethod.Get
+        ///         }
+        ///     ];
+        /// }
+        /// ]]>
+        /// </code>
+        /// </summary>
+        protected virtual IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests() =>
+            Enumerable.Empty<HttpRequestMessage>();
+        #endif
+
         private protected bool IsStarted
         {
             get
@@ -259,6 +304,40 @@ private protected bool IsStarted
             }
         }
 
+        private void AddRegisterBeforeSnapshot()
+        {
+            #if NET8_0_OR_GREATER
+
+            Amazon.Lambda.Core.SnapshotRestore.RegisterBeforeSnapshot(async () =>
+            {
+                var beforeSnapstartRequests = GetBeforeSnapshotRequests();
+
+                foreach (var httpRequest in beforeSnapstartRequests)
+                {
+                    var invokeTimes = 5;
+
+                    var request = await HttpRequestMessageConverter.ConvertToLambdaRequest<TREQUEST>(httpRequest);
+
+                    InvokeFeatures features = new InvokeFeatures();
+                    (features as IItemsFeature).Items = new Dictionary<object, object>();
+                    (features as IServiceProvidersFeature).RequestServices = _hostServices;
+
+                    MarshallRequest(features, request, new SnapStartEmptyLambdaContext());
+
+                    var context = CreateContext(features);
+
+                    for (var i = 0; i < invokeTimes; i++)
+                    {
+                        var lambdaContext = new SnapStartEmptyLambdaContext();
+
+                        await ProcessRequest(lambdaContext, context, features);
+                    }
+                }
+            });
+
+            #endif
+        }
+
         /// <summary>
         /// Should be called in the derived constructor 
         /// </summary>
@@ -284,6 +363,8 @@ protected void Start()
                         "instead of ConfigureWebHostDefaults to make sure the property Lambda services are registered.");
             }
             _logger = ActivatorUtilities.CreateInstance<Logger<AbstractAspNetCoreFunction<TREQUEST, TRESPONSE>>>(this._hostServices);
+
+            AddRegisterBeforeSnapshot();
         }
 
         /// <summary>