Skip to content

Commit 2cf4e1b

Browse files
tomerdfabianfett
tomerd
authored and
fabianfett
committed
fixup
1 parent 26f6a4b commit 2cf4e1b

File tree

3 files changed

+77
-82
lines changed

3 files changed

+77
-82
lines changed

‎Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

+74-68
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
//
1313
//===----------------------------------------------------------------------===//
1414

15+
#if DEBUG
1516
import Dispatch
1617
import Logging
1718
import NIO
@@ -26,8 +27,6 @@ import NIOHTTP1
2627
// callback(.success("Hello, \(payload)!"))
2728
// }
2829
// }
29-
30-
#if DEBUG
3130
extension Lambda {
3231
/// Execute code in the context of a mock Lambda server.
3332
///
@@ -88,20 +87,13 @@ private enum LocalLambda {
8887
}
8988

9089
final class HTTPHandler: ChannelInboundHandler {
91-
92-
enum InvocationState {
93-
case waitingForNextRequest
94-
case idle(EventLoopPromise<Pending>)
95-
case processing(Pending)
96-
}
97-
9890
public typealias InboundIn = HTTPServerRequestPart
9991
public typealias OutboundOut = HTTPServerResponsePart
10092

101-
private var processing = CircularBuffer<(head: HTTPRequestHead, body: ByteBuffer?)>()
102-
103-
private static var queue = [Pending]()
104-
private static var invocationState: InvocationState = .waitingForNextRequest
93+
private var pending = CircularBuffer<(head: HTTPRequestHead, body: ByteBuffer?)>()
94+
95+
private static var invocations = CircularBuffer<Invocation>()
96+
private static var invocationState = InvocationState.waitingForLambdaRequest
10597

10698
private let logger: Logger
10799
private let invocationEndpoint: String
@@ -116,92 +108,100 @@ private enum LocalLambda {
116108

117109
switch requestPart {
118110
case .head(let head):
119-
self.processing.append((head: head, body: nil))
111+
self.pending.append((head: head, body: nil))
120112
case .body(var buffer):
121-
var request = self.processing.removeFirst()
113+
var request = self.pending.removeFirst()
122114
if request.body == nil {
123115
request.body = buffer
124116
} else {
125117
request.body!.writeBuffer(&buffer)
126118
}
127-
self.processing.prepend(request)
119+
self.pending.prepend(request)
128120
case .end:
129-
let request = self.processing.removeFirst()
121+
let request = self.pending.removeFirst()
130122
self.processRequest(context: context, request: request)
131123
}
132124
}
133125

134126
func processRequest(context: ChannelHandlerContext, request: (head: HTTPRequestHead, body: ByteBuffer?)) {
135-
if request.head.uri.hasSuffix(self.invocationEndpoint) {
136-
if let work = request.body {
137-
let requestId = "\(DispatchTime.now().uptimeNanoseconds)" // FIXME:
138-
let promise = context.eventLoop.makePromise(of: Response.self)
139-
promise.futureResult.whenComplete { result in
140-
switch result {
141-
case .success(let response):
142-
self.writeResponse(context: context, response: response)
143-
case .failure:
144-
self.writeResponse(context: context, response: .init(status: .internalServerError))
145-
}
146-
}
147-
let pending = Pending(requestId: requestId, request: work, responsePromise: promise)
148-
switch Self.invocationState {
149-
case .idle(let promise):
150-
promise.succeed(pending)
151-
case .processing(_), .waitingForNextRequest:
152-
Self.queue.append(pending)
127+
switch (request.head.method, request.head.uri) {
128+
// this endpoint is called by the client invoking the lambda
129+
case (.POST, let url) where url.hasSuffix(self.invocationEndpoint):
130+
guard let work = request.body else {
131+
return self.writeResponse(context: context, response: .init(status: .badRequest))
132+
}
133+
let requestID = "\(DispatchTime.now().uptimeNanoseconds)" // FIXME:
134+
let promise = context.eventLoop.makePromise(of: Response.self)
135+
promise.futureResult.whenComplete { result in
136+
switch result {
137+
case .failure(let error):
138+
self.logger.error("invocation error: \(error)")
139+
self.writeResponse(context: context, response: .init(status: .internalServerError))
140+
case .success(let response):
141+
self.writeResponse(context: context, response: response)
153142
}
154143
}
155-
} else if request.head.uri.hasSuffix("/next") {
144+
let invocation = Invocation(requestID: requestID, request: work, responsePromise: promise)
145+
switch Self.invocationState {
146+
case .waitingForInvocation(let promise):
147+
promise.succeed(invocation)
148+
case .waitingForLambdaRequest, .waitingForLambdaResponse:
149+
Self.invocations.append(invocation)
150+
}
151+
// /next endpoint is called by the lambda polling for work
152+
case (.GET, let url) where url.hasSuffix(Consts.requestWorkURLSuffix):
156153
// check if our server is in the correct state
157-
guard case .waitingForNextRequest = Self.invocationState else {
158-
#warning("better error code?!")
159-
self.writeResponse(context: context, response: .init(status: .conflict))
154+
guard case .waitingForLambdaRequest = Self.invocationState else {
155+
self.logger.error("invalid invocation state \(Self.invocationState)")
156+
self.writeResponse(context: context, response: .init(status: .unprocessableEntity))
160157
return
161158
}
162-
159+
163160
// pop the first task from the queue
164-
switch !Self.queue.isEmpty ? Self.queue.removeFirst() : nil {
161+
switch Self.invocations.popFirst() {
165162
case .none:
166-
// if there is nothing in the queue, create a promise that we can succeed,
167-
// when we get a new task
168-
let promise = context.eventLoop.makePromise(of: Pending.self)
169-
promise.futureResult.whenComplete { (result) in
163+
// if there is nothing in the queue,
164+
// create a promise that we can fullfill when we get a new task
165+
let promise = context.eventLoop.makePromise(of: Invocation.self)
166+
promise.futureResult.whenComplete { result in
170167
switch result {
171168
case .failure(let error):
169+
self.logger.error("invocation error: \(error)")
172170
self.writeResponse(context: context, response: .init(status: .internalServerError))
173-
case .success(let pending):
174-
Self.invocationState = .processing(pending)
175-
self.writeResponse(context: context, response: pending.toResponse())
171+
case .success(let invocation):
172+
Self.invocationState = .waitingForLambdaResponse(invocation)
173+
self.writeResponse(context: context, response: invocation.makeResponse())
176174
}
177175
}
178-
Self.invocationState = .idle(promise)
179-
case .some(let pending):
176+
Self.invocationState = .waitingForInvocation(promise)
177+
case .some(let invocation):
180178
// if there is a task pending, we can immediatly respond with it.
181-
Self.invocationState = .processing(pending)
182-
self.writeResponse(context: context, response: pending.toResponse())
179+
Self.invocationState = .waitingForLambdaResponse(invocation)
180+
self.writeResponse(context: context, response: invocation.makeResponse())
183181
}
184-
185-
} else if request.head.uri.hasSuffix("/response") {
182+
// :requestID/response endpoint is called by the lambda posting the response
183+
case (.POST, let url) where url.hasSuffix(Consts.postResponseURLSuffix):
186184
let parts = request.head.uri.split(separator: "/")
187-
guard let requestId = parts.count > 2 ? String(parts[parts.count - 2]) : nil else {
185+
guard let requestID = parts.count > 2 ? String(parts[parts.count - 2]) : nil else {
188186
// the request is malformed, since we were expecting a requestId in the path
189187
return self.writeResponse(context: context, response: .init(status: .badRequest))
190188
}
191-
guard case .processing(let pending) = Self.invocationState else {
189+
guard case .waitingForLambdaResponse(let invocation) = Self.invocationState else {
192190
// a response was send, but we did not expect to receive one
193-
#warning("better error code?!")
194-
return self.writeResponse(context: context, response: .init(status: .conflict))
191+
self.logger.error("invalid invocation state \(Self.invocationState)")
192+
return self.writeResponse(context: context, response: .init(status: .unprocessableEntity))
195193
}
196-
guard requestId == pending.requestId else {
194+
guard requestID == invocation.requestID else {
197195
// the request's requestId is not matching the one we are expecting
196+
self.logger.error("invalid invocation state request ID \(requestID) does not match expected \(invocation.requestID)")
198197
return self.writeResponse(context: context, response: .init(status: .badRequest))
199198
}
200-
201-
pending.responsePromise.succeed(.init(status: .ok, body: request.body))
199+
200+
invocation.responsePromise.succeed(.init(status: .ok, body: request.body))
202201
self.writeResponse(context: context, response: .init(status: .accepted))
203-
Self.invocationState = .waitingForNextRequest
204-
} else {
202+
Self.invocationState = .waitingForLambdaRequest
203+
// unknown call
204+
default:
205205
self.writeResponse(context: context, response: .init(status: .notFound))
206206
}
207207
}
@@ -234,24 +234,30 @@ private enum LocalLambda {
234234
var body: ByteBuffer?
235235
}
236236

237-
struct Pending {
238-
let requestId: String
237+
struct Invocation {
238+
let requestID: String
239239
let request: ByteBuffer
240240
let responsePromise: EventLoopPromise<Response>
241-
242-
func toResponse() -> Response {
241+
242+
func makeResponse() -> Response {
243243
var response = Response()
244244
response.body = self.request
245245
// required headers
246246
response.headers = [
247-
(AmazonHeaders.requestID, self.requestId),
247+
(AmazonHeaders.requestID, self.requestID),
248248
(AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:\(Int16.random(in: Int16.min ... Int16.max)):function:custom-runtime"),
249249
(AmazonHeaders.traceID, "Root=\(Int16.random(in: Int16.min ... Int16.max));Parent=\(Int16.random(in: Int16.min ... Int16.max));Sampled=1"),
250250
(AmazonHeaders.deadline, "\(DispatchWallTime.distantFuture.millisSinceEpoch)"),
251251
]
252252
return response
253253
}
254254
}
255+
256+
enum InvocationState {
257+
case waitingForInvocation(EventLoopPromise<Invocation>)
258+
case waitingForLambdaRequest
259+
case waitingForLambdaResponse(Invocation)
260+
}
255261
}
256262

257263
enum ServerError: Error {

‎Sources/AWSLambdaRuntime/LambdaRunner.swift

+2-11
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,7 @@ extension Lambda {
4646
func run(logger: Logger, handler: Handler) -> EventLoopFuture<Void> {
4747
logger.debug("lambda invocation sequence starting")
4848
// 1. request work from lambda runtime engine
49-
return self.runtimeClient.requestWork(logger: logger).peekError { error -> Void in
50-
if case RuntimeError.badStatusCode(.noContent) = error {
51-
return
52-
}
49+
return self.runtimeClient.requestWork(logger: logger).peekError { error in
5350
logger.error("could not fetch work from lambda runtime engine: \(error)")
5451
}.flatMap { invocation, payload in
5552
// 2. send work to handler
@@ -67,13 +64,7 @@ extension Lambda {
6764
self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in
6865
logger.error("could not report results to lambda runtime engine: \(error)")
6966
}
70-
}.flatMapErrorThrowing { error in
71-
if case RuntimeError.badStatusCode(.noContent) = error {
72-
return ()
73-
}
74-
throw error
75-
}
76-
.always { result in
67+
}.always { result in
7768
// we are done!
7869
logger.log(level: result.successful ? .debug : .warning, "lambda invocation sequence completed \(result.successful ? "successfully" : "with failure")")
7970
}

‎Sources/StringSample/main.swift

+1-3
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,7 @@ struct Handler: EventLoopLambdaHandler {
2626
}
2727
}
2828

29-
try Lambda.withLocalServer {
30-
Lambda.run(Handler())
31-
}
29+
Lambda.run(Handler())
3230

3331
// MARK: - this can also be expressed as a closure:
3432

0 commit comments

Comments
 (0)