WIP - add Snapstart helper to AbstractAspNetCoreFunction

Author: ppittle

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;
@@ -251,6 +253,17 @@ protected virtual IHostBuilder CreateHostBuilder()
             return builder;
         }
 
+        /// <summary>
+        /// TODO - copy from AddAWSLambdaBeforeSnapshotRequest
+        /// Override to register a <see cref="HttpRequestMessage"/> that will be executed
+        /// within <see cref="SnapshotRestore.RegisterBeforeSnapshot"/>.
+        ///
+        /// Improves performance of SnapStart
+        /// </summary>
+        /// <returns></returns>
+        protected virtual IEnumerable<TREQUEST> RegisterBeforeSnapshotRequest() => Enumerable.Empty<TREQUEST>();
+        
+
         private protected bool IsStarted
         {
             get
@@ -284,6 +297,32 @@ protected void Start()
                         "instead of ConfigureWebHostDefaults to make sure the property Lambda services are registered.");
             }
             _logger = ActivatorUtilities.CreateInstance<Logger<AbstractAspNetCoreFunction<TREQUEST, TRESPONSE>>>(this._hostServices);
+
+            #if NET8_0_OR_GREATER
+            var beforeSnapstartRequests = RegisterBeforeSnapshotRequest();
+
+            Amazon.Lambda.Core.SnapshotRestore.RegisterBeforeSnapshot(async () =>
+            {
+                foreach (var request in beforeSnapstartRequests)
+                {
+                    var invokeTimes = 5;
+
+                    InvokeFeatures features = new InvokeFeatures();
+                    MarshallRequest(features, request, new SnapStartEmptyLambdaContext());
+
+                    var context = CreateContext(features);
+
+                    var lambdaContext = new SnapStartEmptyLambdaContext();
+
+                    for (var i = 0; i < invokeTimes; i++)
+                    {
+                        await ProcessRequest(lambdaContext, context, features);
+                    }
+                }
+            });
+
+            #endif
+
         }
 
         /// <summary>

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

@@ -31,6 +31,7 @@
 
   <ItemGroup>
     <FrameworkReference Include="Microsoft.AspNetCore.App" />
+    <ProjectReference Include="..\Amazon.Lambda.RuntimeSupport\Amazon.Lambda.RuntimeSupport.csproj" />
     <ProjectReference Include="..\Amazon.Lambda.Serialization.SystemTextJson\Amazon.Lambda.Serialization.SystemTextJson.csproj" />
   </ItemGroup>
 

Libraries/src/Amazon.Lambda.AspNetCoreServer/Internal/SnapStartEmptyLambdaContext.cs

@@ -0,0 +1,32 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System;
+using System.Collections.Generic;
+using Amazon.Lambda.Core;
+using Amazon.Lambda.RuntimeSupport;
+
+namespace Amazon.Lambda.AspNetCoreServer.Internal;
+
+internal class SnapStartEmptyLambdaContext : ILambdaContext, ICognitoIdentity, IClientContext
+{
+    private LambdaEnvironment _lambdaEnvironment = new();
+
+    public string TraceId => string.Empty;
+    public string AwsRequestId => string.Empty;
+    public IClientContext ClientContext => this;
+    public string FunctionName => _lambdaEnvironment.FunctionName;
+    public string FunctionVersion => _lambdaEnvironment.FunctionVersion;
+    public ICognitoIdentity Identity => this;
+    public string InvokedFunctionArn => string.Empty;
+    public ILambdaLogger Logger => null;
+    public string LogGroupName => _lambdaEnvironment.LogGroupName;
+    public string LogStreamName => _lambdaEnvironment.LogStreamName;
+    public int MemoryLimitInMB => 128;
+    public TimeSpan RemainingTime => TimeSpan.FromMilliseconds(100);
+    public string IdentityId { get; }
+    public string IdentityPoolId { get; }
+    public IDictionary<string, string> Environment { get; } = new Dictionary<string, string>();
+    public IClientApplication Client { get; }
+    public IDictionary<string, string> Custom { get; } = new Dictionary<string, string>();
+}