Skip to content

Commit 7c1c6ff

Browse files
normjnormj
authored
Add missing TraceId property on the ILambdaContext
1 parent af70b07 commit 7c1c6ff

File tree

13 files changed

+187
-26
lines changed

13 files changed

+187
-26
lines changed

‎.autover/changes/6298bf78-28f9-45b6-a978-12313b7ea849.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
{
2+
"Projects": [
3+
{
4+
"Name": "Amazon.Lambda.RuntimeSupport",
5+
"Type": "Patch",
6+
"ChangelogMessages": [
7+
"Add missing TraceId property on the ILambdaContext"
8+
]
9+
},
10+
{
11+
"Name": "Amazon.Lambda.Core",
12+
"Type": "Patch",
13+
"ChangelogMessages": [
14+
"Add missing TraceId property on the ILambdaContext"
15+
]
16+
}
17+
]
18+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,4 +63,5 @@ public SnapStartEmptyLambdaContext()
6363
public IDictionary<string, string> Environment { get; } = new Dictionary<string, string>();
6464
public IClientApplication Client { get; }
6565
public IDictionary<string, string> Custom { get; } = new Dictionary<string, string>();
66+
public string TenantId { get; }
6667
}

‎Libraries/src/Amazon.Lambda.Core/ILambdaContext.cs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,5 +84,15 @@ public interface ILambdaContext
8484
/// used to specify function behavior when nearing the timeout.
8585
/// </summary>
8686
TimeSpan RemainingTime { get; }
87+
88+
/// <summary>
89+
/// Gets the tenant id for the Lambda function.
90+
/// </summary>
91+
string TenantId { get; }
92+
93+
/// <summary>
94+
/// The trace id generated by Lambda for distributed tracing across AWS services.
95+
/// </summary>
96+
string TraceId { get; }
8797
}
88-
}
98+
}

‎Libraries/src/Amazon.Lambda.RuntimeSupport/Client/RuntimeApiClient.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
/*
1+
/*
22
* Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License").
@@ -110,7 +110,7 @@ public async Task<InvocationRequest> GetNextInvocationAsync(CancellationToken ca
110110
SwaggerResponse<Stream> response = await _internalClient.NextAsync(cancellationToken);
111111

112112
var headers = new RuntimeApiHeaders(response.Headers);
113-
_consoleLoggerRedirector.SetCurrentAwsRequestId(headers.AwsRequestId);
113+
_consoleLoggerRedirector.SetRuntimeHeaders(headers);
114114

115115
var lambdaContext = new LambdaContext(headers, LambdaEnvironment, _consoleLoggerRedirector);
116116
return new InvocationRequest
@@ -185,4 +185,4 @@ public async Task SendResponseAsync(string awsRequestId, Stream outputStream, Ca
185185
await _internalClient.ResponseAsync(awsRequestId, outputStream, cancellationToken);
186186
}
187187
}
188-
}
188+
}

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

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,38 @@
1818

1919
namespace Amazon.Lambda.RuntimeSupport
2020
{
21-
internal class RuntimeApiHeaders
21+
/// <summary>
22+
/// Headers that are specific to each invocation that are used in varies parts of RuntimeSupport like the logging system.
23+
/// </summary>
24+
public interface IRuntimeApiHeaders
25+
{
26+
/// <summary>
27+
/// The AWS request ID associated with the request.
28+
/// This is the same ID returned to the client that called invoke().
29+
/// This ID is reused for retries on the same request.
30+
/// </summary>
31+
string AwsRequestId { get; }
32+
33+
/// <summary>
34+
/// The trace id generated by Lambda for distributed tracing across AWS services.
35+
/// </summary>
36+
string TraceId { get; }
37+
38+
/// <summary>
39+
/// Gets the tenant id for the Lambda function.
40+
/// </summary>
41+
string TenantId { get; }
42+
}
43+
44+
internal class RuntimeApiHeaders : IRuntimeApiHeaders
2245
{
2346
internal const string HeaderAwsRequestId = "Lambda-Runtime-Aws-Request-Id";
2447
internal const string HeaderTraceId = "Lambda-Runtime-Trace-Id";
2548
internal const string HeaderClientContext = "Lambda-Runtime-Client-Context";
2649
internal const string HeaderCognitoIdentity = "Lambda-Runtime-Cognito-Identity";
2750
internal const string HeaderDeadlineMs = "Lambda-Runtime-Deadline-Ms";
2851
internal const string HeaderInvokedFunctionArn = "Lambda-Runtime-Invoked-Function-Arn";
52+
internal const string HeaderAwsTenantId = "Lambda-Runtime-Aws-Tenant-Id";
2953

3054
public RuntimeApiHeaders(Dictionary<string, IEnumerable<string>> headers)
3155
{
@@ -37,6 +61,7 @@ public RuntimeApiHeaders(Dictionary<string, IEnumerable<string>> headers)
3761
CognitoIdentityJson = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderCognitoIdentity);
3862
InvokedFunctionArn = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderInvokedFunctionArn);
3963
TraceId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderTraceId);
64+
TenantId = GetHeaderValueOrNull(caseInsensitiveHeaders, HeaderAwsTenantId);
4065
}
4166

4267
public string AwsRequestId { get; private set; }
@@ -45,6 +70,7 @@ public RuntimeApiHeaders(Dictionary<string, IEnumerable<string>> headers)
4570
public string ClientContextJson { get; private set; }
4671
public string CognitoIdentityJson { get; private set; }
4772
public string DeadlineMs { get; private set; }
73+
public string TenantId { get; private set; }
4874

4975
private string GetHeaderValueRequired(Dictionary<string, IEnumerable<string>> headers, string header)
5076
{

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,6 @@ public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lamb
5454
_lambdaEnvironment.SetXAmznTraceId(_runtimeApiHeaders.TraceId);
5555
}
5656

57-
// TODO If/When Amazon.Lambda.Core is major versioned, add this to ILambdaContext.
58-
// Until then function code can access it via the _X_AMZN_TRACE_ID environment variable set by LambdaBootstrap.
5957
public string TraceId => _runtimeApiHeaders.TraceId;
6058

6159
public string AwsRequestId => _runtimeApiHeaders.AwsRequestId;
@@ -79,5 +77,7 @@ public LambdaContext(RuntimeApiHeaders runtimeApiHeaders, LambdaEnvironment lamb
7977
public int MemoryLimitInMB => _memoryLimitInMB;
8078

8179
public TimeSpan RemainingTime => TimeSpan.FromMilliseconds(_deadlineMs - (_dateTimeHelper.UtcNow - UnixEpoch).TotalMilliseconds);
80+
81+
public string TenantId => _runtimeApiHeaders.TenantId;
8282
}
8383
}

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

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,10 @@ namespace Amazon.Lambda.RuntimeSupport.Helpers
3434
public interface IConsoleLoggerWriter
3535
{
3636
/// <summary>
37-
/// The current aws request id
37+
/// The runtime headers for the current event to be added to log messages.
3838
/// </summary>
39-
/// <param name="awsRequestId">The AWS request id for the function invocation added to each log message.</param>
40-
void SetCurrentAwsRequestId(string awsRequestId);
39+
/// <param name="runtimeApiHeaders">The headers containing invocation specific information that can be during log formatting.</param>
40+
void SetRuntimeHeaders(IRuntimeApiHeaders runtimeApiHeaders);
4141

4242
/// <summary>
4343
/// Format message with default log level
@@ -99,7 +99,7 @@ public SimpleLoggerWriter()
9999
}
100100

101101
/// <inheritdoc/>
102-
public void SetCurrentAwsRequestId(string awsRequestId)
102+
public void SetRuntimeHeaders(IRuntimeApiHeaders runtimeApiHeaders)
103103
{
104104
}
105105

@@ -268,10 +268,10 @@ private void ConfigureLoggingActionField()
268268
}
269269

270270
/// <inheritdoc/>
271-
public void SetCurrentAwsRequestId(string awsRequestId)
271+
public void SetRuntimeHeaders(IRuntimeApiHeaders runtimeApiHeaders)
272272
{
273-
_wrappedStdOutWriter.CurrentAwsRequestId = awsRequestId;
274-
_wrappedStdErrorWriter.CurrentAwsRequestId = awsRequestId;
273+
_wrappedStdOutWriter.CurrentRuntimeApiHeaders = runtimeApiHeaders;
274+
_wrappedStdErrorWriter.CurrentRuntimeApiHeaders = runtimeApiHeaders;
275275
}
276276

277277
/// <inheritdoc/>
@@ -311,7 +311,7 @@ enum LogFormatType { Default, Unformatted, Json }
311311

312312
private readonly ILogMessageFormatter _logMessageFormatter;
313313

314-
public string CurrentAwsRequestId { get; set; } = string.Empty;
314+
public IRuntimeApiHeaders CurrentRuntimeApiHeaders { get; set; }
315315

316316
/// <summary>
317317
/// This is typically set to either Console.Out or Console.Error to make sure we acquiring a lock
@@ -388,7 +388,7 @@ internal void FormattedWriteLine(string message)
388388
FormattedWriteLine(_defaultLogLevel, (Exception)null, message);
389389
}
390390

391-
internal void FormattedWriteLine(string level, Exception exeception, string messageTemplate, params object[] args)
391+
internal void FormattedWriteLine(string level, Exception exception, string messageTemplate, params object[] args)
392392
{
393393
lock(LockObject)
394394
{
@@ -404,10 +404,15 @@ internal void FormattedWriteLine(string level, Exception exeception, string mess
404404
messageState.MessageTemplate = messageTemplate ?? string.Empty;
405405
messageState.MessageArguments = args;
406406
messageState.TimeStamp = DateTime.UtcNow;
407-
messageState.AwsRequestId = CurrentAwsRequestId;
408-
messageState.TraceId = Environment.GetEnvironmentVariable(LambdaEnvironment.EnvVarTraceId);
409407
messageState.Level = levelEnum;
410-
messageState.Exception = exeception;
408+
messageState.Exception = exception;
409+
410+
if (CurrentRuntimeApiHeaders != null)
411+
{
412+
messageState.AwsRequestId = CurrentRuntimeApiHeaders.AwsRequestId;
413+
messageState.TenantId = CurrentRuntimeApiHeaders.TenantId;
414+
messageState.TraceId = CurrentRuntimeApiHeaders.TraceId;
415+
}
411416

412417
var message = _logMessageFormatter.FormatMessage(messageState);
413418
_innerWriter.WriteLine(message);

‎Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/Logging/JsonLogMessageFormatter.cs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ public class JsonLogMessageFormatter : AbstractLogMessageFormatter
1919
private static readonly UTF8Encoding UTF8NoBomNoThrow = new UTF8Encoding(encoderShouldEmitUTF8Identifier: false, throwOnInvalidBytes: false);
2020

2121
// Options used when serializing any message property values as a JSON to be added to the structured log message.
22-
private JsonSerializerOptions _jsonSerializationOptions;
22+
private readonly JsonSerializerOptions _jsonSerializationOptions;
2323

2424
/// <summary>
2525
/// Constructs an instance of JsonLogMessageFormatter.
@@ -78,7 +78,12 @@ public override string FormatMessage(MessageState state)
7878
{
7979
writer.WriteString("requestId", state.AwsRequestId);
8080
}
81-
81+
82+
if (!string.IsNullOrEmpty(state.TenantId))
83+
{
84+
writer.WriteString("tenantId", state.TenantId);
85+
}
86+
8287
if (!string.IsNullOrEmpty(state.TraceId))
8388
{
8489
writer.WriteString("traceId", state.TraceId);

‎Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/Logging/MessageState.cs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#if NET6_0_OR_GREATER
1+
#if NET6_0_OR_GREATER
22
using System;
33

44
namespace Amazon.Lambda.RuntimeSupport.Helpers.Logging
@@ -19,6 +19,11 @@ public class MessageState
1919
/// </summary>
2020
public string AwsRequestId { get; set; }
2121

22+
/// <summary>
23+
/// Gets the tenant id for the Lambda function.
24+
/// </summary>
25+
public string TenantId { get; set; }
26+
2227
/// <summary>
2328
/// The current trace id if available.
2429
/// </summary>
@@ -46,4 +51,4 @@ public class MessageState
4651
public Exception Exception { get; set; }
4752
}
4853
}
49-
#endif
54+
#endif

‎Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
/*
1+
/*
22
* Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License").
@@ -47,6 +47,9 @@ public LambdaBootstrapTests()
4747
},
4848
{
4949
RuntimeApiHeaders.HeaderInvokedFunctionArn, new List<string> {"invoked_function_arn"}
50+
},
51+
{
52+
RuntimeApiHeaders.HeaderAwsTenantId, new List<string> {"tenant_id"}
5053
}
5154
};
5255
_testRuntimeApiClient = new TestRuntimeApiClient(_environmentVariables, headers);
@@ -82,6 +85,18 @@ public async Task NoInitializer()
8285
Assert.False(_testFunction.HandlerWasCalled);
8386
}
8487

88+
[Fact]
89+
public async Task ConfirmRuntimeApiHeadersInContext()
90+
{
91+
92+
using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerConfirmContextAsync, null))
93+
{
94+
bootstrap.Client = _testRuntimeApiClient;
95+
await bootstrap.RunAsync(_testFunction.CancellationSource.Token);
96+
}
97+
Assert.True(_testFunction.HandlerWasCalled);
98+
}
99+
85100
[Fact]
86101
public async Task InitializerHandlesExceptions()
87102
{

0 commit comments

Comments
 (0)