Add Deadline (#27)

Motivation:
- As a developer I want to be able to know how much time I have left to execute my lambda, before it times out

Changes:
- storing the deadline on Context and Invocation
- added getRemainingTime() to Context which returns a TimeAmount
- Renamed `Invocation.deadline` to `Invocation.deadlineInMillisSinceEpoch` to better reflect 
- fixed MockLambdaServer to not return “keep-alive” anymore

Co-authored-by: tom doron <tomer@apple.com>

Sources/MockServer/main.swift

@@ -108,11 +108,12 @@ internal final class HTTPHandler: ChannelInboundHandler {
             case .json:
                 responseBody = "{ \"body\": \"\(requestId)\" }"
             }
+            let deadline = Int64(Date(timeIntervalSinceNow: 60).timeIntervalSince1970 * 1000)
             responseHeaders = [
                 (AmazonHeaders.requestID, requestId),
                 (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
                 (AmazonHeaders.traceID, "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"),
-                (AmazonHeaders.deadline, String(Date(timeIntervalSinceNow: 60).timeIntervalSince1970 * 1000)),
+                (AmazonHeaders.deadline, String(deadline)),
             ]
         } else if request.head.uri.hasSuffix("/response") {
             responseStatus = .accepted

Sources/SwiftAwsLambda/Lambda.swift

@@ -78,20 +78,31 @@ public enum Lambda {
     }
 
     public class Context {
-        // from aws
+        /// The request ID, which identifies the request that triggered the function invocation.
         public let requestId: String
+
+        /// The AWS X-Ray tracing header.
         public let traceId: String
+
+        /// The ARN of the Lambda function, version, or alias that's specified in the invocation.
         public let invokedFunctionArn: String
-        public let deadline: String
+
+        /// The timestamp that the function times out
+        public let deadline: DispatchWallTime
+
+        /// For invocations from the AWS Mobile SDK, data about the Amazon Cognito identity provider.
         public let cognitoIdentity: String?
+
+        /// For invocations from the AWS Mobile SDK, data about the client application and device.
         public let clientContext: String?
-        // utility
+
+        /// a logger to log
         public let logger: Logger
 
         internal init(requestId: String,
                       traceId: String,
                       invokedFunctionArn: String,
-                      deadline: String,
+                      deadline: DispatchWallTime,
                       cognitoIdentity: String? = nil,
                       clientContext: String? = nil,
                       logger: Logger) {
@@ -107,6 +118,14 @@ public enum Lambda {
             logger[metadataKey: "awsTraceId"] = .string(traceId)
             self.logger = logger
         }
+
+        public func getRemainingTime() -> TimeAmount {
+            let deadline = self.deadline.millisSinceEpoch
+            let now = DispatchWallTime.now().millisSinceEpoch
+
+            let remaining = deadline - now
+            return .milliseconds(remaining)
+        }
     }
 
     private final class Lifecycle {

Sources/SwiftAwsLambda/LambdaRunner.swift

@@ -130,7 +130,7 @@ private extension Lambda.Context {
         self.init(requestId: invocation.requestId,
                   traceId: invocation.traceId,
                   invokedFunctionArn: invocation.invokedFunctionArn,
-                  deadline: invocation.deadlineDate,
+                  deadline: DispatchWallTime(millisSinceEpoch: invocation.deadlineInMillisSinceEpoch),
                   cognitoIdentity: invocation.cognitoIdentity,
                   clientContext: invocation.clientContext,
                   logger: logger)

Sources/SwiftAwsLambda/LambdaRuntimeClient.swift

@@ -181,7 +181,7 @@ private extension HTTPClient.Response {
 
 internal struct Invocation {
     let requestId: String
-    let deadlineDate: String
+    let deadlineInMillisSinceEpoch: Int64
     let invokedFunctionArn: String
     let traceId: String
     let clientContext: String?
@@ -192,7 +192,8 @@ internal struct Invocation {
             throw LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.requestID)
         }
 
-        guard let unixTimeMilliseconds = headers.first(name: AmazonHeaders.deadline) else {
+        guard let deadline = headers.first(name: AmazonHeaders.deadline),
+            let unixTimeInMilliseconds = Int64(deadline) else {
             throw LambdaRuntimeClientError.invocationMissingHeader(AmazonHeaders.deadline)
         }
 
@@ -205,7 +206,7 @@ internal struct Invocation {
         }
 
         self.requestId = requestId
-        self.deadlineDate = unixTimeMilliseconds
+        self.deadlineInMillisSinceEpoch = unixTimeInMilliseconds
         self.invokedFunctionArn = invokedFunctionArn
         self.traceId = traceId
         self.clientContext = headers["Lambda-Runtime-Client-Context"].first

Sources/SwiftAwsLambda/Utils.swift

@@ -64,3 +64,16 @@ internal enum Signal: Int32 {
     case ALRM = 14
     case TERM = 15
 }
+
+internal extension DispatchWallTime {
+    init(millisSinceEpoch: Int64) {
+        let nanoSinceEpoch = UInt64(millisSinceEpoch) * 1_000_000
+        let seconds = UInt64(nanoSinceEpoch / 1_000_000_000)
+        let nanoseconds = nanoSinceEpoch - (seconds * 1_000_000_000)
+        self.init(timespec: timespec(tv_sec: Int(seconds), tv_nsec: Int(nanoseconds)))
+    }
+
+    var millisSinceEpoch: Int64 {
+        return Int64(bitPattern: self.rawValue) / -1_000_000
+    }
+}

Tests/SwiftAwsLambdaTests/LambdaTest+XCTest.swift

@@ -36,6 +36,8 @@ extension LambdaTest {
             ("testBigPayload", testBigPayload),
             ("testKeepAliveServer", testKeepAliveServer),
             ("testNoKeepAliveServer", testNoKeepAliveServer),
+            ("testDeadline", testDeadline),
+            ("testGetRemainingTime", testGetRemainingTime),
         ]
     }
 }

Tests/SwiftAwsLambdaTests/LambdaTest.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Logging
 import NIO
 @testable import SwiftAwsLambda
 import XCTest
@@ -212,6 +213,58 @@ class LambdaTest: XCTestCase {
         let result = Lambda.run(handler: EchoHandler(), configuration: configuration)
         assertLambdaLifecycleResult(result, shoudHaveRun: maxTimes)
     }
+
+    func testDeadline() {
+        let delta = Int.random(in: 1 ... 600)
+
+        let milli1 = Date(timeIntervalSinceNow: Double(delta)).millisSinceEpoch
+        let milli2 = (DispatchWallTime.now() + .seconds(delta)).millisSinceEpoch
+        XCTAssertEqual(Double(milli1), Double(milli2), accuracy: 2.0)
+
+        let now1 = DispatchWallTime.now()
+        let now2 = DispatchWallTime(millisSinceEpoch: Date().millisSinceEpoch)
+        XCTAssertEqual(Double(now2.rawValue), Double(now1.rawValue), accuracy: 2_000_000.0)
+
+        let future1 = DispatchWallTime.now() + .seconds(delta)
+        let future2 = DispatchWallTime(millisSinceEpoch: Date(timeIntervalSinceNow: Double(delta)).millisSinceEpoch)
+        XCTAssertEqual(Double(future1.rawValue), Double(future2.rawValue), accuracy: 2_000_000.0)
+
+        let past1 = DispatchWallTime.now() - .seconds(delta)
+        let past2 = DispatchWallTime(millisSinceEpoch: Date(timeIntervalSinceNow: Double(-delta)).millisSinceEpoch)
+        XCTAssertEqual(Double(past1.rawValue), Double(past2.rawValue), accuracy: 2_000_000.0)
+
+        let logger = Logger(label: "test")
+        let context = Lambda.Context(requestId: UUID().uuidString,
+                                     traceId: UUID().uuidString,
+                                     invokedFunctionArn: UUID().uuidString,
+                                     deadline: .now() + .seconds(1),
+                                     cognitoIdentity: nil,
+                                     clientContext: nil,
+                                     logger: logger)
+        XCTAssertGreaterThan(context.deadline, .now())
+
+        let expiredContext = Lambda.Context(requestId: UUID().uuidString,
+                                            traceId: UUID().uuidString,
+                                            invokedFunctionArn: UUID().uuidString,
+                                            deadline: .now() - .seconds(1),
+                                            cognitoIdentity: nil,
+                                            clientContext: nil,
+                                            logger: logger)
+        XCTAssertLessThan(expiredContext.deadline, .now())
+    }
+
+    func testGetRemainingTime() {
+        let logger = Logger(label: "test")
+        let context = Lambda.Context(requestId: UUID().uuidString,
+                                     traceId: UUID().uuidString,
+                                     invokedFunctionArn: UUID().uuidString,
+                                     deadline: .now() + .seconds(1),
+                                     cognitoIdentity: nil,
+                                     clientContext: nil,
+                                     logger: logger)
+        XCTAssertLessThanOrEqual(context.getRemainingTime(), .seconds(1))
+        XCTAssertGreaterThan(context.getRemainingTime(), .milliseconds(800))
+    }
 }
 
 private struct Behavior: LambdaServerBehavior {

Tests/SwiftAwsLambdaTests/MockLambdaServer.swift

@@ -140,11 +140,12 @@ internal final class HTTPHandler: ChannelInboundHandler {
                 }
                 responseStatus = .ok
                 responseBody = result
+                let deadline = Date(timeIntervalSinceNow: 60).millisSinceEpoch
                 responseHeaders = [
                     (AmazonHeaders.requestID, requestId),
                     (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:123456789012:function:custom-runtime"),
                     (AmazonHeaders.traceID, "Root=1-5bef4de7-ad49b0e87f6ef6c87fc2e700;Parent=9a9197af755a6419;Sampled=1"),
-                    (AmazonHeaders.deadline, String(Date(timeIntervalSinceNow: 60).timeIntervalSince1970 * 1000)),
+                    (AmazonHeaders.deadline, String(deadline)),
                 ]
             case .failure(let error):
                 responseStatus = .init(statusCode: error.rawValue)
@@ -181,7 +182,9 @@ internal final class HTTPHandler: ChannelInboundHandler {
     func writeResponse(context: ChannelHandlerContext, status: HTTPResponseStatus, headers: [(String, String)]? = nil, body: String? = nil) {
         var headers = HTTPHeaders(headers ?? [])
         headers.add(name: "Content-Length", value: "\(body?.utf8.count ?? 0)")
-        headers.add(name: "Connection", value: self.keepAlive ? "keep-alive" : "close")
+        if !self.keepAlive {
+            headers.add(name: "Connection", value: "close")
+        }
         let head = HTTPResponseHead(version: HTTPVersion(major: 1, minor: 1), status: status, headers: headers)
 
         context.write(wrapOutboundOut(.head(head))).whenFailure { error in

Tests/SwiftAwsLambdaTests/Utils.swift

@@ -93,3 +93,9 @@ struct TestError: Error, Equatable, CustomStringConvertible {
         self.description = description
     }
 }
+
+internal extension Date {
+    var millisSinceEpoch: Int64 {
+        return Int64(self.timeIntervalSince1970 * 1000)
+    }
+}