Skip to content

Commit 3ce0a87

Browse files
fabianfettfabianfett
authored
Attach requestID to logger (#497)
1 parent ae06d59 commit 3ce0a87

File tree

4 files changed

+160
-6
lines changed

4 files changed

+160
-6
lines changed

‎Sources/AWSLambdaRuntime/Lambda.swift

+2
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,11 @@ public enum Lambda {
3737
) async throws where Handler: StreamingLambdaHandler {
3838
var handler = handler
3939

40+
var logger = logger
4041
do {
4142
while !Task.isCancelled {
4243
let (invocation, writer) = try await runtimeClient.nextInvocation()
44+
logger[metadataKey: "aws-request-id"] = "\(invocation.metadata.requestID)"
4345

4446
do {
4547
try await handler.handle(

‎Tests/AWSLambdaRuntimeTests/CollectEverythingLogHandler.swift

+138
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the SwiftAWSLambdaRuntime open source project
4+
//
5+
// Copyright (c) 2025 Apple Inc. and the SwiftAWSLambdaRuntime project authors
6+
// Licensed under Apache License v2.0
7+
//
8+
// See LICENSE.txt for license information
9+
// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
10+
//
11+
// SPDX-License-Identifier: Apache-2.0
12+
//
13+
//===----------------------------------------------------------------------===//
14+
15+
import Logging
16+
import Synchronization
17+
import Testing
18+
19+
struct CollectEverythingLogHandler: LogHandler {
20+
var metadata: Logger.Metadata = [:]
21+
var logLevel: Logger.Level = .info
22+
let logStore: LogStore
23+
24+
final class LogStore: Sendable {
25+
struct Entry: Sendable {
26+
var level: Logger.Level
27+
var message: String
28+
var metadata: [String: String]
29+
}
30+
31+
let logs: Mutex<[Entry]> = .init([])
32+
33+
func append(level: Logger.Level, message: Logger.Message, metadata: Logger.Metadata?) {
34+
self.logs.withLock { entries in
35+
entries.append(
36+
Entry(
37+
level: level,
38+
message: message.description,
39+
metadata: metadata?.mapValues { $0.description } ?? [:]
40+
)
41+
)
42+
}
43+
}
44+
45+
func clear() {
46+
self.logs.withLock {
47+
$0.removeAll()
48+
}
49+
}
50+
51+
enum LogFieldExpectedValue: ExpressibleByStringLiteral, ExpressibleByStringInterpolation {
52+
case exactMatch(String)
53+
case beginsWith(String)
54+
case wildcard
55+
case predicate((String) -> Bool)
56+
57+
init(stringLiteral value: String) {
58+
self = .exactMatch(value)
59+
}
60+
}
61+
62+
@discardableResult
63+
func assertContainsLog(
64+
_ message: String,
65+
_ metadata: (String, LogFieldExpectedValue)...,
66+
sourceLocation: SourceLocation = #_sourceLocation
67+
) -> [Entry] {
68+
var candidates = self.getAllLogsWithMessage(message)
69+
if candidates.isEmpty {
70+
Issue.record("Logs do not contain entry with message: \(message)", sourceLocation: sourceLocation)
71+
return []
72+
}
73+
for (key, value) in metadata {
74+
var errorMsg: String
75+
switch value {
76+
case .wildcard:
77+
candidates = candidates.filter { $0.metadata.contains { $0.key == key } }
78+
errorMsg = "Logs do not contain entry with message: \(message) and metadata: \(key) *"
79+
case .predicate(let predicate):
80+
candidates = candidates.filter { $0.metadata[key].map(predicate) ?? false }
81+
errorMsg =
82+
"Logs do not contain entry with message: \(message) and metadata: \(key) matching predicate"
83+
case .beginsWith(let prefix):
84+
candidates = candidates.filter { $0.metadata[key]?.hasPrefix(prefix) ?? false }
85+
errorMsg = "Logs do not contain entry with message: \(message) and metadata: \(key), \(value)"
86+
case .exactMatch(let value):
87+
candidates = candidates.filter { $0.metadata[key] == value }
88+
errorMsg = "Logs do not contain entry with message: \(message) and metadata: \(key), \(value)"
89+
}
90+
if candidates.isEmpty {
91+
Issue.record("Error: \(errorMsg)", sourceLocation: sourceLocation)
92+
return []
93+
}
94+
}
95+
return candidates
96+
}
97+
98+
func assertDoesNotContainMessage(_ message: String, sourceLocation: SourceLocation = #_sourceLocation) {
99+
let candidates = self.getAllLogsWithMessage(message)
100+
if candidates.count > 0 {
101+
Issue.record("Logs contain entry with message: \(message)", sourceLocation: sourceLocation)
102+
}
103+
}
104+
105+
func getAllLogs() -> [Entry] {
106+
self.logs.withLock { $0 }
107+
}
108+
109+
func getAllLogsWithMessage(_ message: String) -> [Entry] {
110+
self.getAllLogs().filter { $0.message == message }
111+
}
112+
}
113+
114+
init(logStore: LogStore) {
115+
self.logStore = logStore
116+
}
117+
118+
func log(
119+
level: Logger.Level,
120+
message: Logger.Message,
121+
metadata: Logger.Metadata?,
122+
source: String,
123+
file: String,
124+
function: String,
125+
line: UInt
126+
) {
127+
self.logStore.append(level: level, message: message, metadata: self.metadata.merging(metadata ?? [:]) { $1 })
128+
}
129+
130+
subscript(metadataKey key: String) -> Logger.Metadata.Value? {
131+
get {
132+
self.metadata[key]
133+
}
134+
set {
135+
self.metadata[key] = newValue
136+
}
137+
}
138+
}

‎Tests/AWSLambdaRuntimeTests/LambdaMockClient.swift

+2-2
Original file line numberDiff line numberDiff line change
@@ -214,12 +214,12 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
214214
let eventProcessedHandler: CheckedContinuation<ByteBuffer, any Error>
215215
}
216216

217-
func invoke(event: ByteBuffer) async throws -> ByteBuffer {
217+
func invoke(event: ByteBuffer, requestID: String = UUID().uuidString) async throws -> ByteBuffer {
218218
try await withCheckedThrowingContinuation { eventProcessedHandler in
219219
do {
220220
let metadata = try InvocationMetadata(
221221
headers: .init([
222-
("Lambda-Runtime-Aws-Request-Id", "100"), // arbitrary values
222+
("Lambda-Runtime-Aws-Request-Id", "\(requestID)"), // arbitrary values
223223
("Lambda-Runtime-Deadline-Ms", "100"),
224224
("Lambda-Runtime-Invoked-Function-Arn", "100"),
225225
])

‎Tests/AWSLambdaRuntimeTests/LambdaRunLoopTests.swift

+18-4
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ struct LambdaRunLoopTests {
3232
responseWriter: some LambdaResponseStreamWriter,
3333
context: LambdaContext
3434
) async throws {
35+
context.logger.info("Test")
3536
try await responseWriter.writeAndFinish(event)
3637
}
3738
}
@@ -42,6 +43,7 @@ struct LambdaRunLoopTests {
4243
responseWriter: some LambdaResponseStreamWriter,
4344
context: LambdaContext
4445
) async throws {
46+
context.logger.info("Test")
4547
throw LambdaError.handlerError
4648
}
4749
}
@@ -54,16 +56,22 @@ struct LambdaRunLoopTests {
5456
let inputEvent = ByteBuffer(string: "Test Invocation Event")
5557

5658
try await withThrowingTaskGroup(of: Void.self) { group in
59+
let logStore = CollectEverythingLogHandler.LogStore()
5760
group.addTask {
5861
try await Lambda.runLoop(
5962
runtimeClient: self.mockClient,
6063
handler: self.mockEchoHandler,
61-
logger: Logger(label: "RunLoopTest")
64+
logger: Logger(
65+
label: "RunLoopTest",
66+
factory: { _ in CollectEverythingLogHandler(logStore: logStore) }
67+
)
6268
)
6369
}
6470

65-
let response = try await self.mockClient.invoke(event: inputEvent)
71+
let requestID = UUID().uuidString
72+
let response = try await self.mockClient.invoke(event: inputEvent, requestID: requestID)
6673
#expect(response == inputEvent)
74+
logStore.assertContainsLog("Test", ("aws-request-id", .exactMatch(requestID)))
6775

6876
group.cancelAll()
6977
}
@@ -73,20 +81,26 @@ struct LambdaRunLoopTests {
7381
let inputEvent = ByteBuffer(string: "Test Invocation Event")
7482

7583
await withThrowingTaskGroup(of: Void.self) { group in
84+
let logStore = CollectEverythingLogHandler.LogStore()
7685
group.addTask {
7786
try await Lambda.runLoop(
7887
runtimeClient: self.mockClient,
7988
handler: self.failingHandler,
80-
logger: Logger(label: "RunLoopTest")
89+
logger: Logger(
90+
label: "RunLoopTest",
91+
factory: { _ in CollectEverythingLogHandler(logStore: logStore) }
92+
)
8193
)
8294
}
8395

96+
let requestID = UUID().uuidString
8497
await #expect(
8598
throws: LambdaError.handlerError,
8699
performing: {
87-
try await self.mockClient.invoke(event: inputEvent)
100+
try await self.mockClient.invoke(event: inputEvent, requestID: requestID)
88101
}
89102
)
103+
logStore.assertContainsLog("Test", ("aws-request-id", .exactMatch(requestID)))
90104

91105
group.cancelAll()
92106
}

0 commit comments

Comments
 (0)