Attach requestID to logger

Sources/AWSLambdaRuntime/Lambda.swift

@@ -37,9 +37,11 @@ public enum Lambda {
     ) async throws where Handler: StreamingLambdaHandler {
         var handler = handler
 
+        var logger = logger
         do {
             while !Task.isCancelled {
                 let (invocation, writer) = try await runtimeClient.nextInvocation()
+                logger[metadataKey: "aws-request-id"] = "\(invocation.metadata.requestID)"
 
                 do {
                     try await handler.handle(

Tests/AWSLambdaRuntimeTests/CollectEverythingLogHandler.swift

@@ -0,0 +1,138 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright (c) 2025 Apple Inc. and the SwiftAWSLambdaRuntime project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Logging
+import Synchronization
+import Testing
+
+struct CollectEverythingLogHandler: LogHandler {
+    var metadata: Logger.Metadata = [:]
+    var logLevel: Logger.Level = .info
+    let logStore: LogStore
+
+    final class LogStore: Sendable {
+        struct Entry: Sendable {
+            var level: Logger.Level
+            var message: String
+            var metadata: [String: String]
+        }
+
+        let logs: Mutex<[Entry]> = .init([])
+
+        func append(level: Logger.Level, message: Logger.Message, metadata: Logger.Metadata?) {
+            self.logs.withLock { entries in
+                entries.append(
+                    Entry(
+                        level: level,
+                        message: message.description,
+                        metadata: metadata?.mapValues { $0.description } ?? [:]
+                    )
+                )
+            }
+        }
+
+        func clear() {
+            self.logs.withLock {
+                $0.removeAll()
+            }
+        }
+
+        enum LogFieldExpectedValue: ExpressibleByStringLiteral, ExpressibleByStringInterpolation {
+            case exactMatch(String)
+            case beginsWith(String)
+            case wildcard
+            case predicate((String) -> Bool)
+
+            init(stringLiteral value: String) {
+                self = .exactMatch(value)
+            }
+        }
+
+        @discardableResult
+        func assertContainsLog(
+            _ message: String,
+            _ metadata: (String, LogFieldExpectedValue)...,
+            sourceLocation: SourceLocation = #_sourceLocation
+        ) -> [Entry] {
+            var candidates = self.getAllLogsWithMessage(message)
+            if candidates.isEmpty {
+                Issue.record("Logs do not contain entry with message: \(message)", sourceLocation: sourceLocation)
+                return []
+            }
+            for (key, value) in metadata {
+                var errorMsg: String
+                switch value {
+                case .wildcard:
+                    candidates = candidates.filter { $0.metadata.contains { $0.key == key } }
+                    errorMsg = "Logs do not contain entry with message: \(message) and metadata: \(key) *"
+                case .predicate(let predicate):
+                    candidates = candidates.filter { $0.metadata[key].map(predicate) ?? false }
+                    errorMsg =
+                        "Logs do not contain entry with message: \(message) and metadata: \(key) matching predicate"
+                case .beginsWith(let prefix):
+                    candidates = candidates.filter { $0.metadata[key]?.hasPrefix(prefix) ?? false }
+                    errorMsg = "Logs do not contain entry with message: \(message) and metadata: \(key), \(value)"
+                case .exactMatch(let value):
+                    candidates = candidates.filter { $0.metadata[key] == value }
+                    errorMsg = "Logs do not contain entry with message: \(message) and metadata: \(key), \(value)"
+                }
+                if candidates.isEmpty {
+                    Issue.record("Error: \(errorMsg)", sourceLocation: sourceLocation)
+                    return []
+                }
+            }
+            return candidates
+        }
+
+        func assertDoesNotContainMessage(_ message: String, sourceLocation: SourceLocation = #_sourceLocation) {
+            let candidates = self.getAllLogsWithMessage(message)
+            if candidates.count > 0 {
+                Issue.record("Logs contain entry with message: \(message)", sourceLocation: sourceLocation)
+            }
+        }
+
+        func getAllLogs() -> [Entry] {
+            self.logs.withLock { $0 }
+        }
+
+        func getAllLogsWithMessage(_ message: String) -> [Entry] {
+            self.getAllLogs().filter { $0.message == message }
+        }
+    }
+
+    init(logStore: LogStore) {
+        self.logStore = logStore
+    }
+
+    func log(
+        level: Logger.Level,
+        message: Logger.Message,
+        metadata: Logger.Metadata?,
+        source: String,
+        file: String,
+        function: String,
+        line: UInt
+    ) {
+        self.logStore.append(level: level, message: message, metadata: self.metadata.merging(metadata ?? [:]) { $1 })
+    }
+
+    subscript(metadataKey key: String) -> Logger.Metadata.Value? {
+        get {
+            self.metadata[key]
+        }
+        set {
+            self.metadata[key] = newValue
+        }
+    }
+}

Tests/AWSLambdaRuntimeTests/LambdaMockClient.swift

@@ -214,12 +214,12 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
         let eventProcessedHandler: CheckedContinuation<ByteBuffer, any Error>
     }
 
-    func invoke(event: ByteBuffer) async throws -> ByteBuffer {
+    func invoke(event: ByteBuffer, requestID: String = UUID().uuidString) async throws -> ByteBuffer {
         try await withCheckedThrowingContinuation { eventProcessedHandler in
             do {
                 let metadata = try InvocationMetadata(
                     headers: .init([
-                        ("Lambda-Runtime-Aws-Request-Id", "100"),  // arbitrary values
+                        ("Lambda-Runtime-Aws-Request-Id", "\(requestID)"),  // arbitrary values
                         ("Lambda-Runtime-Deadline-Ms", "100"),
                         ("Lambda-Runtime-Invoked-Function-Arn", "100"),
                     ])

Tests/AWSLambdaRuntimeTests/LambdaRunLoopTests.swift

@@ -32,6 +32,7 @@ struct LambdaRunLoopTests {
             responseWriter: some LambdaResponseStreamWriter,
             context: LambdaContext
         ) async throws {
+            context.logger.info("Test")
             try await responseWriter.writeAndFinish(event)
         }
     }
@@ -42,6 +43,7 @@ struct LambdaRunLoopTests {
             responseWriter: some LambdaResponseStreamWriter,
             context: LambdaContext
         ) async throws {
+            context.logger.info("Test")
             throw LambdaError.handlerError
         }
     }
@@ -54,16 +56,22 @@ struct LambdaRunLoopTests {
         let inputEvent = ByteBuffer(string: "Test Invocation Event")
 
         try await withThrowingTaskGroup(of: Void.self) { group in
+            let logStore = CollectEverythingLogHandler.LogStore()
             group.addTask {
                 try await Lambda.runLoop(
                     runtimeClient: self.mockClient,
                     handler: self.mockEchoHandler,
-                    logger: Logger(label: "RunLoopTest")
+                    logger: Logger(
+                        label: "RunLoopTest",
+                        factory: { _ in CollectEverythingLogHandler(logStore: logStore) }
+                    )
                 )
             }
 
-            let response = try await self.mockClient.invoke(event: inputEvent)
+            let requestID = UUID().uuidString
+            let response = try await self.mockClient.invoke(event: inputEvent, requestID: requestID)
             #expect(response == inputEvent)
+            logStore.assertContainsLog("Test", ("aws-request-id", .exactMatch(requestID)))
 
             group.cancelAll()
         }
@@ -73,20 +81,26 @@ struct LambdaRunLoopTests {
         let inputEvent = ByteBuffer(string: "Test Invocation Event")
 
         await withThrowingTaskGroup(of: Void.self) { group in
+            let logStore = CollectEverythingLogHandler.LogStore()
             group.addTask {
                 try await Lambda.runLoop(
                     runtimeClient: self.mockClient,
                     handler: self.failingHandler,
-                    logger: Logger(label: "RunLoopTest")
+                    logger: Logger(
+                        label: "RunLoopTest",
+                        factory: { _ in CollectEverythingLogHandler(logStore: logStore) }
+                    )
                 )
             }
 
+            let requestID = UUID().uuidString
             await #expect(
                 throws: LambdaError.handlerError,
                 performing: {
-                    try await self.mockClient.invoke(event: inputEvent)
+                    try await self.mockClient.invoke(event: inputEvent, requestID: requestID)
                 }
             )
+            logStore.assertContainsLog("Test", ("aws-request-id", .exactMatch(requestID)))
 
             group.cancelAll()
         }