Skip to content

Commit

Permalink
Add shutdownAsync and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
0xTim committed Jul 31, 2024
1 parent 15b3fb7 commit b239bb0
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 0 deletions.
39 changes: 39 additions & 0 deletions Sources/AsyncKit/ConnectionPool/EventLoopGroupConnectionPool.swift
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ public final class EventLoopGroupConnectionPool<Source> where Source: Connection
///
/// - Warning: This method is soft-deprecated. Use `syncShutdownGracefully()` or
/// `shutdownGracefully()` instead.
@available(*, noasync, message: "This calls wait() and should not be used in an async context", renamed: "shutdownAsync()")
public func shutdown() {
// synchronize access to closing
guard self.lock.withLock({
Expand All @@ -190,6 +191,43 @@ public final class EventLoopGroupConnectionPool<Source> where Source: Connection
}
}

/// Closes the connection pool.
///
/// All available connections will be closed immediately. Any connections still in use will be
/// closed as soon as they are returned to the pool. Once closed, the pool can not be used to
/// create new connections.
///
/// Connection pools must be closed before they deinitialize.
///
/// This method shuts down asynchronously, waiting for all connection closures to complete before
/// returning.
///
/// - Warning: The pool is always fully shut down once this method returns, even if an error is
/// thrown. All errors are purely advisory.
public func shutdownAsync() async throws {
// synchronize access to closing
guard self.lock.withLock({
// check to make sure we aren't double closing
guard !self.didShutdown else {
return false
}
self.didShutdown = true
self.logger.debug("Connection pool shutting down, closing each event loop's storage")
return true
}) else {
return
}

// shutdown all pools
for pool in self.storage.values {
do {
try await pool.close().get()
} catch {
self.logger.error("Failed shutting down event loop pool: \(error)")
}
}
}

/// Closes the connection pool.
///
/// All available connections will be closed immediately. Any connections still in use will be
Expand All @@ -203,6 +241,7 @@ public final class EventLoopGroupConnectionPool<Source> where Source: Connection
///
/// - Warning: The pool is always fully shut down once this method returns, even if an error is
/// thrown. All errors are purely advisory.
@available(*, noasync, message: "This calls wait() and should not be used in an async context", renamed: "shutdownAsync()")
public func syncShutdownGracefully() throws {
// - TODO: Does this need to assert "not on any EventLoop", as `EventLoopGroup.syncShutdownGracefully()` does?
var possibleError: Error? = nil
Expand Down
20 changes: 20 additions & 0 deletions Tests/AsyncKitTests/ConnectionPoolTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,26 @@ final class ConnectionPoolTests: AsyncKitTestCase {
XCTAssertEqual($0 as? ConnectionPoolError, ConnectionPoolError.shutdown)
}
}

func testGracefulShutdownAsync() async throws {
let foo = FooDatabase()
let pool = EventLoopGroupConnectionPool(
source: foo,
maxConnectionsPerEventLoop: 2,
on: self.group
)

try await pool.shutdownAsync()
var errorCaught = false

do {
try await pool.shutdownAsync()
} catch {
errorCaught = true
XCTAssertEqual(error as? ConnectionPoolError, ConnectionPoolError.shutdown)
}
XCTAssertTrue(errorCaught)
}

func testGracefulShutdownWithHeldConnection() throws {
let foo = FooDatabase()
Expand Down

0 comments on commit b239bb0

Please sign in to comment.