Reduce allocations in Amazon.Lambda.RuntimeSupport

.autover/changes/a97a0e90-8fc6-4ece-91ec-3038929f09bc.json

@@ -0,0 +1,11 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.RuntimeSupport",
+      "Type": "Patch",
+      "ChangelogMessages": [
+        "Reduce the memory allocated and CPU time used to process each invocation."
+      ]
+    }
+  ]
+}

Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs

@@ -98,6 +98,10 @@ public partial class RuntimeApiSerializationContext : JsonSerializerContext
         private readonly System.Net.Http.HttpClient _httpClient;
         private InternalLogger _logger;
 
+        // The URI most recently used to get the next invocation, which is cached as it is used for every invocation.
+        // Accessed with Volatile.Read/Volatile.Write as it can be read and written by multiple concurrent processing tasks.
+        private NextUri _nextUri;
+
         public InternalRuntimeApiClient(System.Net.Http.HttpClient httpClient)
         {
             _httpClient = httpClient;
@@ -134,10 +138,10 @@ private async System.Threading.Tasks.Task<SwaggerResponse<StatusResponse>> Error
                         request_.Headers.TryAddWithoutValidation("Lambda-Runtime-Function-Error-Type", ConvertToString(lambda_Runtime_Function_Error_Type, System.Globalization.CultureInfo.InvariantCulture));
                     using (var content_ = new System.Net.Http.StringContent(errorJson))
                     {
-                        content_.Headers.ContentType = System.Net.Http.Headers.MediaTypeHeaderValue.Parse(ErrorContentType);
+                        content_.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue(ErrorContentType);
                         request_.Content = content_;
-                        request_.Method = new System.Net.Http.HttpMethod("POST");
-                        request_.Headers.Accept.Add(System.Net.Http.Headers.MediaTypeWithQualityHeaderValue.Parse("application/json"));
+                        request_.Method = System.Net.Http.HttpMethod.Post;
+                        request_.Headers.TryAddWithoutValidation("Accept", "application/json");
 
                         var url_ = urlBuilder_.ToString();
                         request_.RequestUri = new System.Uri(url_, System.UriKind.RelativeOrAbsolute);
@@ -251,17 +255,20 @@ public async Task<SwaggerResponse<StatusResponse>> RestoreErrorAsync(string lamb
             {
                 using (var request_ = new System.Net.Http.HttpRequestMessage())
                 {
-                    request_.Method = new System.Net.Http.HttpMethod("GET");
+                    request_.Method = System.Net.Http.HttpMethod.Get;
 
-                    request_.Headers.Accept.Add(System.Net.Http.Headers.MediaTypeWithQualityHeaderValue.Parse("application/json"));
+                    request_.Headers.TryAddWithoutValidation("Accept", "application/json");
 
-                    var url_ = BaseUrl.TrimEnd('/') + endpointUrl;
-                    request_.RequestUri = new System.Uri(url_, System.UriKind.Absolute);
+                    request_.RequestUri = GetNextUri(endpointUrl);
 
                     var response_ = await client_.SendAsync(request_, System.Net.Http.HttpCompletionOption.ResponseHeadersRead, cancellationToken).ConfigureAwait(false);
                     try
                     {
-                        var headers_ = System.Linq.Enumerable.ToDictionary(response_.Headers, h_ => h_.Key, h_ => h_.Value);
+                        // Size the dictionary for all the headers to avoid resizing it for every invocation
+                        var headers_ = new System.Collections.Generic.Dictionary<string, System.Collections.Generic.IEnumerable<string>>(
+                            response_.Headers.NonValidated.Count + response_.Content.Headers.NonValidated.Count);
+                        foreach (var item_ in response_.Headers)
+                                headers_.Add(item_.Key, item_.Value);
                         foreach (var item_ in response_.Content.Headers)
                                 headers_[item_.Key] = item_.Value;
 
@@ -310,6 +317,47 @@ public async Task<SwaggerResponse<StatusResponse>> RestoreErrorAsync(string lamb
             }
         }
 
+        private System.Uri GetNextUri(string endpointUrl)
+        {
+            var baseUrl = BaseUrl;
+            var nextUri = Volatile.Read(ref _nextUri);
+
+            if (nextUri == null ||
+                !string.Equals(nextUri.BaseUrl, baseUrl, StringComparison.Ordinal) ||
+                !string.Equals(nextUri.EndpointUrl, endpointUrl, StringComparison.Ordinal))
+            {
+                var uri = new System.Uri(baseUrl.TrimEnd('/') + endpointUrl, System.UriKind.Absolute);
+                nextUri = new NextUri(baseUrl, endpointUrl, uri);
+
+                // If multiple processing tasks create the URI concurrently they will all create equivalent values,
+                // so it does not matter which one is published; the important thing is that the instance is
+                // safely published so that no task observes a partially constructed instance.
+                Volatile.Write(ref _nextUri, nextUri);
+            }
+
+            return nextUri.Uri;
+        }
+
+        /// <summary>
+        /// The URI to get the next invocation for a base URL and endpoint URL. This class is immutable so that
+        /// it can be safely read by multiple concurrent processing tasks.
+        /// </summary>
+        private sealed class NextUri
+        {
+            public NextUri(string baseUrl, string endpointUrl, System.Uri uri)
+            {
+                BaseUrl = baseUrl;
+                EndpointUrl = endpointUrl;
+                Uri = uri;
+            }
+
+            public string BaseUrl { get; }
+
+            public string EndpointUrl { get; }
+
+            public System.Uri Uri { get; }
+        }
+
         /// <summary>Runtime makes this request in order to submit a response.</summary>
         /// <returns>Accepted</returns>
         /// <exception cref="RuntimeApiClientException">A server side error occurred.</exception>
@@ -344,10 +392,10 @@ public async System.Threading.Tasks.Task<SwaggerResponse<StatusResponse>> Respon
 
                     try
                     {
-                        content_.Headers.ContentType = System.Net.Http.Headers.MediaTypeHeaderValue.Parse("application/json");
+                        content_.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/json");
                         request_.Content = content_;
-                        request_.Method = new System.Net.Http.HttpMethod("POST");
-                        request_.Headers.Accept.Add(System.Net.Http.Headers.MediaTypeWithQualityHeaderValue.Parse("application/json"));
+                        request_.Method = System.Net.Http.HttpMethod.Post;
+                        request_.Headers.TryAddWithoutValidation("Accept", "application/json");
 
                         // Echo the per-invocation id back for cross-wiring protection. Only sent when the
                         // Runtime API provided it on /next; added to this request message (never the shared
@@ -362,18 +410,19 @@ public async System.Threading.Tasks.Task<SwaggerResponse<StatusResponse>> Respon
                         _logger.LogInformation("Return from SendAsync");
                         try
                         {
+                            if (response_.StatusCode == HttpStatusCode.Accepted)
+                            {
+                                // The headers of a successful response are not used, so do not copy them for every invocation
+                                return new SwaggerResponse<StatusResponse>((int)response_.StatusCode, new System.Collections.Generic.Dictionary<string, System.Collections.Generic.IEnumerable<string>>(), new StatusResponse());
+                            }
+
                             var headers_ = System.Linq.Enumerable.ToDictionary(response_.Headers, h_ => h_.Key, h_ => h_.Value);
                             if (response_.Content != null && response_.Content.Headers != null)
                             {
                                 foreach (var item_ in response_.Content.Headers)
                                     headers_[item_.Key] = item_.Value;
                             }
 
-                            if (response_.StatusCode == HttpStatusCode.Accepted)
-                            {
-                                return new SwaggerResponse<StatusResponse>((int)response_.StatusCode, headers_, new StatusResponse());
-                            }
-                            else
                             if (response_.StatusCode == HttpStatusCode.Gone)
                             {
                                 // 410 Gone means the invocation timed out before this response was submitted.
@@ -512,10 +561,10 @@ public async System.Threading.Tasks.Task<SwaggerResponse<StatusResponse>> ErrorW
 
                     using (var content_ = new System.Net.Http.StringContent(errorJson))
                     {
-                        content_.Headers.ContentType = System.Net.Http.Headers.MediaTypeHeaderValue.Parse(ErrorContentType);
+                        content_.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue(ErrorContentType);
                         request_.Content = content_;
-                        request_.Method = new System.Net.Http.HttpMethod("POST");
-                        request_.Headers.Accept.Add(System.Net.Http.Headers.MediaTypeWithQualityHeaderValue.Parse("application/json"));
+                        request_.Method = System.Net.Http.HttpMethod.Post;
+                        request_.Headers.TryAddWithoutValidation("Accept", "application/json");
 
                         var url_ = urlBuilder_.ToString();
                         request_.RequestUri = new System.Uri(url_, System.UriKind.RelativeOrAbsolute);

Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiHeaders.cs

@@ -63,16 +63,14 @@ internal class RuntimeApiHeaders : IRuntimeApiHeaders
 
         public RuntimeApiHeaders(Dictionary<string, IEnumerable<string>> headers)
         {
-            var caseInsensitiveHeaders = new Dictionary<string, IEnumerable<string>>(headers, StringComparer.OrdinalIgnoreCase);
-
-            DeadlineMs = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderDeadlineMs);
-            AwsRequestId = GetHeaderValueRequired(caseInsensitiveHeaders, HeaderAwsRequestId);
-            ClientContextJson = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderClientContext);
-            CognitoIdentityJson = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderCognitoIdentity);
-            InvokedFunctionArn = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderInvokedFunctionArn);
-            TraceId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderTraceId);
-            TenantId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderAwsTenantId);
-            InvocationId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderInvocationId);
+            DeadlineMs = GetHeaderValueOrNull(headers, HeaderDeadlineMs);
+            AwsRequestId = GetHeaderValueRequired(headers, HeaderAwsRequestId);
+            ClientContextJson = GetHeaderValueOrNull(headers, HeaderClientContext);
+            CognitoIdentityJson = GetHeaderValueOrNull(headers, HeaderCognitoIdentity);
+            InvokedFunctionArn = GetHeaderValueOrNull(headers, HeaderInvokedFunctionArn);
+            TraceId = GetHeaderValueOrNull(headers, HeaderTraceId);
+            TenantId = GetHeaderValueOrNull(headers, HeaderAwsTenantId);
+            InvocationId = GetHeaderValueOrNull(headers, HeaderInvocationId);
         }
 
         public string AwsRequestId { get; private set; }
@@ -86,17 +84,45 @@ public RuntimeApiHeaders(Dictionary<string, IEnumerable<string>> headers)
 
         private string GetHeaderValueRequired(Dictionary<string, IEnumerable<string>> headers, string header)
         {
-            return headers[header].FirstOrDefault();
+            if (!TryGetHeaderValues(headers, header, out var values))
+            {
+                throw new KeyNotFoundException($"The given key '{header}' was not present in the dictionary.");
+            }
+
+            return values.FirstOrDefault();
         }
 
         private string GetHeaderValueOrNull(Dictionary<string, IEnumerable<string>> headers, string header)
         {
-            if (headers.TryGetValue(header, out var values))
+            if (TryGetHeaderValues(headers, header, out var values))
             {
                 return values.FirstOrDefault();
             }
 
             return null;
         }
+
+        private static bool TryGetHeaderValues(Dictionary<string, IEnumerable<string>> headers, string header, out IEnumerable<string> values)
+        {
+            // Header names are case-insensitive. The Runtime API sends the headers with the casing of the
+            // constants, so look the header up directly first and only fall back to searching the headers
+            // ignoring case, rather than copying them into a case-insensitive dictionary for every invocation.
+            if (headers.TryGetValue(header, out values))
+            {
+                return true;
+            }
+
+            foreach (var pair in headers)
+            {
+                if (string.Equals(pair.Key, header, StringComparison.OrdinalIgnoreCase))
+                {
+                    values = pair.Value;
+                    return true;
+                }
+            }
+
+            values = null;
+            return false;
+        }
     }
 }

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

@@ -16,6 +16,7 @@
 using Amazon.Lambda.Core;
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using Amazon.Lambda.RuntimeSupport.Helpers;
 
 namespace Amazon.Lambda.RuntimeSupport
@@ -24,18 +25,27 @@ internal class LambdaContext : ILambdaContext
     {
         internal static readonly DateTime UnixEpoch = new DateTime(1970, 1, 1, 0, 0, 0, DateTimeKind.Utc);
 
+        private static readonly IDateTimeHelper DefaultDateTimeHelper = new DateTimeHelper();
+
+        // Cached delegates and a placeholder for a null value, used to create the values that are parsed
+        // from the headers only when they are first used without allocating for every invocation.
+        private static readonly Func<string, CognitoIdentity> CreateCognitoIdentity = CognitoIdentity.FromJson;
+        private static readonly Func<string, CognitoClientContext> CreateCognitoClientContext = CognitoClientContext.FromJson;
+        private static readonly Func<string, IReadOnlyDictionary<string, string>> CreateW3C = W3CTraceContext.FromClientContextJson;
+        private static readonly object NullValue = new object();
+
         private readonly LambdaEnvironment _lambdaEnvironment;
         private readonly RuntimeApiHeaders _runtimeApiHeaders;
         private readonly IDateTimeHelper _dateTimeHelper;
         private readonly long _deadlineMs;
         private readonly int _memoryLimitInMB;
-        private readonly Lazy<CognitoIdentity> _cognitoIdentityLazy;
-        private readonly Lazy<CognitoClientContext> _cognitoClientContextLazy;
-        private readonly Lazy<IReadOnlyDictionary<string, string>> _w3cLazy;
         private readonly IConsoleLoggerWriter _consoleLogger;
+        private object _cognitoIdentity;
+        private object _cognitoClientContext;
+        private object _w3c;
 
         public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lambdaEnvironment, IConsoleLoggerWriter consoleLogger)
-            : this(runtimeApiHeaders, lambdaEnvironment, new DateTimeHelper(), consoleLogger)
+            : this(runtimeApiHeaders, lambdaEnvironment, DefaultDateTimeHelper, consoleLogger)
         {
         }
 
@@ -49,22 +59,19 @@ public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lamb
 
             int.TryParse(_lambdaEnvironment.FunctionMemorySize, out _memoryLimitInMB);
             long.TryParse(_runtimeApiHeaders.DeadlineMs, out _deadlineMs);
-            _cognitoIdentityLazy = new Lazy<CognitoIdentity>(() => CognitoIdentity.FromJson(runtimeApiHeaders.CognitoIdentityJson));
-            _cognitoClientContextLazy = new Lazy<CognitoClientContext>(() => CognitoClientContext.FromJson(runtimeApiHeaders.ClientContextJson));
-            _w3cLazy = new Lazy<IReadOnlyDictionary<string, string>>(() => W3CTraceContext.FromClientContextJson(runtimeApiHeaders.ClientContextJson));
         }
 
         public string TraceId => _runtimeApiHeaders.TraceId;
 
         public string AwsRequestId => _runtimeApiHeaders.AwsRequestId;
 
-        public IClientContext ClientContext => _cognitoClientContextLazy.Value;
+        public IClientContext ClientContext => GetOrCreate(ref _cognitoClientContext, _runtimeApiHeaders.ClientContextJson, CreateCognitoClientContext);
 
         public string FunctionName => _lambdaEnvironment.FunctionName;
 
         public string FunctionVersion => _lambdaEnvironment.FunctionVersion;
 
-        public ICognitoIdentity Identity => _cognitoIdentityLazy.Value;
+        public ICognitoIdentity Identity => GetOrCreate(ref _cognitoIdentity, _runtimeApiHeaders.CognitoIdentityJson, CreateCognitoIdentity);
 
         public string InvokedFunctionArn => _runtimeApiHeaders.InvokedFunctionArn;
 
@@ -80,7 +87,7 @@ public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lamb
 
         public string TenantId => _runtimeApiHeaders.TenantId;
 
-        public IReadOnlyDictionary<string, string> W3C => _w3cLazy.Value;
+        public IReadOnlyDictionary<string, string> W3C => GetOrCreate(ref _w3c, _runtimeApiHeaders.ClientContextJson, CreateW3C);
 
         /// <summary>
         /// The serializer the Lambda function registered with the runtime, surfaced via
@@ -97,5 +104,24 @@ public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lamb
         /// the <c>Lambda-Runtime-Invocation-Id</c> header, in which case nothing is echoed.
         /// </summary>
         internal string InvocationId => _runtimeApiHeaders.InvocationId;
+
+        /// <summary>
+        /// Gets the value stored in the specified location, creating it from the specified JSON if it has not been
+        /// created yet. The value is created from immutable header values, so if it is created concurrently the
+        /// first value to be stored is the one that is returned to all callers.
+        /// </summary>
+        private static T GetOrCreate<T>(ref object location, string json, Func<string, T> factory)
+            where T : class
+        {
+            var value = Volatile.Read(ref location);
+
+            if (value == null)
+            {
+                value = (object)factory(json) ?? NullValue;
+                value = Interlocked.CompareExchange(ref location, value, null) ?? value;
+            }
+
+            return ReferenceEquals(value, NullValue) ? null : (T)value;
+        }
     }
 }