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(IServiceCollection, Func{IEnumerable{HttpRequestMessage}})"/>
+/// </summary>
+internal class GetBeforeSnapshotRequestsCollector
+{
+    public Func<IEnumerable<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,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
         }
     }
 
@@ -140,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
         }
     }
 
@@ -181,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/Internal/LambdaSnapstartExecuteRequestsBeforeSnapshotHelper.cs

@@ -135,7 +135,11 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
         public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, Func<IEnumerable<HttpRequestMessage>> beforeSnapStartRequests)
         {
             #if NET8_0_OR_GREATER
-            LambdaSnapstartExecuteRequestsBeforeSnapshotHelper.Registrar.Register(beforeSnapStartRequests);
+            services.AddSingleton(new GetBeforeSnapshotRequestsCollector
+            {
+                Requests = beforeSnapStartRequests
+            });
+
             #endif
 
             return services;
@@ -174,11 +178,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;
         }
     }