diff --git a/.autover/changes/a97a0e90-8fc6-4ece-91ec-3038929f09bc.json b/.autover/changes/a97a0e90-8fc6-4ece-91ec-3038929f09bc.json new file mode 100644 index 000000000..0eb37b8d1 --- /dev/null +++ b/.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." + ] + } + ] +} diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs index 1036d6642..82cfd85e1 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/InternalClientAdapted.cs +++ b/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> 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> 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>( + 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> 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; + } + + /// + /// 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. + /// + 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; } + } + /// Runtime makes this request in order to submit a response. /// Accepted /// A server side error occurred. @@ -344,10 +392,10 @@ public async System.Threading.Tasks.Task> 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,6 +410,12 @@ public async System.Threading.Tasks.Task> 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((int)response_.StatusCode, new System.Collections.Generic.Dictionary>(), new StatusResponse()); + } + var headers_ = System.Linq.Enumerable.ToDictionary(response_.Headers, h_ => h_.Key, h_ => h_.Value); if (response_.Content != null && response_.Content.Headers != null) { @@ -369,11 +423,6 @@ public async System.Threading.Tasks.Task> Respon headers_[item_.Key] = item_.Value; } - if (response_.StatusCode == HttpStatusCode.Accepted) - { - return new SwaggerResponse((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> 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); diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiHeaders.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiHeaders.cs index 6dd922914..889db748a 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiHeaders.cs +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiHeaders.cs @@ -63,16 +63,14 @@ internal class RuntimeApiHeaders : IRuntimeApiHeaders public RuntimeApiHeaders(Dictionary> headers) { - var caseInsensitiveHeaders = new Dictionary>(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> headers) private string GetHeaderValueRequired(Dictionary> 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> 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> headers, string header, out IEnumerable 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; + } } } diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaContext.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaContext.cs index f91c80a81..4c2474a2b 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Context/LambdaContext.cs +++ b/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 CreateCognitoIdentity = CognitoIdentity.FromJson; + private static readonly Func CreateCognitoClientContext = CognitoClientContext.FromJson; + private static readonly Func> 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 _cognitoIdentityLazy; - private readonly Lazy _cognitoClientContextLazy; - private readonly Lazy> _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.FromJson(runtimeApiHeaders.CognitoIdentityJson)); - _cognitoClientContextLazy = new Lazy(() => CognitoClientContext.FromJson(runtimeApiHeaders.ClientContextJson)); - _w3cLazy = new Lazy>(() => 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 W3C => _w3cLazy.Value; + public IReadOnlyDictionary W3C => GetOrCreate(ref _w3c, _runtimeApiHeaders.ClientContextJson, CreateW3C); /// /// The serializer the Lambda function registered with the runtime, surfaced via @@ -97,5 +104,24 @@ public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lamb /// the Lambda-Runtime-Invocation-Id header, in which case nothing is echoed. /// internal string InvocationId => _runtimeApiHeaders.InvocationId; + + /// + /// 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. + /// + private static T GetOrCreate(ref object location, string json, Func 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; + } } } diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/InternalLogger.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/InternalLogger.cs index bd301fac7..d018237e5 100644 --- a/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/InternalLogger.cs +++ b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/InternalLogger.cs @@ -22,36 +22,52 @@ internal class InternalLogger private const string DebuggingEnvironmentVariable = "LAMBDA_RUNTIMESUPPORT_DEBUG"; public static readonly InternalLogger ConsoleLogger = new InternalLogger(Console.WriteLine); - public static readonly InternalLogger NoOpLogger = new InternalLogger(message => { }); + public static readonly InternalLogger NoOpLogger = new InternalLogger(message => { }, isEnabled: false); private readonly Action _internalLoggingAction; + // Whether messages are logged. When false, messages are not formatted to avoid allocating + // strings that would be discarded, as the logger is called multiple times for every invocation. + private readonly bool _isEnabled; + /// /// Constructs InternalLogger which logs to the internalLoggingAction. /// /// - private InternalLogger(Action internalLoggingAction) + /// Whether messages are logged. + private InternalLogger(Action internalLoggingAction, bool isEnabled = true) { _internalLoggingAction = internalLoggingAction; if (_internalLoggingAction == null) { throw new ArgumentNullException(nameof(internalLoggingAction)); } + + _isEnabled = isEnabled; } public void LogDebug(string message) { - _internalLoggingAction($"[Debug] {message}"); + if (_isEnabled) + { + _internalLoggingAction($"[Debug] {message}"); + } } public void LogError(Exception exception, string message) { - _internalLoggingAction($"[Error] {message} - {exception.ToString()}"); + if (_isEnabled) + { + _internalLoggingAction($"[Error] {message} - {exception.ToString()}"); + } } public void LogInformation(string message) { - _internalLoggingAction($"[Info] {message}"); + if (_isEnabled) + { + _internalLoggingAction($"[Info] {message}"); + } } /// diff --git a/Libraries/test/Amazon.Lambda.Annotations.SourceGenerators.Tests/SourceGeneratorTests.cs b/Libraries/test/Amazon.Lambda.Annotations.SourceGenerators.Tests/SourceGeneratorTests.cs index 7e757f152..009c6643c 100644 --- a/Libraries/test/Amazon.Lambda.Annotations.SourceGenerators.Tests/SourceGeneratorTests.cs +++ b/Libraries/test/Amazon.Lambda.Annotations.SourceGenerators.Tests/SourceGeneratorTests.cs @@ -2392,15 +2392,15 @@ private static DiagnosticResult[] GetExpectedRuntimeSupportDiagnostics() DiagnosticResult.CompilerError("CS0534").WithSpan(clientFile, 89, 30, 89, 60).WithArguments(runtimeApiContext, "System.Text.Json.Serialization.JsonSerializerContext.GeneratedSerializerOptions.get"), DiagnosticResult.CompilerError("CS0534").WithSpan(clientFile, 89, 30, 89, 60).WithArguments(runtimeApiContext, "System.Text.Json.Serialization.JsonSerializerContext.GetTypeInfo(System.Type)"), DiagnosticResult.CompilerError("CS7036").WithSpan(clientFile, 89, 30, 89, 60).WithArguments("options", "System.Text.Json.Serialization.JsonSerializerContext.JsonSerializerContext(System.Text.Json.JsonSerializerOptions?)"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 161, 136, 161, 143).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 176, 135, 176, 142).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 279, 131, 279, 138).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 393, 135, 393, 142).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 408, 135, 408, 142).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 423, 135, 423, 142).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 539, 136, 539, 143).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 563, 135, 563, 142).WithArguments(runtimeApiContext, "Default"), - DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 578, 135, 578, 142).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 165, 136, 165, 143).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 180, 135, 180, 142).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 286, 131, 286, 138).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 442, 135, 442, 142).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 457, 135, 457, 142).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 472, 135, 472, 142).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 588, 136, 588, 143).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 612, 135, 612, 142).WithArguments(runtimeApiContext, "Default"), + DiagnosticResult.CompilerError("CS0117").WithSpan(clientFile, 627, 135, 627, 142).WithArguments(runtimeApiContext, "Default"), // These are here because the internalvisibleto attribute isn't included in test compilations, so these types are inaccessible. diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaContextTests.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaContextTests.cs index cb6c64e13..3a8671664 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaContextTests.cs +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaContextTests.cs @@ -90,5 +90,69 @@ public void InvocationIdIsNullWhenHeaderAbsent() Assert.Null(runtimeApiHeaders.InvocationId); Assert.Null(context.InvocationId); } + + [Fact] + public void RuntimeApiHeadersAreReadIgnoringCase() + { + var headers = new Dictionary> + { + ["lambda-runtime-aws-request-id"] = new[] { "request-generated-id" }, + ["LAMBDA-RUNTIME-INVOKED-FUNCTION-ARN"] = new[] { "my-function-arn" }, + ["lambda-runtime-trace-id"] = new[] { "trace-id" }, + ["Lambda-Runtime-Client-Context"] = new[] { "{}" }, + ["lambda-runtime-cognito-identity"] = new[] { "{}" }, + ["lambda-runtime-deadline-ms"] = new[] { "1542409706888" }, + ["lambda-runtime-aws-tenant-id"] = new[] { "tenant-generated-id" }, + ["lambda-runtime-invocation-id"] = new[] { "invocation-generated-id" } + }; + + var runtimeApiHeaders = new RuntimeApiHeaders(headers); + + Assert.Equal("request-generated-id", runtimeApiHeaders.AwsRequestId); + Assert.Equal("my-function-arn", runtimeApiHeaders.InvokedFunctionArn); + Assert.Equal("trace-id", runtimeApiHeaders.TraceId); + Assert.Equal("{}", runtimeApiHeaders.ClientContextJson); + Assert.Equal("{}", runtimeApiHeaders.CognitoIdentityJson); + Assert.Equal("1542409706888", runtimeApiHeaders.DeadlineMs); + Assert.Equal("tenant-generated-id", runtimeApiHeaders.TenantId); + Assert.Equal("invocation-generated-id", runtimeApiHeaders.InvocationId); + } + + [Fact] + public void ContextValuesAreParsedFromHeadersOnce() + { + var headers = new Dictionary> + { + ["Lambda-Runtime-Aws-Request-Id"] = new[] { "request-generated-id" }, + ["Lambda-Runtime-Client-Context"] = new[] { "{\"custom\":{\"value\":\"test\"}}" }, + ["Lambda-Runtime-Cognito-Identity"] = new[] { "{\"cognitoIdentityId\":\"identity-id\",\"cognitoIdentityPoolId\":\"identity-pool-id\"}" } + }; + + var runtimeApiHeaders = new RuntimeApiHeaders(headers); + var lambdaEnvironment = new LambdaEnvironment(_environmentVariables); + + var context = new LambdaContext(runtimeApiHeaders, lambdaEnvironment, new Helpers.LogLevelLoggerWriter(new SystemEnvironmentVariables())); + + Assert.Equal("test", context.ClientContext.Custom["value"]); + Assert.Same(context.ClientContext, context.ClientContext); + + Assert.Equal("identity-id", context.Identity.IdentityId); + Assert.Equal("identity-pool-id", context.Identity.IdentityPoolId); + Assert.Same(context.Identity, context.Identity); + + Assert.Empty(context.W3C); + Assert.Same(context.W3C, context.W3C); + } + + [Fact] + public void RuntimeApiHeadersThrowsWhenRequestIdIsAbsent() + { + var headers = new Dictionary> + { + ["Lambda-Runtime-Invoked-Function-Arn"] = new[] { "my-function-arn" } + }; + + Assert.Throws(() => new RuntimeApiHeaders(headers)); + } } } diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/RuntimeApiClientTests.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/RuntimeApiClientTests.cs index f78e8f700..14a93e89f 100644 --- a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/RuntimeApiClientTests.cs +++ b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/RuntimeApiClientTests.cs @@ -294,6 +294,155 @@ public async Task ReportInvocationErrorAsync_410Gone_ThrowsRuntimeApiInvokeTimeo Assert.Equal("req-timeout", ex.AwsRequestId); } + // --- Request messages sent to the Runtime API --- + + /// + /// Mock handler that captures the request and returns a next invocation with the specified headers. + /// + private class NextInvocationMockHttpMessageHandler : HttpMessageHandler + { + public HttpRequestMessage CapturedRequest { get; private set; } + private readonly (string Name, string Value)[] _headers; + + public NextInvocationMockHttpMessageHandler(params (string Name, string Value)[] headers) + { + _headers = headers; + } + + protected override Task SendAsync( + HttpRequestMessage request, CancellationToken cancellationToken) + { + CapturedRequest = request; + var response = new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent("{\"key\":\"value\"}") + }; + + foreach (var (name, value) in _headers) + { + response.Headers.TryAddWithoutValidation(name, value); + } + + return Task.FromResult(response); + } + } + + [Fact] + public async Task GetNextInvocationAsync_SendsExpectedRequest() + { + var handler = new NextInvocationMockHttpMessageHandler(("Lambda-Runtime-Aws-Request-Id", "req-1")); + var httpClient = new HttpClient(handler); + var envVars = new TestEnvironmentVariables(); + envVars.SetEnvironmentVariable("AWS_LAMBDA_RUNTIME_API", "localhost:9001"); + var client = new RuntimeApiClient(envVars, httpClient); + + using var invocation = await client.GetNextInvocationAsync(CancellationToken.None); + + Assert.NotNull(handler.CapturedRequest); + Assert.Equal(HttpMethod.Get, handler.CapturedRequest.Method); + Assert.Equal("http://localhost:9001/2018-06-01/runtime/invocation/next", handler.CapturedRequest.RequestUri.ToString()); + Assert.Equal("application/json", handler.CapturedRequest.Headers.GetValues("Accept").Single()); + } + + [Fact] + public async Task GetNextInvocationAsync_UsesExpectedUriForEachEndpoint() + { + var handler = new NextInvocationMockHttpMessageHandler(("Lambda-Runtime-Aws-Request-Id", "req-1")); + var httpClient = new HttpClient(handler); + var envVars = new TestEnvironmentVariables(); + envVars.SetEnvironmentVariable("AWS_LAMBDA_RUNTIME_API", "localhost:9001"); + var client = new RuntimeApiClient(envVars, httpClient); + + await client.RestoreNextInvocationAsync(CancellationToken.None); + Assert.Equal("http://localhost:9001/2018-06-01/runtime/restore/next", handler.CapturedRequest.RequestUri.ToString()); + + using (await client.GetNextInvocationAsync(CancellationToken.None)) + { + Assert.Equal("http://localhost:9001/2018-06-01/runtime/invocation/next", handler.CapturedRequest.RequestUri.ToString()); + } + + using (await client.GetNextInvocationAsync(CancellationToken.None)) + { + Assert.Equal("http://localhost:9001/2018-06-01/runtime/invocation/next", handler.CapturedRequest.RequestUri.ToString()); + } + + await client.RestoreNextInvocationAsync(CancellationToken.None); + Assert.Equal("http://localhost:9001/2018-06-01/runtime/restore/next", handler.CapturedRequest.RequestUri.ToString()); + } + + [Fact] + public async Task GetNextInvocationAsync_ReadsRuntimeApiHeadersIgnoringCase() + { + var handler = new NextInvocationMockHttpMessageHandler( + ("lambda-runtime-aws-request-id", "req-1"), + ("LAMBDA-RUNTIME-INVOKED-FUNCTION-ARN", "arn:aws:lambda:us-east-2:123456789012:function:custom-runtime"), + ("Lambda-Runtime-Trace-Id", "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"), + ("lambda-runtime-deadline-ms", "1542409706888"), + ("Lambda-Runtime-Aws-Tenant-Id", "tenant-1"), + ("lambda-runtime-invocation-id", "inv-uuid-abc")); + var httpClient = new HttpClient(handler); + var envVars = new TestEnvironmentVariables(); + envVars.SetEnvironmentVariable("AWS_LAMBDA_RUNTIME_API", "localhost:9001"); + var client = new RuntimeApiClient(envVars, httpClient); + + using var invocation = await client.GetNextInvocationAsync(CancellationToken.None); + + var context = Assert.IsType(invocation.LambdaContext); + Assert.Equal("req-1", context.AwsRequestId); + Assert.Equal("arn:aws:lambda:us-east-2:123456789012:function:custom-runtime", context.InvokedFunctionArn); + Assert.Equal("Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1", context.TraceId); + Assert.Equal("tenant-1", context.TenantId); + Assert.Equal("inv-uuid-abc", context.InvocationId); + + using var reader = new StreamReader(invocation.InputStream); + Assert.Equal("{\"key\":\"value\"}", await reader.ReadToEndAsync()); + } + + [Fact] + public async Task SendResponseAsync_SendsExpectedRequest() + { + var handler = new ConfigurableMockHttpMessageHandler(HttpStatusCode.Accepted); + var client = CreateClientWith(handler); + + var outputStream = new MemoryStream(new byte[] { 1, 2, 3 }); + await client.SendResponseAsync("req-1", outputStream, CancellationToken.None); + + Assert.NotNull(handler.CapturedRequest); + Assert.Equal(HttpMethod.Post, handler.CapturedRequest.Method); + Assert.Equal("http://localhost:9001/2018-06-01/runtime/invocation/req-1/response", handler.CapturedRequest.RequestUri.ToString()); + Assert.Equal("application/json", handler.CapturedRequest.Headers.GetValues("Accept").Single()); + Assert.Equal("application/json", handler.CapturedRequest.Content.Headers.GetValues("Content-Type").Single()); + } + + [Fact] + public async Task SendResponseAsync_WithoutOutputStream_SendsExpectedRequest() + { + var handler = new ConfigurableMockHttpMessageHandler(HttpStatusCode.Accepted); + var client = CreateClientWith(handler); + + await client.SendResponseAsync("req-1", null, CancellationToken.None); + + Assert.NotNull(handler.CapturedRequest); + Assert.Equal(HttpMethod.Post, handler.CapturedRequest.Method); + Assert.Equal("application/json", handler.CapturedRequest.Headers.GetValues("Accept").Single()); + Assert.Equal("application/json", handler.CapturedRequest.Content.Headers.GetValues("Content-Type").Single()); + } + + [Fact] + public async Task ReportInvocationErrorAsync_SendsExpectedRequest() + { + var handler = new ConfigurableMockHttpMessageHandler(HttpStatusCode.Accepted, "{}"); + var client = CreateClientWith(handler); + + await client.ReportInvocationErrorAsync("req-1", new Exception("boom"), CancellationToken.None); + + Assert.NotNull(handler.CapturedRequest); + Assert.Equal(HttpMethod.Post, handler.CapturedRequest.Method); + Assert.Equal("http://localhost:9001/2018-06-01/runtime/invocation/req-1/error", handler.CapturedRequest.RequestUri.ToString()); + Assert.Equal("application/json", handler.CapturedRequest.Headers.GetValues("Accept").Single()); + Assert.Equal("application/vnd.aws.lambda.error+json", handler.CapturedRequest.Content.Headers.GetValues("Content-Type").Single()); + } + [Fact] public async Task StartStreamingResponseAsync_NullRequestId_ThrowsArgumentNullException() {