Skip to content

Commit 3aee427

Browse files
sebstosebsto
committed
use Atomic instead of Mutex.
Make the atomix check on run() instead of init()
1 parent 4ebf24f commit 3aee427

File tree

9 files changed

+117
-93
lines changed

9 files changed

+117
-93
lines changed

‎Examples/BackgroundTasks/Sources/main.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,5 +53,5 @@ struct BackgroundProcessingHandler: LambdaWithBackgroundProcessingHandler {
5353
}
5454

5555
let adapter = LambdaCodableAdapter(handler: BackgroundProcessingHandler())
56-
let runtime = try LambdaRuntime.init(handler: adapter)
56+
let runtime = LambdaRuntime.init(handler: adapter)
5757
try await runtime.run()

‎Examples/Streaming/Sources/main.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,5 +32,5 @@ struct SendNumbersWithPause: StreamingLambdaHandler {
3232
}
3333
}
3434

35-
let runtime = try LambdaRuntime.init(handler: SendNumbersWithPause())
35+
let runtime = LambdaRuntime.init(handler: SendNumbersWithPause())
3636
try await runtime.run()

‎Sources/AWSLambdaRuntime/ControlPlaneRequest.swift

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,19 +36,19 @@ package struct InvocationMetadata: Hashable {
3636
package let clientContext: String?
3737
package let cognitoIdentity: String?
3838

39-
package init(headers: HTTPHeaders) throws(LambdaRuntimeClientError) {
40-
guard let requestID = headers.first(name: AmazonHeaders.requestID), !requestID.isEmpty else {
41-
throw LambdaRuntimeClientError(code: .nextInvocationMissingHeaderRequestID)
39+
package init(headers: HTTPHeaders) throws(LambdaRuntimeError) {
40+
guard let requestID: String = headers.first(name: AmazonHeaders.requestID), !requestID.isEmpty else {
41+
throw LambdaRuntimeError(code: .nextInvocationMissingHeaderRequestID)
4242
}
4343

4444
guard let deadline = headers.first(name: AmazonHeaders.deadline),
4545
let unixTimeInMilliseconds = Int64(deadline)
4646
else {
47-
throw LambdaRuntimeClientError(code: .nextInvocationMissingHeaderDeadline)
47+
throw LambdaRuntimeError(code: .nextInvocationMissingHeaderDeadline)
4848
}
4949

5050
guard let invokedFunctionARN = headers.first(name: AmazonHeaders.invokedFunctionARN) else {
51-
throw LambdaRuntimeClientError(code: .nextInvocationMissingHeaderInvokeFuctionARN)
51+
throw LambdaRuntimeError(code: .nextInvocationMissingHeaderInvokeFuctionARN)
5252
}
5353

5454
self.requestID = requestID

‎Sources/AWSLambdaRuntime/FoundationSupport/Lambda+JSON.swift

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -108,11 +108,7 @@ extension LambdaRuntime {
108108
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
109109
)
110110

111-
do {
112-
try self.init(handler: handler)
113-
} catch {
114-
fatalError("Failed to initialize LambdaRuntime: \(error)")
115-
}
111+
self.init(handler: handler)
116112
}
117113

118114
/// Initialize an instance with a `LambdaHandler` defined in the form of a closure **with a `Void` return type**.
@@ -136,11 +132,7 @@ extension LambdaRuntime {
136132
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
137133
)
138134

139-
do {
140-
try self.init(handler: handler)
141-
} catch {
142-
fatalError("Failed to initialize LambdaRuntime: \(error)")
143-
}
135+
self.init(handler: handler)
144136
}
145137
}
146138
#endif // trait: FoundationJSONSupport

‎Sources/AWSLambdaRuntime/LambdaHandlers.swift

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -179,11 +179,7 @@ extension LambdaRuntime {
179179
public convenience init(
180180
body: @Sendable @escaping (ByteBuffer, LambdaResponseStreamWriter, LambdaContext) async throws -> Void
181181
) where Handler == StreamingClosureHandler {
182-
do {
183-
try self.init(handler: StreamingClosureHandler(body: body))
184-
} catch {
185-
fatalError("Failed to initialize LambdaRuntime: \(error)")
186-
}
182+
self.init(handler: StreamingClosureHandler(body: body))
187183
}
188184

189185
/// Initialize an instance with a ``LambdaHandler`` defined in the form of a closure **with a non-`Void` return type**, an encoder, and a decoder.
@@ -217,11 +213,7 @@ extension LambdaRuntime {
217213
decoder: decoder,
218214
handler: streamingAdapter
219215
)
220-
do {
221-
try self.init(handler: codableWrapper)
222-
} catch {
223-
fatalError("Failed to initialize LambdaRuntime: \(error)")
224-
}
216+
self.init(handler: codableWrapper)
225217
}
226218

227219
/// Initialize an instance with a ``LambdaHandler`` defined in the form of a closure **with a `Void` return type**, an encoder, and a decoder.
@@ -245,10 +237,6 @@ extension LambdaRuntime {
245237
decoder: decoder,
246238
handler: LambdaHandlerAdapter(handler: ClosureHandler(body: body))
247239
)
248-
do {
249-
try self.init(handler: handler)
250-
} catch {
251-
fatalError("Failed to initialize LambdaRuntime: \(error)")
252-
}
240+
self.init(handler: handler)
253241
}
254242
}

‎Sources/AWSLambdaRuntime/LambdaRuntime.swift

Lines changed: 36 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,10 @@ import FoundationEssentials
2323
import Foundation
2424
#endif
2525

26-
// This is our gardian to ensure only one LambdaRuntime is initialized
27-
// We use a Mutex here to ensure thread safety
28-
// We use Bool instead of LambdaRuntime<Handler> as the type here, as we don't know the concrete type that will be used
29-
private let _singleton = Mutex<Bool>(false)
30-
public enum LambdaRuntimeError: Error {
31-
case moreThanOneLambdaRuntimeInstance
32-
}
26+
// This is our gardian to ensure only one LambdaRuntime is running at the time
27+
// We use an Atomic here to ensure thread safety
28+
private let _isRunning = Atomic<Bool>(false)
29+
3330
// We need `@unchecked` Sendable here, as `NIOLockedValueBox` does not understand `sending` today.
3431
// We don't want to use `NIOLockedValueBox` here anyway. We would love to use Mutex here, but this
3532
// sadly crashes the compiler today (on Linux).
@@ -43,25 +40,7 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
4340
handler: sending Handler,
4441
eventLoop: EventLoop = Lambda.defaultEventLoop,
4542
logger: Logger = Logger(label: "LambdaRuntime")
46-
) throws {
47-
// technically, this initializer only throws LambdaRuntime Error but the below line crashes the compiler on Linux
48-
// https://github.com/swiftlang/swift/issues/80020
49-
// ) throws(LambdaRuntimeError) {
50-
51-
do {
52-
try _singleton.withLock {
53-
let alreadyCreated = $0
54-
guard alreadyCreated == false else {
55-
throw LambdaRuntimeError.moreThanOneLambdaRuntimeInstance
56-
}
57-
$0 = true
58-
}
59-
} catch _ as LambdaRuntimeError {
60-
throw LambdaRuntimeError.moreThanOneLambdaRuntimeInstance
61-
} catch {
62-
fatalError("An unknown error occurred: \(error)")
63-
}
64-
43+
) {
6544
self.handlerMutex = NIOLockedValueBox(handler)
6645
self.eventLoop = eventLoop
6746

@@ -74,15 +53,44 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
7453
self.logger.debug("LambdaRuntime initialized")
7554
}
7655

56+
/// Make sure only one run() is called at a time
7757
public func run() async throws {
58+
59+
// we use an atomic global variable to ensure only one LambdaRuntime is running at the time
60+
let (_, original) = _isRunning.compareExchange(expected: false, desired: true, ordering: .relaxed)
61+
62+
// if the original value was already true, run() is already running
63+
if original {
64+
throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
65+
}
66+
67+
try await withTaskCancellationHandler {
68+
// call the internal _run() method
69+
do {
70+
try await self._run()
71+
} catch {
72+
// when we catch an error, flip back the global variable to false
73+
_isRunning.store(false, ordering: .relaxed)
74+
throw error
75+
}
76+
} onCancel: {
77+
// when task is cancelled, flip back the global variable to false
78+
_isRunning.store(false, ordering: .relaxed)
79+
}
80+
81+
// when we're done without error and without cancellation, flip back the global variable to false
82+
_isRunning.store(false, ordering: .relaxed)
83+
}
84+
85+
private func _run() async throws {
7886
let handler = self.handlerMutex.withLockedValue { handler in
7987
let result = handler
8088
handler = nil
8189
return result
8290
}
8391

8492
guard let handler else {
85-
throw LambdaRuntimeClientError(code: .runtimeCanOnlyBeStartedOnce)
93+
throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
8694
}
8795

8896
// are we running inside an AWS Lambda runtime environment ?
@@ -92,7 +100,7 @@ public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: St
92100

93101
let ipAndPort = runtimeEndpoint.split(separator: ":", maxSplits: 1)
94102
let ip = String(ipAndPort[0])
95-
guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeClientError(code: .invalidPort) }
103+
guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeError(code: .invalidPort) }
96104

97105
try await LambdaRuntimeClient.withRuntimeClient(
98106
configuration: .init(ip: ip, port: port),

‎Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
134134

135135
case .connecting(let continuations):
136136
for continuation in continuations {
137-
continuation.resume(throwing: LambdaRuntimeClientError(code: .closingRuntimeClient))
137+
continuation.resume(throwing: LambdaRuntimeError(code: .closingRuntimeClient))
138138
}
139139
self.connectionState = .connecting([])
140140

@@ -173,7 +173,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
173173
private func write(_ buffer: NIOCore.ByteBuffer) async throws {
174174
switch self.lambdaState {
175175
case .idle, .sentResponse:
176-
throw LambdaRuntimeClientError(code: .writeAfterFinishHasBeenSent)
176+
throw LambdaRuntimeError(code: .writeAfterFinishHasBeenSent)
177177

178178
case .waitingForNextInvocation:
179179
fatalError("Invalid state: \(self.lambdaState)")
@@ -194,7 +194,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
194194
private func writeAndFinish(_ buffer: NIOCore.ByteBuffer?) async throws {
195195
switch self.lambdaState {
196196
case .idle, .sentResponse:
197-
throw LambdaRuntimeClientError(code: .finishAfterFinishHasBeenSent)
197+
throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)
198198

199199
case .waitingForNextInvocation:
200200
fatalError("Invalid state: \(self.lambdaState)")
@@ -261,7 +261,7 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
261261
case (.connecting(let array), .notClosing):
262262
self.connectionState = .disconnected
263263
for continuation in array {
264-
continuation.resume(throwing: LambdaRuntimeClientError(code: .lostConnectionToControlPlane))
264+
continuation.resume(throwing: LambdaRuntimeError(code: .lostConnectionToControlPlane))
265265
}
266266

267267
case (.connecting(let array), .closing(let continuation)):
@@ -394,7 +394,7 @@ extension LambdaRuntimeClient: LambdaChannelHandlerDelegate {
394394
}
395395

396396
for continuation in continuations {
397-
continuation.resume(throwing: LambdaRuntimeClientError(code: .connectionToControlPlaneLost))
397+
continuation.resume(throwing: LambdaRuntimeError(code: .connectionToControlPlaneLost))
398398
}
399399

400400
case .connected(let stateChannel, _):
@@ -489,7 +489,7 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
489489
fatalError("Invalid state: \(self.state)")
490490

491491
case .disconnected:
492-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
492+
throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
493493
}
494494
}
495495

@@ -528,10 +528,10 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
528528
)
529529

530530
case .disconnected:
531-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
531+
throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
532532

533533
case .closing:
534-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)
534+
throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)
535535
}
536536
}
537537

@@ -553,13 +553,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
553553

554554
case .connected(_, .idle),
555555
.connected(_, .sentResponse):
556-
throw LambdaRuntimeClientError(code: .writeAfterFinishHasBeenSent)
556+
throw LambdaRuntimeError(code: .writeAfterFinishHasBeenSent)
557557

558558
case .disconnected:
559-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
559+
throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
560560

561561
case .closing:
562-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)
562+
throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)
563563
}
564564
}
565565

@@ -586,13 +586,13 @@ private final class LambdaChannelHandler<Delegate: LambdaChannelHandlerDelegate>
586586
}
587587

588588
case .connected(_, .sentResponse):
589-
throw LambdaRuntimeClientError(code: .finishAfterFinishHasBeenSent)
589+
throw LambdaRuntimeError(code: .finishAfterFinishHasBeenSent)
590590

591591
case .disconnected:
592-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneLost)
592+
throw LambdaRuntimeError(code: .connectionToControlPlaneLost)
593593

594594
case .closing:
595-
throw LambdaRuntimeClientError(code: .connectionToControlPlaneGoingAway)
595+
throw LambdaRuntimeError(code: .connectionToControlPlaneGoingAway)
596596
}
597597
}
598598

@@ -759,7 +759,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {
759759
self.delegate.connectionWillClose(channel: context.channel)
760760
context.close(promise: nil)
761761
continuation.resume(
762-
throwing: LambdaRuntimeClientError(code: .invocationMissingMetadata, underlying: error)
762+
throwing: LambdaRuntimeError(code: .invocationMissingMetadata, underlying: error)
763763
)
764764
}
765765

@@ -769,7 +769,7 @@ extension LambdaChannelHandler: ChannelInboundHandler {
769769
continuation.resume()
770770
} else {
771771
self.state = .connected(context, .idle)
772-
continuation.resume(throwing: LambdaRuntimeClientError(code: .unexpectedStatusCodeForRequest))
772+
continuation.resume(throwing: LambdaRuntimeError(code: .unexpectedStatusCodeForRequest))
773773
}
774774

775775
case .disconnected, .closing, .connected(_, _):

‎Sources/AWSLambdaRuntime/LambdaRuntimeClientError.swift renamed to ‎Sources/AWSLambdaRuntime/LambdaRuntimeError.swift

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,10 @@
1212
//
1313
//===----------------------------------------------------------------------===//
1414

15-
package struct LambdaRuntimeClientError: Error {
16-
package enum Code {
15+
public struct LambdaRuntimeError: Error {
16+
public enum Code: Sendable {
17+
18+
/// internal error codes for LambdaRuntimeClient
1719
case closingRuntimeClient
1820

1921
case connectionToControlPlaneLost
@@ -32,14 +34,17 @@ package struct LambdaRuntimeClientError: Error {
3234
case missingLambdaRuntimeAPIEnvironmentVariable
3335
case runtimeCanOnlyBeStartedOnce
3436
case invalidPort
37+
38+
/// public error codes for LambdaRuntime
39+
case moreThanOneLambdaRuntimeInstance
3540
}
3641

3742
package init(code: Code, underlying: (any Error)? = nil) {
3843
self.code = code
3944
self.underlying = underlying
4045
}
4146

42-
package var code: Code
43-
package var underlying: (any Error)?
47+
public var code: Code
48+
public var underlying: (any Error)?
4449

4550
}

0 commit comments

Comments
 (0)