Skip to content

Commit

Permalink
Add Async Lifecycle Handlers (#3193)
Browse files Browse the repository at this point in the history
* Add async functions on Lifecycle handler

* Hook up async shutdown

* Hook up the rest of the lifecycle handler

* Hookup async stuff

* Add some docs for Lifecycle handlers

* Add async tests for lifecycle handler

* Fix the tests

* Clarify some of the docs

* Try and reduce a flaky test

* Fix the tests

* Redisable test as its extremely flaky
  • Loading branch information
0xTim committed May 15, 2024
1 parent 5bc1dfa commit 90da64a
Show file tree
Hide file tree
Showing 8 changed files with 248 additions and 24 deletions.
42 changes: 35 additions & 7 deletions Sources/Vapor/Application.swift
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,7 @@ public final class Application: Sendable {
self._storage = .init(.init(logger: logger))
self._lifecycle = .init(.init())
self.isBooted = .init(false)
self.core.initialize()
self.core.initialize(asyncEnvironment: async)
self.caches.initialize()
self.views.initialize()
self.passwords.use(.bcrypt)
Expand Down Expand Up @@ -218,7 +218,7 @@ public final class Application: Sendable {
/// If you want to run your ``Application`` indefinitely, or until your code shuts the application down,
/// use ``execute()`` instead.
public func startup() async throws {
try self.boot()
try await self.asyncBoot()

let combinedCommands = AsyncCommands(
commands: self.asyncCommands.commands.merging(self.commands.commands) { $1 },
Expand All @@ -231,6 +231,9 @@ public final class Application: Sendable {
try await self.console.run(combinedCommands, with: context)
}


@available(*, noasync, message: "This can potentially block the thread and should not be called in an async context", renamed: "asyncBoot()")
/// Called when the applications starts up, will trigger the lifecycle handlers
public func boot() throws {
try self.isBooted.withLockedValue { booted in
guard !booted else {
Expand All @@ -241,9 +244,31 @@ public final class Application: Sendable {
try self.lifecycle.handlers.forEach { try $0.didBoot(self) }
}
}

/// Called when the applications starts up, will trigger the lifecycle handlers. The asynchronous version of ``boot()``
public func asyncBoot() async throws {
self.isBooted.withLockedValue { booted in
guard !booted else {
return
}
booted = true
}
for handler in self.lifecycle.handlers {
try await handler.willBootAsync(self)
}
for handler in self.lifecycle.handlers {
try await handler.didBootAsync(self)
}
}

@available(*, noasync, message: "This can block the thread and should not be called in an async context", renamed: "asyncShutdown()")
public func shutdown() {
assert(!self.didShutdown, "Application has already shut down")
self.logger.debug("Application shutting down")

self.logger.trace("Shutting down providers")
self.lifecycle.handlers.reversed().forEach { $0.shutdown(self) }

triggerShutdown()

switch self.eventLoopGroupProvider {
Expand All @@ -263,6 +288,14 @@ public final class Application: Sendable {
}

public func asyncShutdown() async throws {
assert(!self.didShutdown, "Application has already shut down")
self.logger.debug("Application shutting down")

self.logger.trace("Shutting down providers")
for handler in self.lifecycle.handlers.reversed() {
await handler.shutdownAsync(self)
}

triggerShutdown()

switch self.eventLoopGroupProvider {
Expand All @@ -282,11 +315,6 @@ public final class Application: Sendable {
}

private func triggerShutdown() {
assert(!self.didShutdown, "Application has already shut down")
self.logger.debug("Application shutting down")

self.logger.trace("Shutting down providers")
self.lifecycle.handlers.reversed().forEach { $0.shutdown(self) }
self.lifecycle.handlers = []

self.logger.trace("Clearing Application storage")
Expand Down
18 changes: 16 additions & 2 deletions Sources/Vapor/Core/Core.swift
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,16 @@ extension Application {
try! application.threadPool.syncShutdownGracefully()
}
}

struct AsyncLifecycleHandler: Vapor.LifecycleHandler {
func shutdownAsync(_ application: Application) async {
do {
try await application.threadPool.shutdownGracefully()
} catch {
application.logger.debug("Failed to shutdown threadpool", metadata: ["error": "\(error)"])
}
}
}

struct Key: StorageKey {
typealias Value = Storage
Expand All @@ -112,9 +122,13 @@ extension Application {
return storage
}

func initialize() {
func initialize(asyncEnvironment: Bool) {
self.application.storage[Key.self] = .init()
self.application.lifecycle.use(LifecycleHandler())
if asyncEnvironment {
self.application.lifecycle.use(AsyncLifecycleHandler())
} else {
self.application.lifecycle.use(LifecycleHandler())
}
}
}
}
63 changes: 63 additions & 0 deletions Sources/Vapor/Utilities/LifecycleHandler.swift
Original file line number Diff line number Diff line change
@@ -1,11 +1,74 @@
/// Provides a way to hook into lifecycle events of a Vapor application. You can register
/// your handlers with the ``Application`` to be notified when the application
/// is about to start up, has started up and is about to shutdown
///
/// For example
/// ```swift
/// struct LifecycleLogger: LifecycleHander {
/// func willBootAsync(_ application: Application) async throws {
/// application.logger.info("Application about to boot up")
/// }
///
/// func didBootAsync(_ application: Application) async throws {
/// application.logger.info("Application has booted up")
/// }
///
/// func shutdownAsync(_ application: Application) async {
/// application.logger.info("Will shutdown")
/// }
/// }
/// ```
///
/// You can then register your handler with the application:
///
/// ```swift
/// application.lifecycle.use(LifecycleLogger())
/// ```
///
public protocol LifecycleHandler: Sendable {
/// Called when the application is about to boot up
func willBoot(_ application: Application) throws
/// Called when the application has booted up
func didBoot(_ application: Application) throws
/// Called when the application is about to shutdown
func shutdown(_ application: Application)
/// Called when the application is about to boot up. This is the asynchronous version
/// of ``willBoot(_:)-9zn``. When adopting the async APIs you should ensure you
/// provide a compatitble implementation for ``willBoot(_:)-8anu6`` as well if you
/// want to support older users still running in a non-async context
/// **Note** your application must be running in an asynchronous context and initialised with
/// ``Application/make(_:_:)`` for this handler to be called
func willBootAsync(_ application: Application) async throws
/// Called when the application is about to boot up. This is the asynchronous version
/// of ``didBoot(_:)-wfef``. When adopting the async APIs you should ensure you
/// provide a compatitble implementation for ``didBoot(_:)-wfef`` as well if you
/// want to support older users still running in a non-async context
/// **Note** your application must be running in an asynchronous context and initialised with
/// ``Application/make(_:_:)`` for this handler to be called
func didBootAsync(_ application: Application) async throws
/// Called when the application is about to boot up. This is the asynchronous version
/// of ``shutdown(_:)-2clwm``. When adopting the async APIs you should ensure you
/// provide a compatitble implementation for ``shutdown(_:)-2clwm`` as well if you
/// want to support older users still running in a non-async context
/// **Note** your application must be running in an asynchronous context and initialised with
/// ``Application/make(_:_:)`` for this handler to be called
func shutdownAsync(_ application: Application) async
}

extension LifecycleHandler {
public func willBoot(_ application: Application) throws { }
public func didBoot(_ application: Application) throws { }
public func shutdown(_ application: Application) { }

public func willBootAsync(_ application: Application) async throws {
try self.willBoot(application)
}

public func didBootAsync(_ application: Application) async throws {
try self.didBoot(application)
}

public func shutdownAsync(_ application: Application) async {
self.shutdown(application)
}
}
101 changes: 101 additions & 0 deletions Tests/VaporTests/ApplicationTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,29 @@ final class ApplicationTests: XCTestCase {
let willBootFlag: NIOLockedValueBox<Bool>
let didBootFlag: NIOLockedValueBox<Bool>
let shutdownFlag: NIOLockedValueBox<Bool>
let willBootAsyncFlag: NIOLockedValueBox<Bool>
let didBootAsyncFlag: NIOLockedValueBox<Bool>
let shutdownAsyncFlag: NIOLockedValueBox<Bool>

init() {
self.willBootFlag = .init(false)
self.didBootFlag = .init(false)
self.shutdownFlag = .init(false)
self.didBootAsyncFlag = .init(false)
self.willBootAsyncFlag = .init(false)
self.shutdownAsyncFlag = .init(false)
}

func willBootAsync(_ application: Application) async throws {
self.willBootAsyncFlag.withLockedValue { $0 = true }
}

func didBootAsync(_ application: Application) async throws {
self.didBootAsyncFlag.withLockedValue { $0 = true }
}

func shutdownAsync(_ application: Application) async {
self.shutdownAsyncFlag.withLockedValue { $0 = true }
}

func willBoot(_ application: Application) throws {
Expand All @@ -55,18 +73,101 @@ final class ApplicationTests: XCTestCase {
XCTAssertEqual(foo.willBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.willBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownAsyncFlag.withLockedValue({ $0 }), false)

try app.boot()

XCTAssertEqual(foo.willBootFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.didBootFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.shutdownFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.willBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownAsyncFlag.withLockedValue({ $0 }), false)

app.shutdown()

XCTAssertEqual(foo.willBootFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.didBootFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.shutdownFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.willBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownAsyncFlag.withLockedValue({ $0 }), false)
}

func testLifecycleHandlerAsync() async throws {
final class Foo: LifecycleHandler {
let willBootFlag: NIOLockedValueBox<Bool>
let didBootFlag: NIOLockedValueBox<Bool>
let shutdownFlag: NIOLockedValueBox<Bool>
let willBootAsyncFlag: NIOLockedValueBox<Bool>
let didBootAsyncFlag: NIOLockedValueBox<Bool>
let shutdownAsyncFlag: NIOLockedValueBox<Bool>

init() {
self.willBootFlag = .init(false)
self.didBootFlag = .init(false)
self.shutdownFlag = .init(false)
self.didBootAsyncFlag = .init(false)
self.willBootAsyncFlag = .init(false)
self.shutdownAsyncFlag = .init(false)
}

func willBootAsync(_ application: Application) async throws {
self.willBootAsyncFlag.withLockedValue { $0 = true }
}

func didBootAsync(_ application: Application) async throws {
self.didBootAsyncFlag.withLockedValue { $0 = true }
}

func shutdownAsync(_ application: Application) async {
self.shutdownAsyncFlag.withLockedValue { $0 = true }
}

func willBoot(_ application: Application) throws {
self.willBootFlag.withLockedValue { $0 = true }
}

func didBoot(_ application: Application) throws {
self.didBootFlag.withLockedValue { $0 = true }
}

func shutdown(_ application: Application) {
self.shutdownFlag.withLockedValue { $0 = true }
}
}

let app = try await Application.make(.testing)

let foo = Foo()
app.lifecycle.use(foo)

XCTAssertEqual(foo.willBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.willBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootAsyncFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownAsyncFlag.withLockedValue({ $0 }), false)

try await app.asyncBoot()

XCTAssertEqual(foo.willBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.willBootAsyncFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.didBootAsyncFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.shutdownAsyncFlag.withLockedValue({ $0 }), false)

try await app.asyncShutdown()

XCTAssertEqual(foo.willBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.didBootFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.shutdownFlag.withLockedValue({ $0 }), false)
XCTAssertEqual(foo.willBootAsyncFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.didBootAsyncFlag.withLockedValue({ $0 }), true)
XCTAssertEqual(foo.shutdownAsyncFlag.withLockedValue({ $0 }), true)
}

func testThrowDoesNotCrash() throws {
Expand Down
6 changes: 3 additions & 3 deletions Tests/VaporTests/AsyncClientTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ final class AsyncClientTests: XCTestCase {
}

remoteApp.environment.arguments = ["serve"]
try remoteApp.boot()
try await remoteApp.asyncBoot()
try await remoteApp.startup()

XCTAssertNotNil(remoteApp.http.server.shared.localAddress)
Expand Down Expand Up @@ -115,7 +115,7 @@ final class AsyncClientTests: XCTestCase {
}

func testClientBeforeSend() async throws {
try app.boot()
try await app.asyncBoot()

let res = try await app.client.post("http://localhost:\(remoteAppPort!)/anything") { req in
try req.content.encode(["hello": "world"])
Expand Down Expand Up @@ -143,7 +143,7 @@ final class AsyncClientTests: XCTestCase {
}

app.environment.arguments = ["serve"]
try app.boot()
try await app.asyncBoot()
try await app.startup()

XCTAssertNotNil(app.http.server.shared.localAddress)
Expand Down
Loading

0 comments on commit 90da64a

Please sign in to comment.