Update the Minimal API extension method to use GetBeforeSnapshotRequests.

Author: ppittle

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

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

Libraries/src/Amazon.Lambda.AspNetCoreServer.Hosting/Amazon.Lambda.AspNetCoreServer.Hosting.csproj

@@ -27,8 +27,6 @@
     <ProjectReference Include="..\Amazon.Lambda.APIGatewayEvents\Amazon.Lambda.APIGatewayEvents.csproj" />
     <ProjectReference Include="..\Amazon.Lambda.AspNetCoreServer\Amazon.Lambda.AspNetCoreServer.csproj" />
     <ProjectReference Include="..\Amazon.Lambda.RuntimeSupport\Amazon.Lambda.RuntimeSupport.csproj" />
-    <ProjectReference Include="..\Amazon.Lambda.Serialization.SystemTextJson\Amazon.Lambda.Serialization.SystemTextJson.csproj" />
-
-   </ItemGroup>
+  </ItemGroup>
 
 </Project>

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"/>
+/// </summary>
+internal class GetBeforeSnapshotRequestsCollector
+{
+    public HttpRequestMessage? Requests { get; set; }
+}
+

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

@@ -17,10 +17,6 @@ public abstract class LambdaRuntimeSupportServer : LambdaServer
     {
         private readonly IServiceProvider _serviceProvider;
 
-        #if NET8_0_OR_GREATER
-        private readonly LambdaSnapstartExecuteRequestsBeforeSnapshotHelper _snapstartInitHelper;
-        #endif
-
         internal ILambdaSerializer Serializer;
 
         /// <summary>
@@ -31,10 +27,6 @@ public LambdaRuntimeSupportServer(IServiceProvider serviceProvider)
         {
             _serviceProvider = serviceProvider;
 
-            #if NET8_0_OR_GREATER
-            _snapstartInitHelper = _serviceProvider.GetRequiredService<LambdaSnapstartExecuteRequestsBeforeSnapshotHelper>();
-            #endif
-
             Serializer = serviceProvider.GetRequiredService<ILambdaSerializer>();
         }
 
@@ -51,12 +43,6 @@ public override Task StartAsync<TContext>(IHttpApplication<TContext> application
 
             var handlerWrapper = CreateHandlerWrapper(_serviceProvider);
 
-            #if NET8_0_OR_GREATER
-
-            _snapstartInitHelper.RegisterInitializerRequests(handlerWrapper);
-
-            #endif
-
             var bootStrap = new LambdaBootstrap(handlerWrapper);
             return bootStrap.RunAsync();
         }
@@ -99,14 +85,26 @@ 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)
+                    if (collector.Requests != null)
+                        yield return collector.Requests;
             }
+            #endif
         }
     }
 
@@ -140,14 +138,26 @@ 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)
+                    if (collector.Requests != null)
+                        yield return collector.Requests;
+            }
+            #endif
         }
     }
 
@@ -181,14 +191,26 @@ 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)
+                    if (collector.Requests != null)
+                        yield return collector.Requests;
             }
+            #endif
         }
     }
 }

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

@@ -90,20 +90,22 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
 
         #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.
+        /// Adds a <see cref="HttpRequestMessage"/>> that will be used to invoke
+        /// Routes in your lambda function in order to initialize the asp.net and lambda pipelines
+        /// during <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>.  This improves 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
+        /// 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.
+        /// 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 />
-        /// If SnapStart is not enabled, then this method is ignored and <paramref name="beforeSnapStartRequests"/> is never invoked.
+        /// This method can be called multiple times.
         /// <para />
         /// Example:
         /// <para />
@@ -115,9 +117,9 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
         /// builder.Services.AddAWSLambdaHosting(LambdaEventSource.HttpApi);
         /// 
         /// // Initialize asp.net pipeline before Snapshot
-        /// builder.Services.AddAWSLambdaBeforeSnapshotRequest(() => [
+        /// builder.Services.AddAWSLambdaBeforeSnapshotRequest(
         ///     new HttpRequestMessage(HttpMethod.Get, "/test")
-        /// });
+        /// );
         /// 
         /// var app = builder.Build();
         /// 
@@ -127,29 +129,16 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
         /// ]]>
         /// </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)
+        public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, HttpRequestMessage beforeSnapStartRequests)
         {
-            #if NET8_0_OR_GREATER
-            LambdaSnapstartExecuteRequestsBeforeSnapshotHelper.Registrar.Register(beforeSnapStartRequests);
-            #endif
+            services.AddSingleton(new GetBeforeSnapshotRequestsCollector
+            {
+                Requests = beforeSnapStartRequests
+            });
 
             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()
-                });
-        }
+        #endif
 
         private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSource eventSource, Action<HostingOptions>? configure, out HostingOptions? hostingOptions)
         {
@@ -174,11 +163,6 @@ private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSourc
 
             Utilities.EnsureLambdaServerRegistered(services, serverType);
 
-            #if NET8_0_OR_GREATER
-            // register a LambdaSnapStartInitializerHttpMessageHandler
-            services.AddSingleton<LambdaSnapstartExecuteRequestsBeforeSnapshotHelper>(new LambdaSnapstartExecuteRequestsBeforeSnapshotHelper(eventSource));
-            #endif
-
             return true;
         }
     }