Skip to content

Commit 220b292

Browse files
ppittleppittle
committed
Add AddAWSLambdaBeforeSnapshotRequest to support warming up the asp.net/lambda pipelines automatically during BeforeSnapshot callback.
1 parent 3430158 commit 220b292

4 files changed

Lines changed: 238 additions & 6 deletions

File tree

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

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
1-
using Amazon.Lambda.AspNetCoreServer.Internal;
1+
using System.Diagnostics.CodeAnalysis;
2+
using System.Text;
3+
using System.Text.Json;
4+
using Amazon.Lambda.AspNetCoreServer.Internal;
25
using Amazon.Lambda.Core;
36
using Amazon.Lambda.RuntimeSupport;
4-
using Amazon.Lambda.Serialization.SystemTextJson;
7+
using Amazon.Lambda.RuntimeSupport.Helpers;
58
using Microsoft.AspNetCore.Hosting.Server;
6-
using Microsoft.AspNetCore.Mvc.ApplicationParts;
79
using Microsoft.Extensions.DependencyInjection;
810

911
namespace Amazon.Lambda.AspNetCoreServer.Hosting.Internal
@@ -16,7 +18,9 @@ namespace Amazon.Lambda.AspNetCoreServer.Hosting.Internal
1618
/// </summary>
1719
public abstract class LambdaRuntimeSupportServer : LambdaServer
1820
{
19-
IServiceProvider _serviceProvider;
21+
private readonly IServiceProvider _serviceProvider;
22+
private readonly LambdaSnapstartExecuteRequestsBeforeSnapshotHelper _snapstartInitHelper;
23+
2024
internal ILambdaSerializer Serializer;
2125

2226
/// <summary>
@@ -26,6 +30,7 @@ public abstract class LambdaRuntimeSupportServer : LambdaServer
2630
public LambdaRuntimeSupportServer(IServiceProvider serviceProvider)
2731
{
2832
_serviceProvider = serviceProvider;
33+
_snapstartInitHelper = _serviceProvider.GetRequiredService<LambdaSnapstartExecuteRequestsBeforeSnapshotHelper>();
2934
Serializer = serviceProvider.GetRequiredService<ILambdaSerializer>();
3035
}
3136

@@ -36,11 +41,19 @@ public LambdaRuntimeSupportServer(IServiceProvider serviceProvider)
3641
/// <param name="application"></param>
3742
/// <param name="cancellationToken"></param>
3843
/// <returns></returns>
44+
[RequiresUnreferencedCode("_snapstartInitHelper Serializes objects to Json")]
3945
public override Task StartAsync<TContext>(IHttpApplication<TContext> application, CancellationToken cancellationToken)
4046
{
4147
base.StartAsync(application, cancellationToken);
4248

4349
var handlerWrapper = CreateHandlerWrapper(_serviceProvider);
50+
51+
#if NET8_0_OR_GREATER
52+
53+
_snapstartInitHelper.RegisterInitializerRequests(handlerWrapper);
54+
55+
#endif
56+
4457
var bootStrap = new LambdaBootstrap(handlerWrapper);
4558
return bootStrap.RunAsync();
4659
}
@@ -175,4 +188,4 @@ public ApplicationLoadBalancerMinimalApi(IServiceProvider serviceProvider)
175188
}
176189
}
177190
}
178-
}
191+
}

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

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System.Diagnostics.CodeAnalysis;
5+
using System.Net;
6+
using System.Text.Json;
7+
using Amazon.Lambda.APIGatewayEvents;
8+
using Amazon.Lambda.RuntimeSupport;
9+
using Amazon.Lambda.RuntimeSupport.Helpers;
10+
using Microsoft.Extensions.DependencyInjection;
11+
12+
namespace Amazon.Lambda.AspNetCoreServer.Hosting.Internal;
13+
14+
/// <summary>
15+
/// Contains the plumbing to register a user provided <see cref="Func{HttpClient, Task}"/> inside
16+
/// <see cref="Amazon.Lambda.Core.SnapshotRestore.RegisterBeforeSnapshot"/>.
17+
/// The function is meant to initialize the asp.net and lambda pipelines during
18+
/// <see cref="Core.SnapshotRestore.RegisterBeforeSnapshot"/> and improve the
19+
/// performance gains offered by SnapStart.
20+
/// <para />
21+
/// It works by construction a specialized <see cref="HttpClient" /> that will intercept requests
22+
/// and saved them inside <see cref="LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequests" />.
23+
/// <para />
24+
/// Intercepted requests are then be processed later by <see cref="SnapstartHelperLambdaRequests.ExecuteSnapstartInitRequests"/>
25+
/// which will route them correctly through a simulated asp.net/lambda pipeline.
26+
/// </summary>
27+
internal class LambdaSnapstartExecuteRequestsBeforeSnapshotHelper
28+
{
29+
private readonly LambdaEventSource _lambdaEventSource;
30+
31+
public LambdaSnapstartExecuteRequestsBeforeSnapshotHelper(LambdaEventSource lambdaEventSource)
32+
{
33+
_lambdaEventSource = lambdaEventSource;
34+
}
35+
36+
/// <inheritdoc cref="RegisterInitializerRequests"/>
37+
[RequiresUnreferencedCode("Serializes object to json")]
38+
public void RegisterInitializerRequests(HandlerWrapper handlerWrapper)
39+
{
40+
#if NET8_0_OR_GREATER
41+
42+
Amazon.Lambda.Core.SnapshotRestore.RegisterBeforeSnapshot(async () =>
43+
{
44+
// Construct specialized HttpClient that will intercept requests and saved them inside
45+
// LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequests.
46+
//
47+
// They will be processed later by SnapstartHelperLambdaRequests.ExecuteSnapstartInitRequests which will
48+
// route them correctly through a simulated lambda pipeline.
49+
var messageHandlerThatCollectsRequests = new LambdaSnapstartInitializerHttpMessageHandler(_lambdaEventSource);
50+
51+
var httpClientThatCollectsRequests = new HttpClient(messageHandlerThatCollectsRequests);
52+
httpClientThatCollectsRequests.BaseAddress = LambdaSnapstartInitializerHttpMessageHandler.BaseUri;
53+
54+
// "Invoke" each registered request function. Requests will be captured inside.
55+
// LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequests.
56+
await Registrar.Execute(httpClientThatCollectsRequests);
57+
58+
// Request are now in CapturedHttpRequests. Serialize each one into a json object
59+
// and execute the request through the lambda pipeline (ie handlerWrapper).
60+
foreach (var req in LambdaSnapstartInitializerHttpMessageHandler.CapturedHttpRequests)
61+
{
62+
var json = JsonSerializer.Serialize(req);
63+
64+
// TODO - inline
65+
await SnapstartHelperLambdaRequests.ExecuteSnapstartInitRequests(json, times: 5, handlerWrapper);
66+
}
67+
});
68+
69+
#endif
70+
}
71+
72+
/// <inheritdoc cref="ServiceCollectionExtensions.AddAWSLambdaBeforeSnapshotRequest"/>
73+
internal static BeforeSnapstartRequestRegistrar Registrar = new();
74+
75+
internal class BeforeSnapstartRequestRegistrar
76+
{
77+
private List<Func<HttpClient, Task>> beforeSnapstartFuncs = new();
78+
79+
public void Register(Func<HttpClient, Task> beforeSnapstartRequest)
80+
{
81+
beforeSnapstartFuncs.Add(beforeSnapstartRequest);
82+
}
83+
84+
internal async Task Execute(HttpClient client)
85+
{
86+
foreach (var f in beforeSnapstartFuncs)
87+
await f(client);
88+
}
89+
}
90+
91+
private class LambdaSnapstartInitializerHttpMessageHandler : HttpMessageHandler
92+
{
93+
private LambdaEventSource _lambdaEventSource;
94+
95+
public static Uri BaseUri { get; } = new Uri("http://localhost");
96+
97+
public static List<object> CapturedHttpRequests { get; } = new();
98+
99+
public LambdaSnapstartInitializerHttpMessageHandler(LambdaEventSource lambdaEventSource)
100+
{
101+
_lambdaEventSource = lambdaEventSource;
102+
}
103+
104+
protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
105+
{
106+
// Copy request to correct request, ie APIGatewayProxyRequest
107+
108+
// TODO - IMPLEMENT
109+
var translatedRequest = new APIGatewayProxyRequest
110+
{
111+
Path = request.RequestUri.MakeRelativeUri(BaseUri).ToString(),
112+
HttpMethod = request.Method.ToString()
113+
};
114+
115+
CapturedHttpRequests.Add(translatedRequest);
116+
117+
return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK));
118+
}
119+
}
120+
}

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

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
using Amazon.Lambda.AspNetCoreServer.Hosting;
1+
using Amazon.Lambda.AspNetCoreServer.Hosting;
22
using Amazon.Lambda.AspNetCoreServer.Internal;
33
using Amazon.Lambda.AspNetCoreServer.Hosting.Internal;
44
using Amazon.Lambda.Core;
@@ -88,6 +88,46 @@ public static IServiceCollection AddAWSLambdaHosting(this IServiceCollection ser
8888
return services;
8989
}
9090

91+
/// <summary>
92+
/// Adds a function meant to initialize the asp.net and lambda pipelines during <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>
93+
/// improving the performance gains offered by SnapStart.
94+
/// <para />
95+
/// Use the passed <see cref="HttpClient"/> to invoke one or more Routes in your lambda function.
96+
/// Be aware that this will invoke your applications function handler code
97+
/// multiple times. It uses a mock <see cref="ILambdaContext"/> which may not be fully populated.
98+
/// <para />
99+
/// This method automatically registers with <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>.
100+
/// <para />
101+
/// If SnapStart is not enabled, then this method is ignored and <paramref name="beforeSnapStartRequest"/> is never invoked.
102+
/// <para />
103+
/// Example:
104+
/// <para />
105+
/// <code>
106+
/// <![CDATA[
107+
/// // Example Minimal Api
108+
/// var builder = WebApplication.CreateSlimBuilder(args);
109+
///
110+
/// // Initialize asp.net pipeline before Snapshot
111+
/// builder.Services.AddAWSLambdaBeforeSnapshotRequest(async httpClient =>
112+
/// {
113+
/// await httpClient.GetAsync($"/test");
114+
/// });
115+
///
116+
/// var app = builder.Build();
117+
///
118+
/// app.MapGet($"/test", () => "Success");
119+
///
120+
/// app.Run();
121+
/// ]]>
122+
/// </code>
123+
/// </summary>
124+
public static IServiceCollection AddAWSLambdaBeforeSnapshotRequest(this IServiceCollection services, Func<HttpClient, Task> beforeSnapStartRequest)
125+
{
126+
LambdaSnapstartExecuteRequestsBeforeSnapshotHelper.Registrar.Register(beforeSnapStartRequest);
127+
128+
return services;
129+
}
130+
91131
private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSource eventSource, Action<HostingOptions>? configure, out HostingOptions? hostingOptions)
92132
{
93133
hostingOptions = null;
@@ -111,6 +151,9 @@ private static bool TryLambdaSetup(IServiceCollection services, LambdaEventSourc
111151

112152
Utilities.EnsureLambdaServerRegistered(services, serverType);
113153

154+
// register a LambdaSnapStartInitializerHttpMessageHandler
155+
services.AddSingleton<LambdaSnapstartExecuteRequestsBeforeSnapshotHelper>(new LambdaSnapstartExecuteRequestsBeforeSnapshotHelper(eventSource));
156+
114157
return true;
115158
}
116159
}

‎Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/SnapstartHelperLambdaRequests.cs‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.IO;
4+
using System.Text;
5+
using System.Threading.Tasks;
6+
7+
namespace Amazon.Lambda.RuntimeSupport.Helpers
8+
{
9+
/// <summary>
10+
/// should not be public
11+
/// this is public to allow LambdaRuntimeSupportServer access,
12+
/// otherwise it would need access to <see cref="InvocationRequest"/> and
13+
/// <see cref="LambdaContext"/>
14+
///
15+
/// TODO - inline to LambdaSnapstartExecuteRequestsBeforeSnapshotHelper
16+
/// </summary>
17+
public static class SnapstartHelperLambdaRequests
18+
{
19+
private static InternalLogger _logger = InternalLogger.GetDefaultLogger();
20+
21+
private static readonly RuntimeApiHeaders _fakeRuntimeApiHeaders = new(new Dictionary<string, IEnumerable<string>>
22+
{
23+
{ RuntimeApiHeaders.HeaderAwsRequestId, new List<string>() },
24+
{ RuntimeApiHeaders.HeaderTraceId, new List<string>() },
25+
{ RuntimeApiHeaders.HeaderClientContext, new List<string>() },
26+
{ RuntimeApiHeaders.HeaderCognitoIdentity, new List<string>() },
27+
{ RuntimeApiHeaders.HeaderDeadlineMs, new List<string>() },
28+
{ RuntimeApiHeaders.HeaderInvokedFunctionArn, new List<string>() },
29+
});
30+
31+
public static async Task ExecuteSnapstartInitRequests(string jsonRequest, int times, HandlerWrapper handlerWrapper)
32+
{
33+
var dummyRequest = new InvocationRequest
34+
{
35+
InputStream = new MemoryStream(Encoding.UTF8.GetBytes(jsonRequest)),
36+
LambdaContext = new LambdaContext(
37+
_fakeRuntimeApiHeaders,
38+
new LambdaEnvironment(),
39+
new SimpleLoggerWriter())
40+
};
41+
42+
for (var i = 0; i < times; i++)
43+
{
44+
try
45+
{
46+
_ = await handlerWrapper.Handler.Invoke(dummyRequest);
47+
}
48+
catch (Exception e)
49+
{
50+
Console.WriteLine("StartAsync: " + e.Message + e.StackTrace);
51+
_logger.LogError(e, "StartAsync: Custom Warmup Failure: " + e.Message + e.StackTrace);
52+
}
53+
}
54+
}
55+
}
56+
}

0 commit comments

Comments
 (0)