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

Author: ppittle

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

@@ -45,24 +45,10 @@ public void RegisterInitializerRequests(HandlerWrapper handlerWrapper)
     {
         Amazon.Lambda.Core.SnapshotRestore.RegisterBeforeSnapshot(async () =>
         {
-            // Construct specialized HttpClient that will intercept requests and saved them inside
-            // LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequests.
-            //
-            // They will be processed later by SnapstartHelperLambdaRequests.ExecuteSnapstartInitRequests which will
-            // route them correctly through a simulated lambda pipeline.
-            var messageHandlerThatCollectsRequests = new LambdaSnapstartInitializerHttpMessageHandler(_lambdaEventSource);
-
-            var httpClientThatCollectsRequests = new HttpClient(messageHandlerThatCollectsRequests);
-            httpClientThatCollectsRequests.BaseAddress = LambdaSnapstartInitializerHttpMessageHandler.BaseUri;
-
-            // "Invoke" each registered request function.  Requests will be captured inside.
-            // LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequests.
-            await Registrar.Execute(httpClientThatCollectsRequests);
-
-            // Request are now in CapturedHttpRequests.  Serialize each one into a json object
-            // and execute the request through the lambda pipeline (ie handlerWrapper).
-            foreach (var json in LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequestsJson)
+            foreach (var req in Registrar.GetAllRequests())
             {
+                var json = await SnapstartHelperLambdaRequests.SerializeToJson(req, _lambdaEventSource);
+
                 await SnapstartHelperLambdaRequests.ExecuteSnapstartInitRequests(json, times: 5, handlerWrapper);
             }
         });
@@ -74,17 +60,19 @@ public void RegisterInitializerRequests(HandlerWrapper handlerWrapper)
 
     internal class BeforeSnapstartRequestRegistrar
     {
-        private readonly List<Func<HttpClient, Task>> _beforeSnapstartFuncs = new();
+        private readonly List<Func<IEnumerable<HttpRequestMessage>>> _beforeSnapstartRequests = new();
 
-        public void Register(Func<HttpClient, Task> beforeSnapstartRequest)
+        
+        public void Register(Func<IEnumerable<HttpRequestMessage>> beforeSnapstartRequests)
         {
-            _beforeSnapstartFuncs.Add(beforeSnapstartRequest);
+            _beforeSnapstartRequests.Add(beforeSnapstartRequests);
         }
 
-        internal async Task Execute(HttpClient client)
+        internal IEnumerable<HttpRequestMessage> GetAllRequests()
         {
-            foreach (var f in _beforeSnapstartFuncs)
-                await f(client);
+            foreach (var batch in _beforeSnapstartRequests)
+                foreach (var r in batch())
+                    yield return r;
         }
     }
 
@@ -113,6 +101,8 @@ private class HelperLambdaContext : ILambdaContext, ICognitoIdentity, IClientCon
 
     private static class SnapstartHelperLambdaRequests
     {
+        private static readonly Uri _baseUri = new Uri("http://localhost");
+
         public static async Task ExecuteSnapstartInitRequests(string jsonRequest, int times, HandlerWrapper handlerWrapper)
         {
             var dummyRequest = new InvocationRequest(
@@ -131,25 +121,21 @@ public static async Task ExecuteSnapstartInitRequests(string jsonRequest, int ti
                 }
             }
         }
-    }
-
-    private class LambdaSnapstartInitializerHttpMessageHandler : HttpMessageHandler
-    {
-        private readonly LambdaEventSource _lambdaEventSource;
-
-        public static Uri BaseUri { get; } = new Uri("http://localhost");
-
-        public static List<string> CapturedHttpRequestsJson { get; } = new();
-
-        public LambdaSnapstartInitializerHttpMessageHandler(LambdaEventSource lambdaEventSource)
-        {
-            _lambdaEventSource = lambdaEventSource;
-        }
 
-        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+        public static async Task<string> SerializeToJson(HttpRequestMessage request, LambdaEventSource lambdaType)
         {
             if (null == request.RequestUri)
-                return new HttpResponseMessage(HttpStatusCode.OK);
+            {
+                throw new ArgumentException($"{nameof(HttpRequestMessage.RequestUri)} must be set.", nameof(request));
+            }
+
+            if (request.RequestUri.IsAbsoluteUri)
+            {
+                throw new ArgumentException($"{nameof(HttpRequestMessage.RequestUri)} must be relative.", nameof(request));
+            }
+
+            // make request absolut (relative to localhost) otherwise parsing the query will not work
+            request.RequestUri = new Uri(_baseUri, request.RequestUri);
 
             var duckRequest = new
             {
@@ -160,12 +146,12 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                         kvp => kvp.Value.FirstOrDefault(),
                         StringComparer.OrdinalIgnoreCase),
                 HttpMethod = request.Method.ToString(),
-                Path = "/" + BaseUri.MakeRelativeUri(request.RequestUri),
+                Path = "/" + _baseUri.MakeRelativeUri(request.RequestUri),
                 RawQuery = request.RequestUri?.Query,
                 Query = QueryHelpers.ParseNullableQuery(request.RequestUri?.Query)
             };
-            
-            string translatedRequestJson = _lambdaEventSource switch
+
+            string translatedRequestJson = lambdaType switch
             {
                 LambdaEventSource.ApplicationLoadBalancer =>
                     JsonSerializer.Serialize(
@@ -213,17 +199,13 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                         },
                         LambdaRequestTypeClasses.Default.APIGatewayProxyRequest),
                 _ => throw new NotImplementedException(
-                    $"Unknown {nameof(LambdaEventSource)}: {Enum.GetName(_lambdaEventSource)}")
+                    $"Unknown {nameof(LambdaEventSource)}: {Enum.GetName(lambdaType)}")
             };
 
-            // NOTE: Any object added to CapturedHttpRequests must have it's type added
-            // to the 
-            CapturedHttpRequestsJson.Add(translatedRequestJson);
-
-            return new HttpResponseMessage(HttpStatusCode.OK);
+            return translatedRequestJson;
         }
 
-        private async Task<string> ReadContent(HttpRequestMessage r)
+        private static async Task<string> ReadContent(HttpRequestMessage r)
         {
             if (r.Content == null)
                 return string.Empty;

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

@@ -88,18 +88,22 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
             return services;
         }
 
-#if NET8_0_OR_GREATER
+        #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 />
-        /// Use the passed <see cref="HttpClient"/> to invoke one or more Routes in your lambda function.
+        /// 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.  It uses a mock <see cref="ILambdaContext"/> which may not be fully populated.
+        /// multiple times.  Additionally, tt 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="beforeSnapStartRequest"/> is never invoked.
+        /// If SnapStart is not enabled, then this method is ignored and <paramref name="beforeSnapStartRequests"/> is never invoked.
         /// <para />
         /// Example:
         /// <para />
@@ -111,14 +115,13 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
         /// builder.Services.AddAWSLambdaHosting(LambdaEventSource.HttpApi);
         /// 
         /// // Initialize asp.net pipeline before Snapshot
-        /// builder.Services.AddAWSLambdaBeforeSnapshotRequest(async httpClient =>
-        /// {
-        ///     await httpClient.GetAsync($"/test");
+        /// builder.Services.AddAWSLambdaBeforeSnapshotRequest(() => [
+        ///     new HttpRequestMessage(HttpMethod.Get, "/test")
         /// });
         /// 
         /// var app = builder.Build();
         /// 
-        /// app.MapGet($"/test", () => "Success");
+        /// app.MapGet("/test", () => "Success");
         /// 
         /// app.Run(); 
         /// ]]>
@@ -129,16 +132,25 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
         /// Snapstart requires your application to target .NET 8 or above.
         /// </summary>
         #endif
-        public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, Func<HttpClient, Task> beforeSnapStartRequest)
+        public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, Func<IEnumerable<HttpRequestMessage>> beforeSnapStartRequests)
         {
-
             #if NET8_0_OR_GREATER
-            LambdaSnapstartExecuteRequestsBeforeSnapshotHelper.Registrar.Register(beforeSnapStartRequest);
+            LambdaSnapstartExecuteRequestsBeforeSnapshotHelper.Registrar.Register(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;

Libraries/test/Amazon.Lambda.AspNetCoreServer.Hosting.Tests/AddAWSLambdaBeforeSnapshotRequestTests.cs

@@ -30,10 +30,9 @@ public async Task VerifyCallbackIsInvoked(LambdaEventSource hostingType)
 
         builder.Services.AddAWSLambdaHosting(hostingType);
         // Initialize asp.net pipeline before Snapshot
-        builder.Services.AddAWSLambdaBeforeSnapshotRequest(async httpClient =>
-        {
-            await httpClient.GetAsync($"/test");
-        });
+        builder.Services.AddAWSLambdaBeforeSnapshotRequest(() => 
+             new HttpRequestMessage(HttpMethod.Get, "/test")
+        );
 
         var app = builder.Build();