Use the new .NET 8 APIs to configure max heap memory size

Libraries/src/Amazon.Lambda.RuntimeSupport/Amazon.Lambda.RuntimeSupport.csproj

@@ -3,10 +3,9 @@
   <Import Project="..\..\..\buildtools\common.props" />
 
   <PropertyGroup>
-    <TargetFrameworks>netstandard2.0;net5.0;net6.0</TargetFrameworks>
-    <VersionPrefix>1.8.8</VersionPrefix>
+    <TargetFrameworks>netstandard2.0;net5.0;net6.0;net8.0</TargetFrameworks>
+    <VersionPrefix>1.9.0</VersionPrefix>
     <Description>Provides a bootstrap and Lambda Runtime API Client to help you to develop custom .NET Core Lambda Runtimes.</Description>
-    <Description>Provides a bootstrap and Lambda Runtime API Client to help you  to develop custom .NET Core Lambda Runtimes.</Description>
     <AssemblyTitle>Amazon.Lambda.RuntimeSupport</AssemblyTitle>
     <AssemblyName>Amazon.Lambda.RuntimeSupport</AssemblyName>
     <PackageId>Amazon.Lambda.RuntimeSupport</PackageId>

Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs

@@ -114,7 +114,11 @@ private LambdaBootstrap(HttpClient httpClient, LambdaBootstrapHandler handler, L
         /// <returns>A Task that represents the operation.</returns>
         public async Task RunAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            if(UserCodeInit.IsCallPreJit())
+#if NET8_0_OR_GREATER
+            AdjustMemorySettings();
+#endif
+
+            if (UserCodeInit.IsCallPreJit())
             {
                 this._logger.LogInformation("PreJit: CultureInfo");
                 UserCodeInit.LoadStringCultureInfo();
@@ -248,6 +252,43 @@ private void WriteUnhandledExceptionToLog(Exception exception)
             Console.Error.WriteLine(exception);
         }
 
+#if NET8_0_OR_GREATER
+        /// <summary>
+        /// The .NET runtime does not recognize the memory limits placed by Lambda via Lambda's cgroups. This method is run during startup to inform the
+        /// .NET runtime the max memory configured for Lambda function. The max memory can be determined using the AWS_LAMBDA_FUNCTION_MEMORY_SIZE environment variable
+        /// which has the memory in MB.
+        /// 
+        /// For additional context on setting the heap size refer to this GitHub issue:
+        /// https://github.com/dotnet/runtime/issues/70601
+        /// </summary>
+        private void AdjustMemorySettings()
+        {
+            try
+            {
+                int lambdaMemoryInMb;
+                if (!int.TryParse(Environment.GetEnvironmentVariable(LambdaEnvironment.EnvVarFunctionMemorySize), out lambdaMemoryInMb))
+                    return;
+
+                ulong memoryInBytes = (ulong)lambdaMemoryInMb * LambdaEnvironment.OneMegabyte;
+
+                // If the user has already configured the hard heap limit to something lower then is available 
+                // then make no adjustments to honor the user's setting.
+                if ((ulong)GC.GetGCMemoryInfo().TotalAvailableMemoryBytes < memoryInBytes)
+                    return;
+
+                AppContext.SetData("GCHeapHardLimit", memoryInBytes);
+
+#pragma warning disable CA2252
+                GC.RefreshMemoryLimit();
+#pragma warning disable CA2252
+            }
+            catch(Exception ex)
+            {
+                _logger.LogError(ex, "Failed to communicate to the .NET runtime the amount of memory configured for the Lambda function via the AWS_LAMBDA_FUNCTION_MEMORY_SIZE environment variable.");
+            }
+        }
+#endif
+
         #region IDisposable Support
         private bool disposedValue = false; // To detect redundant calls
 

Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaEnvironment.cs

@@ -32,12 +32,15 @@ public class LambdaEnvironment
         internal const string EnvVarLogStreamName = "AWS_LAMBDA_LOG_STREAM_NAME";
         internal const string EnvVarServerHostAndPort = "AWS_LAMBDA_RUNTIME_API";
         internal const string EnvVarTraceId = "_X_AMZN_TRACE_ID";
+        internal const string EnvVarFunctionSize = "AWS_LAMBDA_FUNCTION_MEMORY_SIZE";
 
         internal const string AwsLambdaDotnetCustomRuntime = "AWS_Lambda_dotnet_custom";
         internal const string AmazonLambdaRuntimeSupportMarker = "amazonlambdaruntimesupport";
 
         private IEnvironmentVariables _environmentVariables;
 
+        internal const int OneMegabyte = 1024 * 1024;
+
         public LambdaEnvironment() : this(new SystemEnvironmentVariables()) { }
 
         internal LambdaEnvironment(IEnvironmentVariables environmentVariables)

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/Amazon.Lambda.RuntimeSupport.IntegrationTests.csproj

@@ -43,16 +43,20 @@
     <Exec Command="dotnet tool install -g Amazon.Lambda.Tools" IgnoreExitCode="true" />
 
 	<Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Command="dotnet restore" />
-    <Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release" />
-    <Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --function-architecture arm64" />
+    <Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release --framework net6.0" />
+    <Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --framework net6.0 --function-architecture arm64" />
+
+	<Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Command="dotnet restore" />
+    <Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release --framework net8.0" />
+    <Exec WorkingDirectory="..\CustomRuntimeFunctionTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --framework net8.0 --function-architecture arm64" />	  
 
 	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiTest" Command="dotnet restore" />
-	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release" />
-	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --function-architecture arm64" />
+	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release --framework net6.0" />
+	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --framework net6.0 --function-architecture arm64" />
 
 	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest" Command="dotnet restore" />
-	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release" />
-	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --function-architecture arm64" />	  
+	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest" Condition="'$(Architecture)'=='' or '$(Architecture)'=='x86'" Command="dotnet lambda package -c Release --framework net6.0" />
+	<Exec WorkingDirectory="..\CustomRuntimeAspNetCoreMinimalApiCustomSerializerTest" Condition="'$(Architecture)'=='arm64'" Command="dotnet lambda package -c Release --framework net6.0 --function-architecture arm64" />	  
   </Target>
 
 </Project>

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/BaseCustomRuntimeTest.cs

@@ -15,6 +15,8 @@ namespace Amazon.Lambda.RuntimeSupport.IntegrationTests
 {
     public class BaseCustomRuntimeTest
     {
+        public const int FUNCTION_MEMORY_MB = 512;
+
         protected static readonly RegionEndpoint TestRegion = RegionEndpoint.USWest2;
         protected static readonly string LAMBDA_ASSUME_ROLE_POLICY =
         @"
@@ -240,7 +242,7 @@ protected async Task CreateFunctionAsync(IAmazonLambda lambdaClient, string buck
                     S3Key = DeploymentZipKey
                 },
                 Handler = this.Handler,
-                MemorySize = 512,
+                MemorySize = FUNCTION_MEMORY_MB,
                 Timeout = 30,
                 Runtime = Runtime.Dotnet6,
                 Role = ExecutionRoleArn

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/CustomRuntimeTests.cs

@@ -26,14 +26,36 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
+using static Amazon.Lambda.RuntimeSupport.IntegrationTests.CustomRuntimeTests;
 
 namespace Amazon.Lambda.RuntimeSupport.IntegrationTests
 {
+    public class CustomRuntimeNET6Tests : CustomRuntimeTests
+    {
+        public CustomRuntimeNET6Tests()
+            : base("CustomRuntimeNET6FunctionTest-" + DateTime.Now.Ticks, "CustomRuntimeFunctionTest.zip", @"CustomRuntimeFunctionTest\bin\Release\net6.0\CustomRuntimeFunctionTest.zip", "CustomRuntimeFunctionTest", TargetFramework.NET6)
+        {
+        }
+    }
+
+    public class CustomRuntimeNET8Tests : CustomRuntimeTests
+    {
+        public CustomRuntimeNET8Tests()
+            : base("CustomRuntimeNET8FunctionTest-" + DateTime.Now.Ticks, "CustomRuntimeFunctionTest.zip", @"CustomRuntimeFunctionTest\bin\Release\net8.0\CustomRuntimeFunctionTest.zip", "CustomRuntimeFunctionTest", TargetFramework.NET8)
+        {
+        }
+    }
+
     public class CustomRuntimeTests : BaseCustomRuntimeTest
     {
-        public CustomRuntimeTests() 
-            : base("CustomRuntimeFunctionTest-" + DateTime.Now.Ticks, "CustomRuntimeFunctionTest.zip", @"CustomRuntimeFunctionTest\bin\Release\net6.0\CustomRuntimeFunctionTest.zip", "CustomRuntimeFunctionTest")
+        public enum TargetFramework { NET6, NET8}
+
+        private TargetFramework _targetFramework;
+
+        public CustomRuntimeTests(string functionName, string deploymentZipKey, string deploymentPackageZipRelativePath, string handler, TargetFramework targetFramework) 
+            : base(functionName, deploymentZipKey, deploymentPackageZipRelativePath, handler)
         {
+            _targetFramework = targetFramework;
         }
 
 #if SKIP_RUNTIME_SUPPORT_INTEG_TESTS
@@ -54,6 +76,13 @@ public async Task TestAllHandlersAsync()
                 {
                     roleAlreadyExisted = await PrepareTestResources(s3Client, lambdaClient, iamClient);
 
+                    // .NET API to address setting memory constraint was added for .NET 8
+                    if(_targetFramework == TargetFramework.NET8)
+                    {
+                        await RunMaxHeapMemoryCheck(lambdaClient, "GetTotalAvailableMemoryBytes");
+                        await RunMaxHeapMemoryCheckWithCustomMemorySettings(lambdaClient, "GetTotalAvailableMemoryBytes");
+                    }
+
                     await RunTestExceptionAsync(lambdaClient, "ExceptionNonAsciiCharacterUnwrappedAsync", "", "Exception", "Unhandled exception with non ASCII character: ♂");
                     await RunTestSuccessAsync(lambdaClient, "UnintendedDisposeTest", "not-used", "UnintendedDisposeTest-SUCCESS");
                     await RunTestSuccessAsync(lambdaClient, "LoggingStressTest", "not-used", "LoggingStressTest-success");
@@ -91,6 +120,40 @@ public async Task TestAllHandlersAsync()
             }
         }
 
+        private async Task RunMaxHeapMemoryCheck(AmazonLambdaClient lambdaClient, string handler)
+        {
+            await UpdateHandlerAsync(lambdaClient, handler);
+            var invokeResponse = await InvokeFunctionAsync(lambdaClient, JsonConvert.SerializeObject(""));
+            using (var responseStream = invokeResponse.Payload)
+            using (var sr = new StreamReader(responseStream))
+            {
+                string payloadStr = (await sr.ReadToEndAsync()).Replace("\"", "");
+                // Function payload response will have format {Handler}-{MemorySize}.
+                // To check memory split on the - and grab the second token representing the memory size.
+                var tokens = payloadStr.Split('-');
+                var memory = long.Parse(tokens[1]);
+                Assert.True(memory <= BaseCustomRuntimeTest.FUNCTION_MEMORY_MB * 1048576);
+            }
+        }
+
+        private async Task RunMaxHeapMemoryCheckWithCustomMemorySettings(AmazonLambdaClient lambdaClient, string handler)
+        {
+            // Set the .NET GC environment variable to say there is 256 MB of memory. The function is deployed with 512 but since the user set
+            // it to 256 Lambda should not make any adjustments.
+            await UpdateHandlerAsync(lambdaClient, handler, new Dictionary<string, string> { { "DOTNET_GCHeapHardLimit", "0x10000000" } });
+            var invokeResponse = await InvokeFunctionAsync(lambdaClient, JsonConvert.SerializeObject(""));
+            using (var responseStream = invokeResponse.Payload)
+            using (var sr = new StreamReader(responseStream))
+            {
+                string payloadStr = (await sr.ReadToEndAsync()).Replace("\"", "");
+                // Function payload response will have format {Handler}-{MemorySize}.
+                // To check memory split on the - and grab the second token representing the memory size.
+                var tokens = payloadStr.Split('-');
+                var memory = long.Parse(tokens[1]);
+                Assert.True(memory <= 256 * 1048576);
+            }
+        }
+
         private async Task RunTestExceptionAsync(AmazonLambdaClient lambdaClient, string handler, string input,
             string expectedErrorType, string expectedErrorMessage)
         {

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/CustomRuntimeFunctionTest/CustomRuntimeFunction.cs

@@ -36,6 +36,9 @@ private static async Task Main(string[] args)
             {
                 switch (handler)
                 {
+                    case nameof(GetTotalAvailableMemoryBytes):
+                        bootstrap = new LambdaBootstrap(GetTotalAvailableMemoryBytes);
+                        break;
                     case nameof(ExceptionNonAsciiCharacterUnwrappedAsync):
                         bootstrap = new LambdaBootstrap(ExceptionNonAsciiCharacterUnwrappedAsync);
                         break;
@@ -426,6 +429,11 @@ private static Task<InvocationResponse> GetTimezoneNameAsync(InvocationRequest i
             return Task.FromResult(GetInvocationResponse(nameof(GetTimezoneNameAsync), TimeZoneInfo.Local.Id));
         }
 
+        private static async Task<InvocationResponse> GetTotalAvailableMemoryBytes(InvocationRequest invocation)
+        {
+            return GetInvocationResponse(nameof(GetTotalAvailableMemoryBytes), GC.GetGCMemoryInfo().TotalAvailableMemoryBytes.ToString());
+        }
+
         #region Helpers
         private static void AssertNotNull(object value, string valueName)
         {

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/CustomRuntimeFunctionTest/CustomRuntimeFunctionTest.csproj

@@ -2,7 +2,7 @@
 
   <PropertyGroup>
     <OutputType>Exe</OutputType>
-    <TargetFramework>net6.0</TargetFramework>
+    <TargetFrameworks>net6.0;net8.0</TargetFrameworks>
   </PropertyGroup>
 
   <ItemGroup>