diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 00000000..6413432f --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1 @@ +* @fabianfett @gwynne diff --git a/.github/workflows/api-docs.yml b/.github/workflows/api-docs.yml index 80291c6f..dc2e0634 100644 --- a/.github/workflows/api-docs.yml +++ b/.github/workflows/api-docs.yml @@ -11,4 +11,4 @@ jobs: with: package_name: postgres-nio modules: PostgresNIO - pathsToInvalidate: /postgresnio + pathsToInvalidate: /postgresnio/* diff --git a/.github/workflows/projectboard.yml b/.github/workflows/projectboard.yml deleted file mode 100644 index b857f6ee..00000000 --- a/.github/workflows/projectboard.yml +++ /dev/null @@ -1,31 +0,0 @@ -name: issue-to-project-board-workflow -on: - # Trigger when an issue gets labeled or deleted - issues: - types: [reopened, closed, labeled, unlabeled, assigned, unassigned] - -jobs: - setup_matrix_input: - runs-on: ubuntu-latest - - steps: - - id: set-matrix - run: | - output=$(curl ${{ github.event.issue.url }}/labels | jq '.[] | .name') - - echo '======================' - echo 'Process incoming data' - echo '======================' - json=$(echo $output | sed 's/"\s"/","/g') - echo $json - echo "::set-output name=matrix::$(echo $json)" - outputs: - issueTags: ${{ steps.set-matrix.outputs.matrix }} - - Manage_project_issues: - needs: setup_matrix_input - uses: vapor/ci/.github/workflows/issues-to-project-board.yml@main - with: - labelsJson: ${{ needs.setup_matrix_input.outputs.issueTags }} - secrets: - PROJECT_BOARD_AUTOMATION_PAT: "${{ secrets.PROJECT_BOARD_AUTOMATION_PAT }}" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 66516611..8364e8ae 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -9,70 +9,59 @@ on: pull_request: branches: - "*" +env: + LOG_LEVEL: info jobs: linux-unit: strategy: fail-fast: false matrix: - container: - - swift:5.6-focal - - swift:5.7-jammy - - swift:5.8-jammy - - swiftlang/swift:nightly-5.9-jammy + swift-image: + - swift:5.9-jammy + - swift:5.10-noble + - swift:6.0-noble - swiftlang/swift:nightly-main-jammy - include: - - coverage: true - # https://github.com/apple/swift-package-manager/issues/5853 - - container: swift:5.8-jammy - coverage: false - # https://github.com/apple/swift/issues/65064 - - container: swiftlang/swift:nightly-main-jammy - coverage: false - container: ${{ matrix.container }} + container: ${{ matrix.swift-image }} runs-on: ubuntu-latest - env: - LOG_LEVEL: debug steps: - - name: Check out package - uses: actions/checkout@v3 - - name: Run unit tests with code coverage and Thread Sanitizer + - name: Display OS and Swift versions shell: bash run: | - coverage=$( [[ '${{ matrix.coverage }}' == 'true' ]] && echo -n '--enable-code-coverage' || true ) - swift test --filter=^PostgresNIOTests --sanitize=thread ${coverage} - - name: Submit coverage report to Codecov.io - if: ${{ matrix.coverage }} - uses: vapor/swift-codecov-action@v0.2 + [[ -z "${SWIFT_PLATFORM}" ]] && SWIFT_PLATFORM="$(. /etc/os-release && echo "${ID}${VERSION_ID}")" + [[ -z "${SWIFT_VERSION}" ]] && SWIFT_VERSION="$(cat /.swift_tag 2>/dev/null || true)" + printf 'OS: %s\nTag: %s\nVersion:\n' "${SWIFT_PLATFORM}-${RUNNER_ARCH}" "${SWIFT_VERSION}" + swift --version + - name: Check out package + uses: actions/checkout@v4 + - name: Run unit tests with Thread Sanitizer + run: | + swift test --filter='^(PostgresNIOTests|ConnectionPoolModuleTests)' --sanitize=thread --enable-code-coverage + - name: Submit code coverage + uses: vapor/swift-codecov-action@v0.3 with: - cc_flags: 'unittests' - cc_env_vars: 'SWIFT_VERSION,SWIFT_PLATFORM,RUNNER_OS,RUNNER_ARCH' - cc_fail_ci_if_error: true - cc_verbose: true - cc_dry_run: false + codecov_token: ${{ secrets.CODECOV_TOKEN }} linux-integration-and-dependencies: - if: github.event_name == 'pull_request' strategy: fail-fast: false matrix: - dbimage: + postgres-image: + - postgres:17 - postgres:15 - - postgres:13 - - postgres:11 + - postgres:12 include: - - dbimage: postgres:15 - dbauth: scram-sha-256 - - dbimage: postgres:13 - dbauth: md5 - - dbimage: postgres:11 - dbauth: trust + - postgres-image: postgres:17 + postgres-auth: scram-sha-256 + - postgres-image: postgres:15 + postgres-auth: md5 + - postgres-image: postgres:12 + postgres-auth: trust container: - image: swift:5.8-jammy + image: swift:5.10-noble volumes: [ 'pgrunshare:/var/run/postgresql' ] runs-on: ubuntu-latest env: - LOG_LEVEL: debug # Unfortunately, fluent-postgres-driver details leak through here POSTGRES_DB: 'test_database' POSTGRES_DB_A: 'test_database' @@ -87,37 +76,42 @@ jobs: POSTGRES_HOSTNAME_A: 'psql-a' POSTGRES_HOSTNAME_B: 'psql-b' POSTGRES_SOCKET: '/var/run/postgresql/.s.PGSQL.5432' - POSTGRES_HOST_AUTH_METHOD: ${{ matrix.dbauth }} + POSTGRES_HOST_AUTH_METHOD: ${{ matrix.postgres-auth }} services: psql-a: - image: ${{ matrix.dbimage }} + image: ${{ matrix.postgres-image }} volumes: [ 'pgrunshare:/var/run/postgresql' ] env: POSTGRES_USER: 'test_username' POSTGRES_DB: 'test_database' POSTGRES_PASSWORD: 'test_password' - POSTGRES_HOST_AUTH_METHOD: ${{ matrix.dbauth }} - POSTGRES_INITDB_ARGS: --auth-host=${{ matrix.dbauth }} + POSTGRES_HOST_AUTH_METHOD: ${{ matrix.postgres-auth }} + POSTGRES_INITDB_ARGS: --auth-host=${{ matrix.postgres-auth }} psql-b: - image: ${{ matrix.dbimage }} + image: ${{ matrix.postgres-image }} volumes: [ 'pgrunshare:/var/run/postgresql' ] env: POSTGRES_USER: 'test_username' POSTGRES_DB: 'test_database' POSTGRES_PASSWORD: 'test_password' - POSTGRES_HOST_AUTH_METHOD: ${{ matrix.dbauth }} - POSTGRES_INITDB_ARGS: --auth-host=${{ matrix.dbauth }} + POSTGRES_HOST_AUTH_METHOD: ${{ matrix.postgres-auth }} + POSTGRES_INITDB_ARGS: --auth-host=${{ matrix.postgres-auth }} steps: + - name: Display OS and Swift versions + run: | + [[ -z "${SWIFT_PLATFORM}" ]] && SWIFT_PLATFORM="$(. /etc/os-release && echo "${ID}${VERSION_ID}")" + [[ -z "${SWIFT_VERSION}" ]] && SWIFT_VERSION="$(cat /.swift_tag 2>/dev/null || true)" + printf 'OS: %s\nTag: %s\nVersion:\n' "${SWIFT_PLATFORM}-${RUNNER_ARCH}" "${SWIFT_VERSION}" && swift --version - name: Check out package - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: { path: 'postgres-nio' } - name: Run integration tests run: swift test --package-path postgres-nio --filter=^IntegrationTests - name: Check out postgres-kit dependent - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: { repository: 'vapor/postgres-kit', path: 'postgres-kit' } - name: Check out fluent-postgres-driver dependent - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: { repository: 'vapor/fluent-postgres-driver', path: 'fluent-postgres-driver' } - name: Use local package in dependents run: | @@ -129,67 +123,82 @@ jobs: run: swift test --package-path fluent-postgres-driver macos-all: - if: github.event_name == 'pull_request' strategy: fail-fast: false matrix: - dbimage: + postgres-formula: # Only test one version on macOS, let Linux do the rest - - postgresql@14 - dbauth: + - postgresql@16 + postgres-auth: # Only test one auth method on macOS, Linux tests will cover the others - scram-sha-256 - xcode: - - latest-stable - runs-on: macos-12 + xcode-version: + - '~15' + include: + - xcode-version: '~15' + macos-version: 'macos-14' + runs-on: ${{ matrix.macos-version }} env: - LOG_LEVEL: debug POSTGRES_HOSTNAME: 127.0.0.1 POSTGRES_USER: 'test_username' POSTGRES_PASSWORD: 'test_password' POSTGRES_DB: 'postgres' - POSTGRES_HOST_AUTH_METHOD: ${{ matrix.dbauth }} + POSTGRES_AUTH_METHOD: ${{ matrix.postgres-auth }} POSTGRES_SOCKET: '/tmp/.s.PGSQL.5432' + POSTGRES_FORMULA: ${{ matrix.postgres-formula }} steps: - name: Select latest available Xcode uses: maxim-lobanov/setup-xcode@v1 with: - xcode-version: ${{ matrix.xcode }} + xcode-version: ${{ matrix.xcode-version }} - name: Install Postgres, setup DB and auth, and wait for server start run: | - export PATH="$(brew --prefix)/opt/${{ matrix.dbimage }}/bin:$PATH" PGDATA=/tmp/vapor-postgres-test - (brew unlink postgresql || true) && brew install ${{ matrix.dbimage }} && brew link --force ${{ matrix.dbimage }} - initdb --locale=C --auth-host ${{ matrix.dbauth }} -U $POSTGRES_USER --pwfile=<(echo $POSTGRES_PASSWORD) + export PATH="$(brew --prefix)/opt/${POSTGRES_FORMULA}/bin:$PATH" PGDATA=/tmp/vapor-postgres-test + # ** BEGIN ** Work around bug in both Homebrew and GHA + (brew upgrade python@3.11 || true) && (brew link --force --overwrite python@3.11 || true) + (brew upgrade python@3.12 || true) && (brew link --force --overwrite python@3.12 || true) + (brew upgrade || true) + # ** END ** Work around bug in both Homebrew and GHA + brew install --overwrite "${POSTGRES_FORMULA}" + brew link --overwrite --force "${POSTGRES_FORMULA}" + initdb --locale=C --auth-host "${POSTGRES_AUTH_METHOD}" -U "${POSTGRES_USER}" --pwfile=<(echo "${POSTGRES_PASSWORD}") pg_ctl start --wait - timeout-minutes: 2 + timeout-minutes: 15 - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Run all tests run: swift test - + api-breakage: if: github.event_name == 'pull_request' runs-on: ubuntu-latest - container: swift:5.8-jammy - steps: - - name: Checkout - uses: actions/checkout@v3 - with: - fetch-depth: 0 - # https://github.com/actions/checkout/issues/766 - - name: Mark the workspace as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: API breaking changes - run: swift package diagnose-api-breaking-changes origin/main - - test-exports: - name: Test exports - runs-on: ubuntu-latest - container: swift:5.8-jammy + container: swift:noble steps: - - name: Check out package - uses: actions/checkout@v3 + - name: Checkout + uses: actions/checkout@v4 with: fetch-depth: 0 - - name: Build - run: swift build -Xswiftc -DBUILDING_DOCC + # https://github.com/actions/checkout/issues/766 + - name: API breaking changes + run: | + git config --global --add safe.directory "${GITHUB_WORKSPACE}" + swift package diagnose-api-breaking-changes origin/main + +# gh-codeql: +# if: ${{ false }} +# runs-on: ubuntu-latest +# container: swift:noble +# permissions: { actions: write, contents: read, security-events: write } +# steps: +# - name: Check out code +# uses: actions/checkout@v4 +# - name: Mark repo safe in non-fake global config +# run: git config --global --add safe.directory "${GITHUB_WORKSPACE}" +# - name: Initialize CodeQL +# uses: github/codeql-action/init@v3 +# with: +# languages: swift +# - name: Perform build +# run: swift build +# - name: Run CodeQL analyze +# uses: github/codeql-action/analyze@v3 diff --git a/Benchmarks/.gitignore b/Benchmarks/.gitignore new file mode 100644 index 00000000..24e5b0a1 --- /dev/null +++ b/Benchmarks/.gitignore @@ -0,0 +1 @@ +.build diff --git a/Benchmarks/Benchmarks/ConnectionPoolBenchmarks/ConnectionPoolBenchmarks.swift b/Benchmarks/Benchmarks/ConnectionPoolBenchmarks/ConnectionPoolBenchmarks.swift new file mode 100644 index 00000000..98f21f62 --- /dev/null +++ b/Benchmarks/Benchmarks/ConnectionPoolBenchmarks/ConnectionPoolBenchmarks.swift @@ -0,0 +1,51 @@ +import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import Benchmark + +let benchmarks: @Sendable () -> Void = { + Benchmark("Minimal benchmark", configuration: .init(scalingFactor: .kilo)) { benchmark in + let clock = MockClock() + let factory = MockConnectionFactory(autoMaxStreams: 1) + var configuration = ConnectionPoolConfiguration() + configuration.maximumConnectionSoftLimit = 50 + configuration.maximumConnectionHardLimit = 50 + + let pool = ConnectionPool( + configuration: configuration, + idGenerator: ConnectionIDGenerator(), + keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self), + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + await withTaskGroup { taskGroup in + + taskGroup.addTask { + await pool.run() + } + + let sequential = benchmark.scaledIterations.upperBound / configuration.maximumConnectionSoftLimit + + for parallel in 0.. - -[![SSWG Incubating Badge](https://img.shields.io/badge/sswg-incubating-green.svg)][SSWG Incubation] -[![Documentation](http://img.shields.io/badge/read_the-docs-2196f3.svg)][Documentation] -[![Team Chat](https://img.shields.io/discord/431917998102675485.svg)][Team Chat] -[![MIT License](http://img.shields.io/badge/license-MIT-brightgreen.svg)][MIT License] -[![Continuous Integration](https://github.com/vapor/postgres-nio/actions/workflows/test.yml/badge.svg)][Continuous Integration] -[![Swift 5.6](http://img.shields.io/badge/swift-5.6-brightgreen.svg)][Swift 5.6] +

+ + + + PostgresNIO +

+ + Documentation + + + MIT License + + + Continuous Integration + + + Swift 5.8+ + + + SSWG Incubation Level: Graduated + +

🐘 Non-blocking, event-driven Swift client for PostgreSQL built on [SwiftNIO]. Features: - A [`PostgresConnection`] which allows you to connect to, authorize with, query, and retrieve results from a PostgreSQL server +- A [`PostgresClient`] which pools and manages connections - An async/await interface that supports backpressure - Automatic conversions between Swift primitive types and the Postgres wire format -- Integrated with the Swift server ecosystem, including use of [SwiftLog]. +- Integrated with the Swift server ecosystem, including use of [SwiftLog] and [ServiceLifecycle]. - Designed to run efficiently on all supported platforms (tested extensively on Linux and Darwin systems) - Support for `Network.framework` when available (e.g. on Apple platforms) - Supports running on Unix Domain Sockets -PostgresNIO does not provide a `ConnectionPool` as of today, but this is a [feature high on our list](https://github.com/vapor/postgres-nio/issues/256). If you need a `ConnectionPool` today, please have a look at Vapor's [PostgresKit]. - ## API Docs Check out the [PostgresNIO API docs][Documentation] for a @@ -30,13 +43,16 @@ detailed look at all of the classes, structs, protocols, and more. ## Getting started +Interested in an example? We prepared a simple [Birthday example](https://github.com/vapor/postgres-nio/blob/main/Snippets/Birthdays.swift) +in the Snippets folder. + #### Adding the dependency Add `PostgresNIO` as dependency to your `Package.swift`: ```swift dependencies: [ - .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.14.0"), + .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.21.0"), ... ] ``` @@ -50,14 +66,14 @@ Add `PostgresNIO` to the target you want to use it in: ] ``` -#### Creating a connection +#### Creating a client -To create a connection, first create a connection configuration object: +To create a [`PostgresClient`], which pools connections for you, first create a configuration object: ```swift import PostgresNIO -let config = PostgresConnection.Configuration( +let config = PostgresClient.Configuration( host: "localhost", port: 5432, username: "my_username", @@ -67,68 +83,35 @@ let config = PostgresConnection.Configuration( ) ``` -A connection must be created on a SwiftNIO `EventLoop`. In most server use cases, an -`EventLoopGroup` is created at app startup and closed during app shutdown. - +Next you can create you client with it: ```swift -import NIOPosix - -let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) - -// Much later -try await eventLoopGroup.shutdownGracefully() +let client = PostgresClient(configuration: config) ``` -A [`Logger`] is also required. - +Once you have create your client, you must [`run()`] it: ```swift -import Logging - -let logger = Logger(label: "postgres-logger") -``` +await withTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await client.run() // !important + } -Now we can put it together: + // You can use the client while the `client.run()` method is not cancelled. -```swift -import PostgresNIO -import NIOPosix -import Logging - -let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) -let logger = Logger(label: "postgres-logger") - -let config = PostgresConnection.Configuration( - host: "localhost", - port: 5432, - username: "my_username", - password: "my_password", - database: "my_database", - tls: .disable -) - -let connection = try await PostgresConnection.connect( - on: eventLoopGroup.next(), - configuration: config, - id: 1, - logger: logger -) - -// Close your connection once done -try await connection.close() - -// Shutdown the EventLoopGroup, once all connections are closed. -try await eventLoopGroup.shutdownGracefully() + // To shutdown the client, cancel its run method, by cancelling the taskGroup. + taskGroup.cancelAll() +} ``` #### Querying -Once a connection is established, queries can be sent to the server. This is very straightforward: +Once a client is running, queries can be sent to the server. This is straightforward: ```swift -let rows = try await connection.query("SELECT id, username, birthday FROM users", logger: logger) +let rows = try await client.query("SELECT id, username, birthday FROM users") ``` -The query will return a [`PostgresRowSequence`], which is an AsyncSequence of [`PostgresRow`]s. The rows can be iterated one-by-one: +The query will return a [`PostgresRowSequence`], which is an AsyncSequence of [`PostgresRow`]s. +The rows can be iterated one-by-one: ```swift for try await row in rows { @@ -164,7 +147,7 @@ Sending parameterized queries to the database is also supported (in the coolest let id = 1 let username = "fancyuser" let birthday = Date() -try await connection.query(""" +try await client.query(""" INSERT INTO users (id, username, birthday) VALUES (\(id), \(username), \(birthday)) """, logger: logger @@ -180,23 +163,24 @@ Some queries do not receive any rows from the server (most often `INSERT`, `UPDA Please see [SECURITY.md] for details on the security process. [SSWG Incubation]: https://github.com/swift-server/sswg/blob/main/process/incubation.md#graduated-level -[Documentation]: https://swiftpackageindex.com/vapor/postgres-nio/documentation +[Documentation]: https://api.vapor.codes/postgresnio/documentation/postgresnio [Team Chat]: https://discord.gg/vapor [MIT License]: LICENSE [Continuous Integration]: https://github.com/vapor/postgres-nio/actions -[Swift 5.6]: https://swift.org +[Swift 5.8]: https://swift.org [Security.md]: https://github.com/vapor/.github/blob/main/SECURITY.md -[`PostgresConnection`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresconnection/ -[`query(_:logger:)`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresconnection/query(_:logger:file:line:)-9mkfn -[`PostgresQuery`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresquery/ -[`PostgresRow`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresrow/ -[`PostgresRowSequence`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresrowsequence/ -[`PostgresDecodable`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresdecodable/ -[`PostgresEncodable`]: https://swiftpackageindex.com/vapor/postgres-nio/documentation/postgresnio/postgresencodable/ - -[PostgresKit]: https://github.com/vapor/postgres-kit - +[`PostgresConnection`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresconnection +[`PostgresClient`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresclient +[`run()`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresclient/run() +[`query(_:logger:)`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresconnection/query(_:logger:file:line:)-9mkfn +[`PostgresQuery`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresquery +[`PostgresRow`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresrow +[`PostgresRowSequence`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresrowsequence +[`PostgresDecodable`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresdecodable +[`PostgresEncodable`]: https://api.vapor.codes/postgresnio/documentation/postgresnio/postgresencodable [SwiftNIO]: https://github.com/apple/swift-nio +[PostgresKit]: https://github.com/vapor/postgres-kit [SwiftLog]: https://github.com/apple/swift-log +[ServiceLifecycle]: https://github.com/swift-server/swift-service-lifecycle [`Logger`]: https://apple.github.io/swift-log/docs/current/Logging/Structs/Logger.html diff --git a/Snippets/Birthdays.swift b/Snippets/Birthdays.swift new file mode 100644 index 00000000..60516aa1 --- /dev/null +++ b/Snippets/Birthdays.swift @@ -0,0 +1,74 @@ +import PostgresNIO +import Foundation + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +enum Birthday { + static func main() async throws { + // 1. Create a configuration to match server's parameters + let config = PostgresClient.Configuration( + host: "localhost", + port: 5432, + username: "test_username", + password: "test_password", + database: "test_database", + tls: .disable + ) + + // 2. Create a client + let client = PostgresClient(configuration: config) + + // 3. Run the client + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await client.run() // !important + } + + // 4. Create a friends table to store data into + try await client.query(""" + CREATE TABLE IF NOT EXISTS "friends" ( + id SERIAL PRIMARY KEY, + given_name TEXT, + last_name TEXT, + birthday TIMESTAMP WITH TIME ZONE + ) + """ + ) + + // 5. Create a Swift friend representation + struct Friend { + var firstName: String + var lastName: String + var birthday: Date + } + + // 6. Create John Appleseed with special birthday + let dateFormatter = DateFormatter() + dateFormatter.dateFormat = "yyyy-MM-dd" + let johnsBirthday = dateFormatter.date(from: "1960-09-26")! + let friend = Friend(firstName: "Hans", lastName: "Müller", birthday: johnsBirthday) + + // 7. Store friend into the database + try await client.query(""" + INSERT INTO "friends" (given_name, last_name, birthday) + VALUES + (\(friend.firstName), \(friend.lastName), \(friend.birthday)); + """ + ) + + // 8. Query database for the friend we just inserted + let rows = try await client.query(""" + SELECT id, given_name, last_name, birthday FROM "friends" WHERE given_name = \(friend.firstName) + """ + ) + + // 9. Iterate the returned rows, decoding the rows into Swift primitives + for try await (id, firstName, lastName, birthday) in rows.decode((Int, String, String, Date).self) { + print("\(id) | \(firstName) \(lastName), \(birthday)") + } + + // 10. Shutdown the client, by cancelling its run method, through cancelling the taskGroup. + taskGroup.cancelAll() + } + } +} + diff --git a/Snippets/PostgresClient.swift b/Snippets/PostgresClient.swift new file mode 100644 index 00000000..9bfacc28 --- /dev/null +++ b/Snippets/PostgresClient.swift @@ -0,0 +1,40 @@ +import PostgresNIO +import struct Foundation.UUID + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +enum Runner { + static func main() async throws { + +// snippet.configuration +let config = PostgresClient.Configuration( + host: "localhost", + port: 5432, + username: "my_username", + password: "my_password", + database: "my_database", + tls: .disable +) +// snippet.end + +// snippet.makeClient +let client = PostgresClient(configuration: config) +// snippet.end + + } + + static func runAndCancel(client: PostgresClient) async { +// snippet.run +await withTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await client.run() // !important + } + + // You can use the client while the `client.run()` method is not cancelled. + + // To shutdown the client, cancel its run method, by cancelling the taskGroup. + taskGroup.cancelAll() +} +// snippet.end + } +} + diff --git a/Sources/ConnectionPoolModule/ConnectionIDGenerator.swift b/Sources/ConnectionPoolModule/ConnectionIDGenerator.swift new file mode 100644 index 00000000..b428d805 --- /dev/null +++ b/Sources/ConnectionPoolModule/ConnectionIDGenerator.swift @@ -0,0 +1,15 @@ +import Atomics + +public struct ConnectionIDGenerator: ConnectionIDGeneratorProtocol { + static let globalGenerator = ConnectionIDGenerator() + + private let atomic: ManagedAtomic + + public init() { + self.atomic = .init(0) + } + + public func next() -> Int { + return self.atomic.loadThenWrappingIncrement(ordering: .relaxed) + } +} diff --git a/Sources/ConnectionPoolModule/ConnectionPool.swift b/Sources/ConnectionPoolModule/ConnectionPool.swift new file mode 100644 index 00000000..5cdb980d --- /dev/null +++ b/Sources/ConnectionPoolModule/ConnectionPool.swift @@ -0,0 +1,597 @@ + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public struct ConnectionAndMetadata: Sendable { + + public var connection: Connection + + public var maximalStreamsOnConnection: UInt16 + + public init(connection: Connection, maximalStreamsOnConnection: UInt16) { + self.connection = connection + self.maximalStreamsOnConnection = maximalStreamsOnConnection + } +} + +/// A connection that can be pooled in a ``ConnectionPool`` +public protocol PooledConnection: AnyObject, Sendable { + /// The connections identifier type. + associatedtype ID: Hashable & Sendable + + /// The connections identifier. The identifier is passed to + /// the connection factory method and must stay attached to + /// the connection at all times. It must not change during + /// the connections lifetime. + var id: ID { get } + + /// A method to register closures that are invoked when the + /// connection is closed. If the connection closed unexpectedly + /// the closure shall be called with the underlying error. + /// In most NIO clients this can be easily implemented by + /// attaching to the `channel.closeFuture`: + /// ``` + /// func onClose( + /// _ closure: @escaping @Sendable ((any Error)?) -> () + /// ) { + /// channel.closeFuture.whenComplete { _ in + /// closure(previousError) + /// } + /// } + /// ``` + func onClose(_ closure: @escaping @Sendable ((any Error)?) -> ()) + + /// Close the running connection. Once the close has completed + /// closures that were registered in `onClose` must be + /// invoked. + func close() +} + +/// A connection id generator. Its returned connection IDs will +/// be used when creating new ``PooledConnection``s +public protocol ConnectionIDGeneratorProtocol: Sendable { + /// The connections identifier type. + associatedtype ID: Hashable & Sendable + + /// The next connection ID that shall be used. + func next() -> ID +} + +/// A keep alive behavior for connections maintained by the pool +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public protocol ConnectionKeepAliveBehavior: Sendable { + /// the connection type + associatedtype Connection: PooledConnection + + /// The time after which a keep-alive shall + /// be triggered. + /// If nil is returned, keep-alive is deactivated + var keepAliveFrequency: Duration? { get } + + /// This method is invoked when the keep-alive shall be + /// run. + func runKeepAlive(for connection: Connection) async throws +} + +/// A request to get a connection from the `ConnectionPool` +public protocol ConnectionRequestProtocol: Sendable { + /// A connection lease request ID type. + associatedtype ID: Hashable & Sendable + /// The leased connection type + associatedtype Connection: PooledConnection + + /// A connection lease request ID. This ID must be generated + /// by users of the `ConnectionPool` outside the + /// `ConnectionPool`. It is not generated inside the pool like + /// the `ConnectionID`s. The lease request ID must be unique + /// and must not change, if your implementing type is a + /// reference type. + var id: ID { get } + + /// A function that is called with a connection or a + /// `PoolError`. + func complete(with: Result) +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public struct ConnectionPoolConfiguration: Sendable { + /// The minimum number of connections to preserve in the pool. + /// + /// If the pool is mostly idle and the remote servers closes + /// idle connections, + /// the `ConnectionPool` will initiate new outbound + /// connections proactively to avoid the number of available + /// connections dropping below this number. + public var minimumConnectionCount: Int + + /// Between the `minimumConnectionCount` and + /// `maximumConnectionSoftLimit` the connection pool creates + /// _preserved_ connections. Preserved connections are closed + /// if they have been idle for ``idleTimeout``. + public var maximumConnectionSoftLimit: Int + + /// The maximum number of connections for this pool, that can + /// exist at any point in time. The pool can create _overflow_ + /// connections, if all connections are leased, and the + /// `maximumConnectionHardLimit` > `maximumConnectionSoftLimit ` + /// Overflow connections are closed immediately as soon as they + /// become idle. + public var maximumConnectionHardLimit: Int + + /// The time that a _preserved_ idle connection stays in the + /// pool before it is closed. + public var idleTimeout: Duration + + /// initializer + public init() { + self.minimumConnectionCount = 0 + self.maximumConnectionSoftLimit = 16 + self.maximumConnectionHardLimit = 16 + self.idleTimeout = .seconds(60) + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public final class ConnectionPool< + Connection: PooledConnection, + ConnectionID: Hashable & Sendable, + ConnectionIDGenerator: ConnectionIDGeneratorProtocol, + Request: ConnectionRequestProtocol, + RequestID: Hashable & Sendable, + KeepAliveBehavior: ConnectionKeepAliveBehavior, + ObservabilityDelegate: ConnectionPoolObservabilityDelegate, + Clock: _Concurrency.Clock +>: Sendable where + Connection.ID == ConnectionID, + ConnectionIDGenerator.ID == ConnectionID, + Request.Connection == Connection, + Request.ID == RequestID, + KeepAliveBehavior.Connection == Connection, + ObservabilityDelegate.ConnectionID == ConnectionID, + Clock.Duration == Duration +{ + public typealias ConnectionFactory = @Sendable (ConnectionID, ConnectionPool) async throws -> ConnectionAndMetadata + + @usableFromInline + typealias StateMachine = PoolStateMachine> + + @usableFromInline + let factory: ConnectionFactory + + @usableFromInline + let keepAliveBehavior: KeepAliveBehavior + + @usableFromInline + let observabilityDelegate: ObservabilityDelegate + + @usableFromInline + let clock: Clock + + @usableFromInline + let configuration: ConnectionPoolConfiguration + + @usableFromInline + struct State: Sendable { + @usableFromInline + var stateMachine: StateMachine + @usableFromInline + var lastConnectError: (any Error)? + } + + @usableFromInline let stateBox: NIOLockedValueBox + + private let requestIDGenerator = _ConnectionPoolModule.ConnectionIDGenerator() + + @usableFromInline + let eventStream: AsyncStream + + @usableFromInline + let eventContinuation: AsyncStream.Continuation + + public init( + configuration: ConnectionPoolConfiguration, + idGenerator: ConnectionIDGenerator, + requestType: Request.Type, + keepAliveBehavior: KeepAliveBehavior, + observabilityDelegate: ObservabilityDelegate, + clock: Clock, + connectionFactory: @escaping ConnectionFactory + ) { + self.clock = clock + self.factory = connectionFactory + self.keepAliveBehavior = keepAliveBehavior + self.observabilityDelegate = observabilityDelegate + self.configuration = configuration + var stateMachine = StateMachine( + configuration: .init(configuration, keepAliveBehavior: keepAliveBehavior), + generator: idGenerator, + timerCancellationTokenType: CheckedContinuation.self + ) + + let (stream, continuation) = AsyncStream.makeStream(of: NewPoolActions.self) + self.eventStream = stream + self.eventContinuation = continuation + + let connectionRequests = stateMachine.refillConnections() + + self.stateBox = NIOLockedValueBox(.init(stateMachine: stateMachine)) + + for request in connectionRequests { + self.eventContinuation.yield(.makeConnection(request)) + } + } + + @inlinable + public func releaseConnection(_ connection: Connection, streams: UInt16 = 1) { + self.modifyStateAndRunActions { state in + state.stateMachine.releaseConnection(connection, streams: streams) + } + } + + @inlinable + public func leaseConnection(_ request: Request) { + self.modifyStateAndRunActions { state in + state.stateMachine.leaseConnection(request) + } + } + + @inlinable + public func leaseConnections(_ requests: some Collection) { + let actions = self.stateBox.withLockedValue { state in + var actions = [StateMachine.Action]() + actions.reserveCapacity(requests.count) + + for request in requests { + let stateMachineAction = state.stateMachine.leaseConnection(request) + actions.append(stateMachineAction) + } + + return actions + } + + for action in actions { + self.runRequestAction(action.request) + self.runConnectionAction(action.connection) + } + } + + public func cancelLeaseConnection(_ requestID: RequestID) { + self.modifyStateAndRunActions { state in + state.stateMachine.cancelRequest(id: requestID) + } + } + + /// Mark a connection as going away. Connection implementors have to call this method if the connection + /// has received a close intent from the server. For example: an HTTP/2 GOWAY frame. + public func connectionWillClose(_ connection: Connection) { + + } + + public func connectionReceivedNewMaxStreamSetting(_ connection: Connection, newMaxStreamSetting maxStreams: UInt16) { + self.modifyStateAndRunActions { state in + state.stateMachine.connectionReceivedNewMaxStreamSetting(connection.id, newMaxStreamSetting: maxStreams) + } + } + + public func run() async { + await withTaskCancellationHandler { + #if os(Linux) || compiler(>=5.9) + if #available(macOS 14.0, iOS 17.0, tvOS 17.0, watchOS 10.0, *) { + return await withDiscardingTaskGroup() { taskGroup in + await self.run(in: &taskGroup) + } + } + #endif + return await withTaskGroup(of: Void.self) { taskGroup in + await self.run(in: &taskGroup) + } + } onCancel: { + let actions = self.stateBox.withLockedValue { state in + state.stateMachine.triggerForceShutdown() + } + + self.runStateMachineActions(actions) + } + } + + // MARK: - Private Methods - + + @inlinable + func connectionDidClose(_ connection: Connection, error: (any Error)?) { + self.observabilityDelegate.connectionClosed(id: connection.id, error: error) + + self.modifyStateAndRunActions { state in + state.stateMachine.connectionClosed(connection) + } + } + + // MARK: Events + + @usableFromInline + enum NewPoolActions: Sendable { + case makeConnection(StateMachine.ConnectionRequest) + case runKeepAlive(Connection) + + case scheduleTimer(StateMachine.Timer) + } + + #if os(Linux) || compiler(>=5.9) + @available(macOS 14.0, iOS 17.0, tvOS 17.0, watchOS 10.0, *) + private func run(in taskGroup: inout DiscardingTaskGroup) async { + for await event in self.eventStream { + self.runEvent(event, in: &taskGroup) + } + } + #endif + + private func run(in taskGroup: inout TaskGroup) async { + var running = 0 + for await event in self.eventStream { + running += 1 + self.runEvent(event, in: &taskGroup) + + if running == 100 { + _ = await taskGroup.next() + running -= 1 + } + } + } + + private func runEvent(_ event: NewPoolActions, in taskGroup: inout some TaskGroupProtocol) { + switch event { + case .makeConnection(let request): + self.makeConnection(for: request, in: &taskGroup) + + case .runKeepAlive(let connection): + self.runKeepAlive(connection, in: &taskGroup) + + case .scheduleTimer(let timer): + self.runTimer(timer, in: &taskGroup) + } + } + + // MARK: Run actions + + @inlinable + /*private*/ func modifyStateAndRunActions(_ closure: (inout State) -> StateMachine.Action) { + let actions = self.stateBox.withLockedValue { state -> StateMachine.Action in + closure(&state) + } + self.runStateMachineActions(actions) + } + + @inlinable + /*private*/ func runStateMachineActions(_ actions: StateMachine.Action) { + self.runConnectionAction(actions.connection) + self.runRequestAction(actions.request) + } + + @inlinable + /*private*/ func runConnectionAction(_ action: StateMachine.ConnectionAction) { + switch action { + case .makeConnection(let request, let timers): + self.cancelTimers(timers) + self.eventContinuation.yield(.makeConnection(request)) + + case .runKeepAlive(let connection, let cancelContinuation): + cancelContinuation?.resume(returning: ()) + self.eventContinuation.yield(.runKeepAlive(connection)) + + case .scheduleTimers(let timers): + for timer in timers { + self.eventContinuation.yield(.scheduleTimer(timer)) + } + + case .cancelTimers(let timers): + self.cancelTimers(timers) + + case .closeConnection(let connection, let timers): + self.closeConnection(connection) + self.cancelTimers(timers) + + case .shutdown(let cleanup): + for connection in cleanup.connections { + self.closeConnection(connection) + } + self.cancelTimers(cleanup.timersToCancel) + + case .none: + break + } + } + + @inlinable + /*private*/ func runRequestAction(_ action: StateMachine.RequestAction) { + switch action { + case .leaseConnection(let requests, let connection): + for request in requests { + request.complete(with: .success(connection)) + } + + case .failRequest(let request, let error): + request.complete(with: .failure(error)) + + case .failRequests(let requests, let error): + for request in requests { request.complete(with: .failure(error)) } + + case .none: + break + } + } + + @inlinable + /*private*/ func makeConnection(for request: StateMachine.ConnectionRequest, in taskGroup: inout some TaskGroupProtocol) { + taskGroup.addTask_ { + self.observabilityDelegate.startedConnecting(id: request.connectionID) + + do { + let bundle = try await self.factory(request.connectionID, self) + self.connectionEstablished(bundle) + + // after the connection has been established, we keep the task open. This ensures + // that the pools run method can not be exited before all connections have been + // closed. + await withCheckedContinuation { (continuation: CheckedContinuation) in + bundle.connection.onClose { + self.connectionDidClose(bundle.connection, error: $0) + continuation.resume() + } + } + } catch { + self.connectionEstablishFailed(error, for: request) + } + } + } + + @inlinable + /*private*/ func connectionEstablished(_ connectionBundle: ConnectionAndMetadata) { + self.observabilityDelegate.connectSucceeded(id: connectionBundle.connection.id, streamCapacity: connectionBundle.maximalStreamsOnConnection) + + self.modifyStateAndRunActions { state in + state.lastConnectError = nil + return state.stateMachine.connectionEstablished( + connectionBundle.connection, + maxStreams: connectionBundle.maximalStreamsOnConnection + ) + } + } + + @inlinable + /*private*/ func connectionEstablishFailed(_ error: Error, for request: StateMachine.ConnectionRequest) { + self.observabilityDelegate.connectFailed(id: request.connectionID, error: error) + + self.modifyStateAndRunActions { state in + state.lastConnectError = error + return state.stateMachine.connectionEstablishFailed(error, for: request) + } + } + + @inlinable + /*private*/ func runKeepAlive(_ connection: Connection, in taskGroup: inout some TaskGroupProtocol) { + self.observabilityDelegate.keepAliveTriggered(id: connection.id) + + taskGroup.addTask_ { + do { + try await self.keepAliveBehavior.runKeepAlive(for: connection) + + self.observabilityDelegate.keepAliveSucceeded(id: connection.id) + + self.modifyStateAndRunActions { state in + state.stateMachine.connectionKeepAliveDone(connection) + } + } catch { + self.observabilityDelegate.keepAliveFailed(id: connection.id, error: error) + + self.modifyStateAndRunActions { state in + state.stateMachine.connectionKeepAliveFailed(connection.id) + } + } + } + } + + @inlinable + /*private*/ func closeConnection(_ connection: Connection) { + self.observabilityDelegate.connectionClosing(id: connection.id) + + connection.close() + } + + @usableFromInline + enum TimerRunResult: Sendable { + case timerTriggered + case timerCancelled + case cancellationContinuationFinished + } + + @inlinable + /*private*/ func runTimer(_ timer: StateMachine.Timer, in poolGroup: inout some TaskGroupProtocol) { + poolGroup.addTask_ { () async -> () in + await withTaskGroup(of: TimerRunResult.self, returning: Void.self) { taskGroup in + taskGroup.addTask { + do { + #if os(Linux) || compiler(>=5.9) + try await self.clock.sleep(for: timer.duration) + #else + try await self.clock.sleep(until: self.clock.now.advanced(by: timer.duration), tolerance: nil) + #endif + return .timerTriggered + } catch { + return .timerCancelled + } + } + + taskGroup.addTask { + await withCheckedContinuation { (continuation: CheckedContinuation) in + let continuation = self.stateBox.withLockedValue { state in + state.stateMachine.timerScheduled(timer, cancelContinuation: continuation) + } + + continuation?.resume(returning: ()) + } + + return .cancellationContinuationFinished + } + + switch await taskGroup.next()! { + case .cancellationContinuationFinished: + taskGroup.cancelAll() + + case .timerTriggered: + let action = self.stateBox.withLockedValue { state in + state.stateMachine.timerTriggered(timer) + } + + self.runStateMachineActions(action) + + case .timerCancelled: + // the only way to reach this, is if the state machine decided to cancel the + // timer. therefore we don't need to report it back! + break + } + + return + } + } + } + + @inlinable + /*private*/ func cancelTimers(_ cancellationTokens: some Sequence>) { + for token in cancellationTokens { + token.resume() + } + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolConfiguration { + init(_ configuration: ConnectionPoolConfiguration, keepAliveBehavior: KeepAliveBehavior) { + self.minimumConnectionCount = configuration.minimumConnectionCount + self.maximumConnectionSoftLimit = configuration.maximumConnectionSoftLimit + self.maximumConnectionHardLimit = configuration.maximumConnectionHardLimit + self.keepAliveDuration = keepAliveBehavior.keepAliveFrequency + self.idleTimeoutDuration = configuration.idleTimeout + } +} + +@usableFromInline +protocol TaskGroupProtocol { + // We need to call this `addTask_` because some Swift versions define this + // under exactly this name and others have different attributes. So let's pick + // a name that doesn't clash anywhere and implement it using the standard `addTask`. + mutating func addTask_(operation: @escaping @Sendable () async -> Void) +} + +#if os(Linux) || swift(>=5.9) +@available(macOS 14.0, iOS 17.0, tvOS 17.0, watchOS 10.0, *) +extension DiscardingTaskGroup: TaskGroupProtocol { + @inlinable + mutating func addTask_(operation: @escaping @Sendable () async -> Void) { + self.addTask(priority: nil, operation: operation) + } +} +#endif + +extension TaskGroup: TaskGroupProtocol { + @inlinable + mutating func addTask_(operation: @escaping @Sendable () async -> Void) { + self.addTask(priority: nil, operation: operation) + } +} diff --git a/Sources/ConnectionPoolModule/ConnectionPoolError.swift b/Sources/ConnectionPoolModule/ConnectionPoolError.swift new file mode 100644 index 00000000..1f1e1d2c --- /dev/null +++ b/Sources/ConnectionPoolModule/ConnectionPoolError.swift @@ -0,0 +1,16 @@ + +public struct ConnectionPoolError: Error, Hashable { + enum Base: Error, Hashable { + case requestCancelled + case poolShutdown + } + + private let base: Base + + init(_ base: Base) { self.base = base } + + /// The connection requests got cancelled + public static let requestCancelled = ConnectionPoolError(.requestCancelled) + /// The connection requests can't be fulfilled as the pool has already been shutdown + public static let poolShutdown = ConnectionPoolError(.poolShutdown) +} diff --git a/Sources/ConnectionPoolModule/ConnectionPoolObservabilityDelegate.swift b/Sources/ConnectionPoolModule/ConnectionPoolObservabilityDelegate.swift new file mode 100644 index 00000000..fc1e300c --- /dev/null +++ b/Sources/ConnectionPoolModule/ConnectionPoolObservabilityDelegate.swift @@ -0,0 +1,62 @@ + +public protocol ConnectionPoolObservabilityDelegate: Sendable { + associatedtype ConnectionID: Hashable & Sendable + + /// The connection with the given ID has started trying to establish a connection. The outcome + /// of the connection will be reported as either ``connectSucceeded(id:streamCapacity:)`` or + /// ``connectFailed(id:error:)``. + func startedConnecting(id: ConnectionID) + + /// A connection attempt failed with the given error. After some period of + /// time ``startedConnecting(id:)`` may be called again. + func connectFailed(id: ConnectionID, error: Error) + + /// A connection was established on the connection with the given ID. `streamCapacity` streams are + /// available to use on the connection. The maximum number of available streams may change over + /// time and is reported via ````. The + func connectSucceeded(id: ConnectionID, streamCapacity: UInt16) + + /// The utlization of the connection changed; a stream may have been used, returned or the + /// maximum number of concurrent streams available on the connection changed. + func connectionUtilizationChanged(id:ConnectionID, streamsUsed: UInt16, streamCapacity: UInt16) + + func keepAliveTriggered(id: ConnectionID) + + func keepAliveSucceeded(id: ConnectionID) + + func keepAliveFailed(id: ConnectionID, error: Error) + + /// The remote peer is quiescing the connection: no new streams will be created on it. The + /// connection will eventually be closed and removed from the pool. + func connectionClosing(id: ConnectionID) + + /// The connection was closed. The connection may be established again in the future (notified + /// via ``startedConnecting(id:)``). + func connectionClosed(id: ConnectionID, error: Error?) + + func requestQueueDepthChanged(_ newDepth: Int) +} + +public struct NoOpConnectionPoolMetrics: ConnectionPoolObservabilityDelegate { + public init(connectionIDType: ConnectionID.Type) {} + + public func startedConnecting(id: ConnectionID) {} + + public func connectFailed(id: ConnectionID, error: Error) {} + + public func connectSucceeded(id: ConnectionID, streamCapacity: UInt16) {} + + public func connectionUtilizationChanged(id: ConnectionID, streamsUsed: UInt16, streamCapacity: UInt16) {} + + public func keepAliveTriggered(id: ConnectionID) {} + + public func keepAliveSucceeded(id: ConnectionID) {} + + public func keepAliveFailed(id: ConnectionID, error: Error) {} + + public func connectionClosing(id: ConnectionID) {} + + public func connectionClosed(id: ConnectionID, error: Error?) {} + + public func requestQueueDepthChanged(_ newDepth: Int) {} +} diff --git a/Sources/ConnectionPoolModule/ConnectionRequest.swift b/Sources/ConnectionPoolModule/ConnectionRequest.swift new file mode 100644 index 00000000..1d1c55da --- /dev/null +++ b/Sources/ConnectionPoolModule/ConnectionRequest.swift @@ -0,0 +1,78 @@ + +public struct ConnectionRequest: ConnectionRequestProtocol { + public typealias ID = Int + + public var id: ID + + @usableFromInline + private(set) var continuation: CheckedContinuation + + @inlinable + init( + id: Int, + continuation: CheckedContinuation + ) { + self.id = id + self.continuation = continuation + } + + public func complete(with result: Result) { + self.continuation.resume(with: result) + } +} + +@usableFromInline +let requestIDGenerator = _ConnectionPoolModule.ConnectionIDGenerator() + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension ConnectionPool where Request == ConnectionRequest { + public convenience init( + configuration: ConnectionPoolConfiguration, + idGenerator: ConnectionIDGenerator = _ConnectionPoolModule.ConnectionIDGenerator(), + keepAliveBehavior: KeepAliveBehavior, + observabilityDelegate: ObservabilityDelegate, + clock: Clock = ContinuousClock(), + connectionFactory: @escaping ConnectionFactory + ) { + self.init( + configuration: configuration, + idGenerator: idGenerator, + requestType: ConnectionRequest.self, + keepAliveBehavior: keepAliveBehavior, + observabilityDelegate: observabilityDelegate, + clock: clock, + connectionFactory: connectionFactory + ) + } + + @inlinable + public func leaseConnection() async throws -> Connection { + let requestID = requestIDGenerator.next() + + let connection = try await withTaskCancellationHandler { + if Task.isCancelled { + throw CancellationError() + } + + return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + let request = Request( + id: requestID, + continuation: continuation + ) + + self.leaseConnection(request) + } + } onCancel: { + self.cancelLeaseConnection(requestID) + } + + return connection + } + + @inlinable + public func withConnection(_ closure: (Connection) async throws -> Result) async throws -> Result { + let connection = try await self.leaseConnection() + defer { self.releaseConnection(connection) } + return try await closure(connection) + } +} diff --git a/Sources/ConnectionPoolModule/Max2Sequence.swift b/Sources/ConnectionPoolModule/Max2Sequence.swift new file mode 100644 index 00000000..9b7d972b --- /dev/null +++ b/Sources/ConnectionPoolModule/Max2Sequence.swift @@ -0,0 +1,104 @@ +// A `Sequence` that can contain at most two elements. However it does not heap allocate. +@usableFromInline +struct Max2Sequence: Sequence { + @usableFromInline + private(set) var first: Element? + @usableFromInline + private(set) var second: Element? + + @inlinable + var count: Int { + if self.first == nil { return 0 } + if self.second == nil { return 1 } + return 2 + } + + @inlinable + var isEmpty: Bool { + self.first == nil + } + + @inlinable + init(_ first: Element?, _ second: Element? = nil) { + if let first = first { + self.first = first + self.second = second + } else { + self.first = second + self.second = nil + } + } + + @inlinable + init() { + self.first = nil + self.second = nil + } + + @inlinable + func makeIterator() -> Iterator { + Iterator(first: self.first, second: self.second) + } + + @usableFromInline + struct Iterator: IteratorProtocol { + @usableFromInline + let first: Element? + @usableFromInline + let second: Element? + + @usableFromInline + private(set) var index: UInt8 = 0 + + @inlinable + init(first: Element?, second: Element?) { + self.first = first + self.second = second + self.index = 0 + } + + @inlinable + mutating func next() -> Element? { + switch self.index { + case 0: + self.index += 1 + return self.first + case 1: + self.index += 1 + return self.second + default: + return nil + } + } + } + + @inlinable + mutating func append(_ element: Element) { + precondition(self.second == nil) + if self.first == nil { + self.first = element + } else if self.second == nil { + self.second = element + } else { + fatalError("Max2Sequence can only hold two Elements.") + } + } + + @inlinable + func map(_ transform: (Element) throws -> (NewElement)) rethrows -> Max2Sequence { + try Max2Sequence(self.first.flatMap(transform), self.second.flatMap(transform)) + } +} + +extension Max2Sequence: ExpressibleByArrayLiteral { + @inlinable + init(arrayLiteral elements: Element...) { + precondition(elements.count <= 2) + var iterator = elements.makeIterator() + self.init(iterator.next(), iterator.next()) + } +} + +extension Max2Sequence: Equatable where Element: Equatable {} +extension Max2Sequence: Hashable where Element: Hashable {} +extension Max2Sequence: Sendable where Element: Sendable {} diff --git a/Sources/ConnectionPoolModule/NIOLock.swift b/Sources/ConnectionPoolModule/NIOLock.swift new file mode 100644 index 00000000..b6cd7164 --- /dev/null +++ b/Sources/ConnectionPoolModule/NIOLock.swift @@ -0,0 +1,279 @@ +// Implementation vendored from SwiftNIO: +// https://github.com/apple/swift-nio + +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftNIO open source project +// +// Copyright (c) 2017-2022 Apple Inc. and the SwiftNIO project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftNIO project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +#if canImport(Darwin) +import Darwin +#elseif os(Windows) +import ucrt +import WinSDK +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#elseif canImport(Bionic) +import Bionic +#elseif canImport(WASILibc) +import WASILibc +#if canImport(wasi_pthread) +import wasi_pthread +#endif +#else +#error("The concurrency NIOLock module was unable to identify your C library.") +#endif + +#if os(Windows) +@usableFromInline +typealias LockPrimitive = SRWLOCK +#else +@usableFromInline +typealias LockPrimitive = pthread_mutex_t +#endif + +@usableFromInline +enum LockOperations {} + +extension LockOperations { + @inlinable + static func create(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + #if os(Windows) + InitializeSRWLock(mutex) + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + var attr = pthread_mutexattr_t() + pthread_mutexattr_init(&attr) + debugOnly { + pthread_mutexattr_settype(&attr, .init(PTHREAD_MUTEX_ERRORCHECK)) + } + + let err = pthread_mutex_init(mutex, &attr) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + #endif + } + + @inlinable + static func destroy(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + #if os(Windows) + // SRWLOCK does not need to be free'd + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + let err = pthread_mutex_destroy(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + #endif + } + + @inlinable + static func lock(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + #if os(Windows) + AcquireSRWLockExclusive(mutex) + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + let err = pthread_mutex_lock(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + #endif + } + + @inlinable + static func unlock(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + #if os(Windows) + ReleaseSRWLockExclusive(mutex) + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + let err = pthread_mutex_unlock(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + #endif + } +} + +// Tail allocate both the mutex and a generic value using ManagedBuffer. +// Both the header pointer and the elements pointer are stable for +// the class's entire lifetime. +// +// However, for safety reasons, we elect to place the lock in the "elements" +// section of the buffer instead of the head. The reasoning here is subtle, +// so buckle in. +// +// _As a practical matter_, the implementation of ManagedBuffer ensures that +// the pointer to the header is stable across the lifetime of the class, and so +// each time you call `withUnsafeMutablePointers` or `withUnsafeMutablePointerToHeader` +// the value of the header pointer will be the same. This is because ManagedBuffer uses +// `Builtin.addressOf` to load the value of the header, and that does ~magic~ to ensure +// that it does not invoke any weird Swift accessors that might copy the value. +// +// _However_, the header is also available via the `.header` field on the ManagedBuffer. +// This presents a problem! The reason there's an issue is that `Builtin.addressOf` and friends +// do not interact with Swift's exclusivity model. That is, the various `with` functions do not +// conceptually trigger a mutating access to `.header`. For elements this isn't a concern because +// there's literally no other way to perform the access, but for `.header` it's entirely possible +// to accidentally recursively read it. +// +// Our implementation is free from these issues, so we don't _really_ need to worry about it. +// However, out of an abundance of caution, we store the Value in the header, and the LockPrimitive +// in the trailing elements. We still don't use `.header`, but it's better to be safe than sorry, +// and future maintainers will be happier that we were cautious. +// +// See also: https://github.com/apple/swift/pull/40000 +@usableFromInline +final class LockStorage: ManagedBuffer { + + @inlinable + static func create(value: Value) -> Self { + let buffer = Self.create(minimumCapacity: 1) { _ in + value + } + // Intentionally using a force cast here to avoid a miss compiliation in 5.10. + // This is as fast as an unsafeDownCast since ManagedBuffer is inlined and the optimizer + // can eliminate the upcast/downcast pair + let storage = buffer as! Self + + storage.withUnsafeMutablePointers { _, lockPtr in + LockOperations.create(lockPtr) + } + + return storage + } + + @inlinable + func lock() { + self.withUnsafeMutablePointerToElements { lockPtr in + LockOperations.lock(lockPtr) + } + } + + @inlinable + func unlock() { + self.withUnsafeMutablePointerToElements { lockPtr in + LockOperations.unlock(lockPtr) + } + } + + @inlinable + deinit { + self.withUnsafeMutablePointerToElements { lockPtr in + LockOperations.destroy(lockPtr) + } + } + + @inlinable + func withLockPrimitive(_ body: (UnsafeMutablePointer) throws -> T) rethrows -> T { + try self.withUnsafeMutablePointerToElements { lockPtr in + try body(lockPtr) + } + } + + @inlinable + func withLockedValue(_ mutate: (inout Value) throws -> T) rethrows -> T { + try self.withUnsafeMutablePointers { valuePtr, lockPtr in + LockOperations.lock(lockPtr) + defer { LockOperations.unlock(lockPtr) } + return try mutate(&valuePtr.pointee) + } + } +} + +/// A threading lock based on `libpthread` instead of `libdispatch`. +/// +/// - Note: ``NIOLock`` has reference semantics. +/// +/// This object provides a lock on top of a single `pthread_mutex_t`. This kind +/// of lock is safe to use with `libpthread`-based threading models, such as the +/// one used by NIO. On Windows, the lock is based on the substantially similar +/// `SRWLOCK` type. +struct NIOLock { + @usableFromInline + internal let _storage: LockStorage + + /// Create a new lock. + @inlinable + init() { + self._storage = .create(value: ()) + } + + /// Acquire the lock. + /// + /// Whenever possible, consider using `withLock` instead of this method and + /// `unlock`, to simplify lock handling. + @inlinable + func lock() { + self._storage.lock() + } + + /// Release the lock. + /// + /// Whenever possible, consider using `withLock` instead of this method and + /// `lock`, to simplify lock handling. + @inlinable + func unlock() { + self._storage.unlock() + } + + @inlinable + internal func withLockPrimitive(_ body: (UnsafeMutablePointer) throws -> T) rethrows -> T { + try self._storage.withLockPrimitive(body) + } +} + +extension NIOLock { + /// Acquire the lock for the duration of the given block. + /// + /// This convenience method should be preferred to `lock` and `unlock` in + /// most situations, as it ensures that the lock will be released regardless + /// of how `body` exits. + /// + /// - Parameter body: The block to execute while holding the lock. + /// - Returns: The value returned by the block. + @inlinable + func withLock(_ body: () throws -> T) rethrows -> T { + self.lock() + defer { + self.unlock() + } + return try body() + } + + @inlinable + func withLockVoid(_ body: () throws -> Void) rethrows { + try self.withLock(body) + } +} + +extension NIOLock: @unchecked Sendable {} + +extension UnsafeMutablePointer { + @inlinable + func assertValidAlignment() { + assert(UInt(bitPattern: self) % UInt(MemoryLayout.alignment) == 0) + } +} + +/// A utility function that runs the body code only in debug builds, without +/// emitting compiler warnings. +/// +/// This is currently the only way to do this in Swift: see +/// https://forums.swift.org/t/support-debug-only-code/11037 for a discussion. +@inlinable +internal func debugOnly(_ body: () -> Void) { + assert( + { + body() + return true + }() + ) +} diff --git a/Sources/ConnectionPoolModule/NIOLockedValueBox.swift b/Sources/ConnectionPoolModule/NIOLockedValueBox.swift new file mode 100644 index 00000000..c9cd89e0 --- /dev/null +++ b/Sources/ConnectionPoolModule/NIOLockedValueBox.swift @@ -0,0 +1,86 @@ +// Implementation vendored from SwiftNIO: +// https://github.com/apple/swift-nio + +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftNIO open source project +// +// Copyright (c) 2022 Apple Inc. and the SwiftNIO project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftNIO project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +/// Provides locked access to `Value`. +/// +/// - Note: ``NIOLockedValueBox`` has reference semantics and holds the `Value` +/// alongside a lock behind a reference. +/// +/// This is no different than creating a ``Lock`` and protecting all +/// accesses to a value using the lock. But it's easy to forget to actually +/// acquire/release the lock in the correct place. ``NIOLockedValueBox`` makes +/// that much easier. +@usableFromInline +struct NIOLockedValueBox { + + @usableFromInline + internal let _storage: LockStorage + + /// Initialize the `Value`. + @inlinable + init(_ value: Value) { + self._storage = .create(value: value) + } + + /// Access the `Value`, allowing mutation of it. + @inlinable + func withLockedValue(_ mutate: (inout Value) throws -> T) rethrows -> T { + try self._storage.withLockedValue(mutate) + } + + /// Provides an unsafe view over the lock and its value. + /// + /// This can be beneficial when you require fine grained control over the lock in some + /// situations but don't want lose the benefits of ``withLockedValue(_:)`` in others by + /// switching to ``NIOLock``. + var unsafe: Unsafe { + Unsafe(_storage: self._storage) + } + + /// Provides an unsafe view over the lock and its value. + struct Unsafe { + @usableFromInline + let _storage: LockStorage + + /// Manually acquire the lock. + @inlinable + func lock() { + self._storage.lock() + } + + /// Manually release the lock. + @inlinable + func unlock() { + self._storage.unlock() + } + + /// Mutate the value, assuming the lock has been acquired manually. + /// + /// - Parameter mutate: A closure with scoped access to the value. + /// - Returns: The result of the `mutate` closure. + @inlinable + func withValueAssumingLockIsAcquired( + _ mutate: (_ value: inout Value) throws -> Result + ) rethrows -> Result { + try self._storage.withUnsafeMutablePointerToHeader { value in + try mutate(&value.pointee) + } + } + } +} + +extension NIOLockedValueBox: @unchecked Sendable where Value: Sendable {} diff --git a/Sources/ConnectionPoolModule/NoKeepAliveBehavior.swift b/Sources/ConnectionPoolModule/NoKeepAliveBehavior.swift new file mode 100644 index 00000000..0a7b2dee --- /dev/null +++ b/Sources/ConnectionPoolModule/NoKeepAliveBehavior.swift @@ -0,0 +1,8 @@ +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public struct NoOpKeepAliveBehavior: ConnectionKeepAliveBehavior { + public var keepAliveFrequency: Duration? { nil } + + public func runKeepAlive(for connection: Connection) async throws {} + + public init(connectionType: Connection.Type) {} +} diff --git a/Sources/ConnectionPoolModule/PoolStateMachine+ConnectionGroup.swift b/Sources/ConnectionPoolModule/PoolStateMachine+ConnectionGroup.swift new file mode 100644 index 00000000..a8e97ffd --- /dev/null +++ b/Sources/ConnectionPoolModule/PoolStateMachine+ConnectionGroup.swift @@ -0,0 +1,733 @@ +import Atomics + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine { + + @usableFromInline + struct LeaseResult { + @usableFromInline + var connection: Connection + @usableFromInline + var timersToCancel: Max2Sequence + @usableFromInline + var wasIdle: Bool + @usableFromInline + var use: ConnectionGroup.ConnectionUse + + @inlinable + init( + connection: Connection, + timersToCancel: Max2Sequence, + wasIdle: Bool, + use: ConnectionGroup.ConnectionUse + ) { + self.connection = connection + self.timersToCancel = timersToCancel + self.wasIdle = wasIdle + self.use = use + } + } + + @usableFromInline + struct ConnectionGroup: Sendable { + @usableFromInline + struct Stats: Hashable, Sendable { + @usableFromInline var connecting: UInt16 = 0 + @usableFromInline var backingOff: UInt16 = 0 + @usableFromInline var idle: UInt16 = 0 + @usableFromInline var leased: UInt16 = 0 + @usableFromInline var runningKeepAlive: UInt16 = 0 + @usableFromInline var closing: UInt16 = 0 + + @usableFromInline var availableStreams: UInt16 = 0 + @usableFromInline var leasedStreams: UInt16 = 0 + + @usableFromInline var soonAvailable: UInt16 { + self.connecting + self.backingOff + self.runningKeepAlive + } + + @usableFromInline var active: UInt16 { + self.idle + self.leased + self.connecting + self.backingOff + } + } + + /// The minimum number of connections + @usableFromInline + let minimumConcurrentConnections: Int + + /// The maximum number of preserved connections + @usableFromInline + let maximumConcurrentConnectionSoftLimit: Int + + /// The absolute maximum number of connections + @usableFromInline + let maximumConcurrentConnectionHardLimit: Int + + @usableFromInline + let keepAlive: Bool + + @usableFromInline + let keepAliveReducesAvailableStreams: Bool + + /// A connectionID generator. + @usableFromInline + let generator: ConnectionIDGenerator + + /// The connections states + @usableFromInline + private(set) var connections: [ConnectionState] + + @usableFromInline + private(set) var stats = Stats() + + @inlinable + init( + generator: ConnectionIDGenerator, + minimumConcurrentConnections: Int, + maximumConcurrentConnectionSoftLimit: Int, + maximumConcurrentConnectionHardLimit: Int, + keepAlive: Bool, + keepAliveReducesAvailableStreams: Bool + ) { + self.generator = generator + self.connections = [] + self.minimumConcurrentConnections = minimumConcurrentConnections + self.maximumConcurrentConnectionSoftLimit = maximumConcurrentConnectionSoftLimit + self.maximumConcurrentConnectionHardLimit = maximumConcurrentConnectionHardLimit + self.keepAlive = keepAlive + self.keepAliveReducesAvailableStreams = keepAliveReducesAvailableStreams + } + + var isEmpty: Bool { + self.connections.isEmpty + } + + @usableFromInline + var canGrow: Bool { + self.stats.active < self.maximumConcurrentConnectionHardLimit + } + + @usableFromInline + var soonAvailableConnections: UInt16 { + self.stats.soonAvailable + } + + // MARK: - Mutations - + + /// A connection's use. Is it persisted or an overflow connection? + @usableFromInline + enum ConnectionUse: Equatable { + case persisted + case demand + case overflow + } + + /// Information around an idle connection. + @usableFromInline + struct AvailableConnectionContext { + /// The connection's use. Either general purpose or for requests with `EventLoop` + /// requirements. + @usableFromInline + var use: ConnectionUse + + @usableFromInline + var info: ConnectionAvailableInfo + + @inlinable + init(use: ConnectionUse, info: ConnectionAvailableInfo) { + self.use = use + self.info = info + } + } + + mutating func refillConnections() -> [ConnectionRequest] { + let existingConnections = self.stats.active + let missingConnection = self.minimumConcurrentConnections - Int(existingConnections) + guard missingConnection > 0 else { + return [] + } + + var requests = [ConnectionRequest]() + requests.reserveCapacity(missingConnection) + + for _ in 0.. ConnectionRequest? { + precondition(self.minimumConcurrentConnections <= self.stats.active) + guard self.maximumConcurrentConnectionSoftLimit > self.stats.active else { + return nil + } + return self.createNewConnection() + } + + @inlinable + mutating func createNewOverflowConnectionIfPossible() -> ConnectionRequest? { + precondition(self.maximumConcurrentConnectionSoftLimit <= self.stats.active) + guard self.maximumConcurrentConnectionHardLimit > self.stats.active else { + return nil + } + return self.createNewConnection() + } + + @inlinable + /*private*/ mutating func createNewConnection() -> ConnectionRequest { + precondition(self.canGrow) + self.stats.connecting += 1 + let connectionID = self.generator.next() + let connection = ConnectionState(id: connectionID) + self.connections.append(connection) + return ConnectionRequest(connectionID: connectionID) + } + + /// A new ``Connection`` was established. + /// + /// This will put the connection into the idle state. + /// + /// - Parameter connection: The new established connection. + /// - Returns: An index and an IdleConnectionContext to determine the next action for the now idle connection. + /// Call ``parkConnection(at:)``, ``leaseConnection(at:)`` or ``closeConnection(at:)`` + /// with the supplied index after this. + @inlinable + mutating func newConnectionEstablished(_ connection: Connection, maxStreams: UInt16) -> (Int, AvailableConnectionContext) { + guard let index = self.connections.firstIndex(where: { $0.id == connection.id }) else { + preconditionFailure("There is a new connection that we didn't request!") + } + self.stats.connecting -= 1 + self.stats.idle += 1 + self.stats.availableStreams += maxStreams + let connectionInfo = self.connections[index].connected(connection, maxStreams: maxStreams) + // TODO: If this is an overflow connection, but we are currently also creating a + // persisted connection, we might want to swap those. + let context = self.makeAvailableConnectionContextForConnection(at: index, info: connectionInfo) + return (index, context) + } + + @inlinable + mutating func backoffNextConnectionAttempt(_ connectionID: Connection.ID) -> ConnectionTimer { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + preconditionFailure("We tried to create a new connection that we know nothing about?") + } + + self.stats.connecting -= 1 + self.stats.backingOff += 1 + + return self.connections[index].failedToConnect() + } + + @usableFromInline + enum BackoffDoneAction { + case createConnection(ConnectionRequest, TimerCancellationToken?) + case cancelTimers(Max2Sequence) + } + + @inlinable + mutating func backoffDone(_ connectionID: Connection.ID, retry: Bool) -> BackoffDoneAction { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + preconditionFailure("We tried to create a new connection that we know nothing about?") + } + + self.stats.backingOff -= 1 + + if retry || self.stats.active < self.minimumConcurrentConnections { + self.stats.connecting += 1 + let backoffTimerCancellation = self.connections[index].retryConnect() + return .createConnection(.init(connectionID: connectionID), backoffTimerCancellation) + } + + let backoffTimerCancellation = self.connections[index].destroyBackingOffConnection() + var timerCancellations = Max2Sequence(backoffTimerCancellation) + + if let timerCancellationToken = self.swapForDeletion(index: index) { + timerCancellations.append(timerCancellationToken) + } + return .cancelTimers(timerCancellations) + } + + @inlinable + mutating func timerScheduled( + _ timer: ConnectionTimer, + cancelContinuation: TimerCancellationToken + ) -> TimerCancellationToken? { + guard let index = self.connections.firstIndex(where: { $0.id == timer.connectionID }) else { + return cancelContinuation + } + + return self.connections[index].timerScheduled(timer, cancelContinuation: cancelContinuation) + } + + // MARK: Changes at runtime + + @usableFromInline + struct NewMaxStreamInfo { + + @usableFromInline + var index: Int + + @usableFromInline + var newMaxStreams: UInt16 + + @usableFromInline + var oldMaxStreams: UInt16 + + @usableFromInline + var usedStreams: UInt16 + + @inlinable + init(index: Int, info: ConnectionState.NewMaxStreamInfo) { + self.index = index + self.newMaxStreams = info.newMaxStreams + self.oldMaxStreams = info.oldMaxStreams + self.usedStreams = info.usedStreams + } + } + + @inlinable + mutating func connectionReceivedNewMaxStreamSetting( + _ connectionID: ConnectionID, + newMaxStreamSetting maxStreams: UInt16 + ) -> NewMaxStreamInfo? { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + return nil + } + + guard let info = self.connections[index].newMaxStreamSetting(maxStreams) else { + return nil + } + + self.stats.availableStreams += maxStreams - info.oldMaxStreams + + return NewMaxStreamInfo(index: index, info: info) + } + + // MARK: Leasing and releasing + + /// Lease a connection, if an idle connection is available. + /// + /// - Returns: A connection to execute a request on. + @inlinable + mutating func leaseConnection() -> LeaseResult? { + if self.stats.availableStreams == 0 { + return nil + } + + guard let index = self.findAvailableConnection() else { + preconditionFailure("Stats and actual count are of.") + } + + return self.leaseConnection(at: index, streams: 1) + } + + @usableFromInline + enum LeasedConnectionOrStartingCount { + case leasedConnection(LeaseResult) + case startingCount(UInt16) + } + + @inlinable + mutating func leaseConnectionOrSoonAvailableConnectionCount() -> LeasedConnectionOrStartingCount { + if let result = self.leaseConnection() { + return .leasedConnection(result) + } + return .startingCount(self.stats.soonAvailable) + } + + @inlinable + mutating func leaseConnection(at index: Int, streams: UInt16) -> LeaseResult { + let leaseResult = self.connections[index].lease(streams: streams) + let use = self.getConnectionUse(index: index) + + if leaseResult.wasIdle { + self.stats.idle -= 1 + self.stats.leased += 1 + } + self.stats.leasedStreams += streams + self.stats.availableStreams -= streams + return LeaseResult( + connection: leaseResult.connection, + timersToCancel: leaseResult.timersToCancel, + wasIdle: leaseResult.wasIdle, + use: use + ) + } + + @inlinable + mutating func parkConnection(at index: Int, hasBecomeIdle newIdle: Bool) -> Max2Sequence { + let scheduleIdleTimeoutTimer: Bool + switch index { + case 0.. (Int, AvailableConnectionContext)? { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + return nil + } + + let connectionInfo = self.connections[index].release(streams: streams) + self.stats.availableStreams += streams + self.stats.leasedStreams -= streams + switch connectionInfo { + case .idle: + self.stats.idle += 1 + self.stats.leased -= 1 + case .leased: + break + } + + let context = self.makeAvailableConnectionContextForConnection(at: index, info: connectionInfo) + return (index, context) + } + + @inlinable + mutating func keepAliveIfIdle(_ connectionID: Connection.ID) -> KeepAliveAction? { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + // because of a race this connection (connection close runs against trigger of ping pong) + // was already removed from the state machine. + return nil + } + + guard let action = self.connections[index].runKeepAliveIfIdle(reducesAvailableStreams: self.keepAliveReducesAvailableStreams) else { + return nil + } + + self.stats.runningKeepAlive += 1 + if self.keepAliveReducesAvailableStreams { + self.stats.availableStreams -= 1 + } + + return action + } + + @inlinable + mutating func keepAliveSucceeded(_ connectionID: Connection.ID) -> (Int, AvailableConnectionContext)? { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + // keepAliveSucceeded can race against, closeIfIdle, shutdowns or connection errors + return nil + } + + guard let connectionInfo = self.connections[index].keepAliveSucceeded() else { + // if we don't get connection info here this means, that the connection already was + // transitioned to closing. when we did this we already decremented the + // runningKeepAlive timer. + return nil + } + + self.stats.runningKeepAlive -= 1 + if self.keepAliveReducesAvailableStreams { + self.stats.availableStreams += 1 + } + + let context = self.makeAvailableConnectionContextForConnection(at: index, info: connectionInfo) + return (index, context) + } + + @inlinable + mutating func keepAliveFailed(_ connectionID: Connection.ID) -> CloseAction? { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + // Connection has already been closed + return nil + } + + guard let closeAction = self.connections[index].keepAliveFailed() else { + return nil + } + + self.stats.idle -= 1 + self.stats.closing += 1 + self.stats.runningKeepAlive -= closeAction.runningKeepAlive ? 1 : 0 + self.stats.availableStreams -= closeAction.maxStreams - closeAction.usedStreams + + // force unwrapping the connection is fine, because a close action due to failed + // keepAlive cannot happen without a connection + return CloseAction( + connection: closeAction.connection!, + timersToCancel: closeAction.cancelTimers + ) + } + + // MARK: Connection close/removal + + @usableFromInline + struct CloseAction { + @usableFromInline + private(set) var connection: Connection + + @usableFromInline + private(set) var timersToCancel: Max2Sequence + + @inlinable + init(connection: Connection, timersToCancel: Max2Sequence) { + self.connection = connection + self.timersToCancel = timersToCancel + } + } + + /// Closes the connection at the given index. + @inlinable + mutating func closeConnectionIfIdle(at index: Int) -> CloseAction? { + guard let closeAction = self.connections[index].closeIfIdle() else { + return nil // apparently the connection isn't idle + } + + self.stats.idle -= 1 + self.stats.closing += 1 + self.stats.runningKeepAlive -= closeAction.runningKeepAlive ? 1 : 0 + self.stats.availableStreams -= closeAction.maxStreams - closeAction.usedStreams + + return CloseAction( + connection: closeAction.connection!, + timersToCancel: closeAction.cancelTimers + ) + } + + @inlinable + mutating func closeConnectionIfIdle(_ connectionID: Connection.ID) -> CloseAction? { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + // because of a race this connection (connection close runs against trigger of timeout) + // was already removed from the state machine. + return nil + } + + if index < self.minimumConcurrentConnections { + // because of a race a connection might receive a idle timeout after it was moved into + // the persisted connections. If a connection is now persisted, we now need to ignore + // the trigger + return nil + } + + return self.closeConnectionIfIdle(at: index) + } + + /// Information around the failed/closed connection. + @usableFromInline + struct ClosedAction { + /// Connections that are currently starting + @usableFromInline + var connectionsStarting: Int + + @usableFromInline + var timersToCancel: TinyFastSequence + + @usableFromInline + var newConnectionRequest: ConnectionRequest? + + @inlinable + init( + connectionsStarting: Int, + timersToCancel: TinyFastSequence, + newConnectionRequest: ConnectionRequest? = nil + ) { + self.connectionsStarting = connectionsStarting + self.timersToCancel = timersToCancel + self.newConnectionRequest = newConnectionRequest + } + } + + /// Connection closed. Call this method, if a connection is closed. + /// + /// This will put the position into the closed state. + /// + /// - Parameter connectionID: The failed connection's id. + /// - Returns: An optional index and an IdleConnectionContext to determine the next action for the closed connection. + /// You must call ``removeConnection(at:)`` or ``replaceConnection(at:)`` with the + /// supplied index after this. If nil is returned the connection was closed by the state machine and was + /// therefore already removed. + @inlinable + mutating func connectionClosed(_ connectionID: Connection.ID) -> ClosedAction { + guard let index = self.connections.firstIndex(where: { $0.id == connectionID }) else { + preconditionFailure("All connections that have been created should say goodbye exactly once!") + } + + let closedAction = self.connections[index].closed() + var timersToCancel = TinyFastSequence(closedAction.cancelTimers) + + if closedAction.wasRunningKeepAlive { + self.stats.runningKeepAlive -= 1 + } + self.stats.leasedStreams -= closedAction.usedStreams + self.stats.availableStreams -= closedAction.maxStreams - closedAction.usedStreams + + switch closedAction.previousConnectionState { + case .idle: + self.stats.idle -= 1 + + case .leased: + self.stats.leased -= 1 + + case .closing: + self.stats.closing -= 1 + } + + if let cancellationTimer = self.swapForDeletion(index: index) { + timersToCancel.append(cancellationTimer) + } + + let newConnectionRequest: ConnectionRequest? + if self.connections.count < self.minimumConcurrentConnections { + newConnectionRequest = self.createNewConnection() + } else { + newConnectionRequest = .none + } + + return ClosedAction( + connectionsStarting: 0, + timersToCancel: timersToCancel, + newConnectionRequest: newConnectionRequest + ) + } + + // MARK: Shutdown + + mutating func triggerForceShutdown(_ cleanup: inout ConnectionAction.Shutdown) { + for var connectionState in self.connections { + guard let closeAction = connectionState.close() else { + continue + } + + if let connection = closeAction.connection { + cleanup.connections.append(connection) + } + cleanup.timersToCancel.append(contentsOf: closeAction.cancelTimers) + } + + self.connections = [] + } + + // MARK: - Private functions - + + @inlinable + /*private*/ func getConnectionUse(index: Int) -> ConnectionUse { + switch index { + case 0.. AvailableConnectionContext { + precondition(self.connections[index].isAvailable) + let use = self.getConnectionUse(index: index) + return AvailableConnectionContext(use: use, info: info) + } + + @inlinable + /*private*/ func findAvailableConnection() -> Int? { + return self.connections.firstIndex(where: { $0.isAvailable }) + } + + @inlinable + /*private*/ mutating func swapForDeletion(index indexToDelete: Int) -> TimerCancellationToken? { + let maybeLastConnectedIndex = self.connections.lastIndex(where: { $0.isConnected }) + + if maybeLastConnectedIndex == nil || maybeLastConnectedIndex! < indexToDelete { + self.removeO1(indexToDelete) + return nil + } + + // if maybeLastConnectedIndex == nil, we return early in the above if case. + let lastConnectedIndex = maybeLastConnectedIndex! + + switch indexToDelete { + case 0.. TimerCancellationToken? { + switch self { + case .scheduled(let timer): + self = .notScheduled + return timer.cancellationContinuation + case .running, .notScheduled: + return nil + } + } + } + + @usableFromInline + struct Timer: Sendable { + @usableFromInline + let timerID: Int + + @usableFromInline + private(set) var cancellationContinuation: TimerCancellationToken? + + @inlinable + init(id: Int) { + self.timerID = id + self.cancellationContinuation = nil + } + + @inlinable + mutating func registerCancellationContinuation(_ continuation: TimerCancellationToken) { + precondition(self.cancellationContinuation == nil) + self.cancellationContinuation = continuation + } + } + + /// The pool is creating a connection. Valid transitions are to: `.backingOff`, `.idle`, and `.closed` + case starting + /// The pool is waiting to retry establishing a connection. Valid transitions are to: `.closed`. + /// This means, the connection can be removed from the connections without cancelling external + /// state. The connection state can then be replaced by a new one. + case backingOff(Timer) + /// The connection is `idle` and ready to execute a new query. Valid transitions to: `.pingpong`, `.leased`, + /// `.closing` and `.closed` + case idle(Connection, maxStreams: UInt16, keepAlive: KeepAlive, idleTimer: Timer?) + /// The connection is leased and executing a query. Valid transitions to: `.idle` and `.closed` + case leased(Connection, usedStreams: UInt16, maxStreams: UInt16, keepAlive: KeepAlive) + /// The connection is closing. Valid transitions to: `.closed` + case closing(Connection) + /// The connection is closed. Final state. + case closed + } + + @usableFromInline + let id: Connection.ID + + @usableFromInline + private(set) var state: State = .starting + + @usableFromInline + private(set) var nextTimerID: Int = 0 + + @inlinable + init(id: Connection.ID) { + self.id = id + } + + @inlinable + var isIdle: Bool { + switch self.state { + case .idle(_, _, .notScheduled, _), .idle(_, _, .scheduled, _): + return true + case .idle(_, _, .running, _): + return false + case .backingOff, .starting, .closed, .closing, .leased: + return false + } + } + + @inlinable + var isAvailable: Bool { + switch self.state { + case .idle(_, let maxStreams, .running(true), _): + return maxStreams > 1 + case .idle(_, let maxStreams, let keepAlive, _): + return keepAlive.usedStreams < maxStreams + case .leased(_, let usedStreams, let maxStreams, let keepAlive): + return usedStreams + keepAlive.usedStreams < maxStreams + case .backingOff, .starting, .closed, .closing: + return false + } + } + + @inlinable + var isLeased: Bool { + switch self.state { + case .leased: + return true + case .backingOff, .starting, .closed, .closing, .idle: + return false + } + } + + @inlinable + var isConnected: Bool { + switch self.state { + case .idle, .leased: + return true + case .backingOff, .starting, .closed, .closing: + return false + } + } + + @inlinable + mutating func connected(_ connection: Connection, maxStreams: UInt16) -> ConnectionAvailableInfo { + switch self.state { + case .starting: + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: .notScheduled, idleTimer: nil) + return .idle(availableStreams: maxStreams, newIdle: true) + case .backingOff, .idle, .leased, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + } + } + + @usableFromInline + struct NewMaxStreamInfo { + @usableFromInline + var newMaxStreams: UInt16 + + @usableFromInline + var oldMaxStreams: UInt16 + + @usableFromInline + var usedStreams: UInt16 + + @inlinable + init(newMaxStreams: UInt16, oldMaxStreams: UInt16, usedStreams: UInt16) { + self.newMaxStreams = newMaxStreams + self.oldMaxStreams = oldMaxStreams + self.usedStreams = usedStreams + } + } + + @inlinable + mutating func newMaxStreamSetting(_ newMaxStreams: UInt16) -> NewMaxStreamInfo? { + switch self.state { + case .starting, .backingOff: + preconditionFailure("Invalid state: \(self.state)") + + case .idle(let connection, let oldMaxStreams, let keepAlive, idleTimer: let idleTimer): + self.state = .idle(connection, maxStreams: newMaxStreams, keepAlive: keepAlive, idleTimer: idleTimer) + return NewMaxStreamInfo( + newMaxStreams: newMaxStreams, + oldMaxStreams: oldMaxStreams, + usedStreams: keepAlive.usedStreams + ) + + case .leased(let connection, let usedStreams, let oldMaxStreams, let keepAlive): + self.state = .leased(connection, usedStreams: usedStreams, maxStreams: newMaxStreams, keepAlive: keepAlive) + return NewMaxStreamInfo( + newMaxStreams: newMaxStreams, + oldMaxStreams: oldMaxStreams, + usedStreams: usedStreams + keepAlive.usedStreams + ) + + case .closing, .closed: + return nil + } + } + + + @inlinable + mutating func parkConnection(scheduleKeepAliveTimer: Bool, scheduleIdleTimeoutTimer: Bool) -> Max2Sequence { + var keepAliveTimer: ConnectionTimer? + var keepAliveTimerState: State.Timer? + var idleTimer: ConnectionTimer? + var idleTimerState: State.Timer? + + switch self.state { + case .backingOff, .starting, .leased, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + + case .idle(let connection, let maxStreams, .notScheduled, .none): + let keepAlive: State.KeepAlive + if scheduleKeepAliveTimer { + keepAliveTimerState = self._nextTimer() + keepAliveTimer = ConnectionTimer(timerID: keepAliveTimerState!.timerID, connectionID: self.id, usecase: .keepAlive) + keepAlive = .scheduled(keepAliveTimerState!) + } else { + keepAlive = .notScheduled + } + if scheduleIdleTimeoutTimer { + idleTimerState = self._nextTimer() + idleTimer = ConnectionTimer(timerID: idleTimerState!.timerID, connectionID: self.id, usecase: .idleTimeout) + } + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: keepAlive, idleTimer: idleTimerState) + return Max2Sequence(keepAliveTimer, idleTimer) + + case .idle(_, _, .scheduled, .some): + precondition(!scheduleKeepAliveTimer) + precondition(!scheduleIdleTimeoutTimer) + return Max2Sequence() + + case .idle(let connection, let maxStreams, .notScheduled, let idleTimer): + precondition(!scheduleIdleTimeoutTimer) + let keepAlive: State.KeepAlive + if scheduleKeepAliveTimer { + keepAliveTimerState = self._nextTimer() + keepAliveTimer = ConnectionTimer(timerID: keepAliveTimerState!.timerID, connectionID: self.id, usecase: .keepAlive) + keepAlive = .scheduled(keepAliveTimerState!) + } else { + keepAlive = .notScheduled + } + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: keepAlive, idleTimer: idleTimer) + return Max2Sequence(keepAliveTimer) + + case .idle(let connection, let maxStreams, .scheduled(let keepAliveTimer), .none): + precondition(!scheduleKeepAliveTimer) + + if scheduleIdleTimeoutTimer { + idleTimerState = self._nextTimer() + idleTimer = ConnectionTimer(timerID: idleTimerState!.timerID, connectionID: self.id, usecase: .keepAlive) + } + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: .scheduled(keepAliveTimer), idleTimer: idleTimerState) + return Max2Sequence(idleTimer, nil) + + case .idle(let connection, let maxStreams, keepAlive: .running(let usingStream), idleTimer: .none): + if scheduleIdleTimeoutTimer { + idleTimerState = self._nextTimer() + idleTimer = ConnectionTimer(timerID: idleTimerState!.timerID, connectionID: self.id, usecase: .keepAlive) + } + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: .running(usingStream), idleTimer: idleTimerState) + return Max2Sequence(keepAliveTimer, idleTimer) + + case .idle(_, _, keepAlive: .running(_), idleTimer: .some): + precondition(!scheduleKeepAliveTimer) + precondition(!scheduleIdleTimeoutTimer) + return Max2Sequence() + } + } + + /// The connection failed to start + @inlinable + mutating func failedToConnect() -> ConnectionTimer { + switch self.state { + case .starting: + let backoffTimerState = self._nextTimer() + self.state = .backingOff(backoffTimerState) + return ConnectionTimer(timerID: backoffTimerState.timerID, connectionID: self.id, usecase: .backoff) + + case .backingOff, .idle, .leased, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + } + } + + /// Moves a connection, that has previously ``failedToConnect()`` back into the connecting state. + /// + /// - Returns: A ``TimerCancellationToken`` that was previously registered with the state machine + /// for the ``ConnectionTimer`` returned in ``failedToConnect()``. If no token was registered + /// nil is returned. + @inlinable + mutating func retryConnect() -> TimerCancellationToken? { + switch self.state { + case .backingOff(let timer): + self.state = .starting + return timer.cancellationContinuation + case .starting, .idle, .leased, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + } + } + + @inlinable + mutating func destroyBackingOffConnection() -> TimerCancellationToken? { + switch self.state { + case .backingOff(let timer): + self.state = .closed + return timer.cancellationContinuation + case .starting, .idle, .leased, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + } + } + + @usableFromInline + struct LeaseAction { + @usableFromInline + var connection: Connection + @usableFromInline + var timersToCancel: Max2Sequence + @usableFromInline + var wasIdle: Bool + + @inlinable + init(connection: Connection, timersToCancel: Max2Sequence, wasIdle: Bool) { + self.connection = connection + self.timersToCancel = timersToCancel + self.wasIdle = wasIdle + } + } + + @inlinable + mutating func lease(streams newLeasedStreams: UInt16 = 1) -> LeaseAction { + switch self.state { + case .idle(let connection, let maxStreams, var keepAlive, let idleTimer): + var cancel = Max2Sequence() + if let token = idleTimer?.cancellationContinuation { + cancel.append(token) + } + if let token = keepAlive.cancelTimerIfScheduled() { + cancel.append(token) + } + precondition(maxStreams >= newLeasedStreams + keepAlive.usedStreams, "Invalid state: \(self.state)") + self.state = .leased(connection, usedStreams: newLeasedStreams, maxStreams: maxStreams, keepAlive: keepAlive) + return LeaseAction(connection: connection, timersToCancel: cancel, wasIdle: true) + + case .leased(let connection, let usedStreams, let maxStreams, let keepAlive): + precondition(maxStreams >= usedStreams + newLeasedStreams + keepAlive.usedStreams, "Invalid state: \(self.state)") + self.state = .leased(connection, usedStreams: usedStreams + newLeasedStreams, maxStreams: maxStreams, keepAlive: keepAlive) + return LeaseAction(connection: connection, timersToCancel: .init(), wasIdle: false) + + case .backingOff, .starting, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + } + } + + @inlinable + mutating func release(streams returnedStreams: UInt16) -> ConnectionAvailableInfo { + switch self.state { + case .leased(let connection, let usedStreams, let maxStreams, let keepAlive): + precondition(usedStreams >= returnedStreams) + let newUsedStreams = usedStreams - returnedStreams + let availableStreams = maxStreams - (newUsedStreams + keepAlive.usedStreams) + if newUsedStreams == 0 { + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: keepAlive, idleTimer: nil) + return .idle(availableStreams: availableStreams, newIdle: true) + } else { + self.state = .leased(connection, usedStreams: newUsedStreams, maxStreams: maxStreams, keepAlive: keepAlive) + return .leased(availableStreams: availableStreams) + } + case .backingOff, .starting, .idle, .closing, .closed: + preconditionFailure("Invalid state: \(self.state)") + } + } + + @inlinable + mutating func runKeepAliveIfIdle(reducesAvailableStreams: Bool) -> KeepAliveAction? { + switch self.state { + case .idle(let connection, let maxStreams, .scheduled(let timer), let idleTimer): + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: .running(reducesAvailableStreams), idleTimer: idleTimer) + return KeepAliveAction( + connection: connection, + keepAliveTimerCancellationContinuation: timer.cancellationContinuation + ) + + case .leased, .closed, .closing: + return nil + + case .backingOff, .starting, .idle(_, _, .running, _), .idle(_, _, .notScheduled, _): + preconditionFailure("Invalid state: \(self.state)") + } + } + + @inlinable + mutating func keepAliveSucceeded() -> ConnectionAvailableInfo? { + switch self.state { + case .idle(let connection, let maxStreams, .running, let idleTimer): + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: .notScheduled, idleTimer: idleTimer) + return .idle(availableStreams: maxStreams, newIdle: false) + + case .leased(let connection, let usedStreams, let maxStreams, .running): + self.state = .leased(connection, usedStreams: usedStreams, maxStreams: maxStreams, keepAlive: .notScheduled) + return .leased(availableStreams: maxStreams - usedStreams) + + case .closed, .closing: + return nil + + case .backingOff, .starting, + .leased(_, _, _, .notScheduled), + .leased(_, _, _, .scheduled), + .idle(_, _, .notScheduled, _), + .idle(_, _, .scheduled, _): + preconditionFailure("Invalid state: \(self.state)") + } + } + + @inlinable + mutating func keepAliveFailed() -> CloseAction? { + return self.close() + } + + @inlinable + mutating func timerScheduled( + _ timer: ConnectionTimer, + cancelContinuation: TimerCancellationToken + ) -> TimerCancellationToken? { + switch timer.usecase { + case .backoff: + switch self.state { + case .backingOff(var timerState): + if timerState.timerID == timer.timerID { + timerState.registerCancellationContinuation(cancelContinuation) + self.state = .backingOff(timerState) + return nil + } else { + return cancelContinuation + } + + case .starting, .idle, .leased, .closing, .closed: + return cancelContinuation + } + + case .idleTimeout: + switch self.state { + case .idle(let connection, let maxStreams, let keepAlive, let idleTimerState): + if var idleTimerState = idleTimerState, idleTimerState.timerID == timer.timerID { + idleTimerState.registerCancellationContinuation(cancelContinuation) + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: keepAlive, idleTimer: idleTimerState) + return nil + } else { + return cancelContinuation + } + + case .starting, .backingOff, .leased, .closing, .closed: + return cancelContinuation + } + + case .keepAlive: + switch self.state { + case .idle(let connection, let maxStreams, .scheduled(var keepAliveTimerState), let idleTimerState): + if keepAliveTimerState.timerID == timer.timerID { + keepAliveTimerState.registerCancellationContinuation(cancelContinuation) + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: .scheduled(keepAliveTimerState), idleTimer: idleTimerState) + return nil + } else { + return cancelContinuation + } + + case .starting, .backingOff, .leased, .closing, .closed, + .idle(_, _, .running, _), + .idle(_, _, .notScheduled, _): + return cancelContinuation + } + } + } + + @inlinable + mutating func cancelIdleTimer() -> TimerCancellationToken? { + switch self.state { + case .starting, .backingOff, .leased, .closing, .closed: + return nil + + case .idle(let connection, let maxStreams, let keepAlive, let idleTimer): + self.state = .idle(connection, maxStreams: maxStreams, keepAlive: keepAlive, idleTimer: nil) + return idleTimer?.cancellationContinuation + } + } + + @usableFromInline + struct CloseAction { + + @usableFromInline + enum PreviousConnectionState { + case idle + case leased + case closing + case backingOff + } + + @usableFromInline + var connection: Connection? + @usableFromInline + var previousConnectionState: PreviousConnectionState + @usableFromInline + var cancelTimers: Max2Sequence + @usableFromInline + var usedStreams: UInt16 + @usableFromInline + var maxStreams: UInt16 + @usableFromInline + var runningKeepAlive: Bool + + + @inlinable + init( + connection: Connection?, + previousConnectionState: PreviousConnectionState, + cancelTimers: Max2Sequence, + usedStreams: UInt16, + maxStreams: UInt16, + runningKeepAlive: Bool + ) { + self.connection = connection + self.previousConnectionState = previousConnectionState + self.cancelTimers = cancelTimers + self.usedStreams = usedStreams + self.maxStreams = maxStreams + self.runningKeepAlive = runningKeepAlive + } + } + + @inlinable + mutating func closeIfIdle() -> CloseAction? { + switch self.state { + case .idle(let connection, let maxStreams, var keepAlive, let idleTimerState): + self.state = .closing(connection) + return CloseAction( + connection: connection, + previousConnectionState: .idle, + cancelTimers: Max2Sequence( + keepAlive.cancelTimerIfScheduled(), + idleTimerState?.cancellationContinuation + ), + usedStreams: keepAlive.usedStreams, + maxStreams: maxStreams, + runningKeepAlive: keepAlive.isRunning + ) + + case .leased, .closed: + return nil + + case .backingOff, .starting, .closing: + preconditionFailure("Invalid state: \(self.state)") + } + } + + @inlinable + mutating func close() -> CloseAction? { + switch self.state { + case .starting: + // If we are currently starting, there is nothing we can do about it right now. + // Only once the connection has come up, or failed, we can actually act. + return nil + + case .closing, .closed: + // If we are already closing, we can't do anything else. + return nil + + case .idle(let connection, let maxStreams, var keepAlive, let idleTimerState): + self.state = .closing(connection) + return CloseAction( + connection: connection, + previousConnectionState: .idle, + cancelTimers: Max2Sequence( + keepAlive.cancelTimerIfScheduled(), + idleTimerState?.cancellationContinuation + ), + usedStreams: keepAlive.usedStreams, + maxStreams: maxStreams, + runningKeepAlive: keepAlive.isRunning + ) + + case .leased(let connection, usedStreams: let usedStreams, maxStreams: let maxStreams, var keepAlive): + self.state = .closing(connection) + return CloseAction( + connection: connection, + previousConnectionState: .leased, + cancelTimers: Max2Sequence( + keepAlive.cancelTimerIfScheduled() + ), + usedStreams: keepAlive.usedStreams + usedStreams, + maxStreams: maxStreams, + runningKeepAlive: keepAlive.isRunning + ) + + case .backingOff(let timer): + self.state = .closed + return CloseAction( + connection: nil, + previousConnectionState: .backingOff, + cancelTimers: Max2Sequence(timer.cancellationContinuation), + usedStreams: 0, + maxStreams: 0, + runningKeepAlive: false + ) + } + } + + @usableFromInline + struct ClosedAction { + + @usableFromInline + enum PreviousConnectionState { + case idle + case leased + case closing + } + + @usableFromInline + var previousConnectionState: PreviousConnectionState + @usableFromInline + var cancelTimers: Max2Sequence + @usableFromInline + var maxStreams: UInt16 + @usableFromInline + var usedStreams: UInt16 + @usableFromInline + var wasRunningKeepAlive: Bool + + @inlinable + init( + previousConnectionState: PreviousConnectionState, + cancelTimers: Max2Sequence, + maxStreams: UInt16, + usedStreams: UInt16, + wasRunningKeepAlive: Bool + ) { + self.previousConnectionState = previousConnectionState + self.cancelTimers = cancelTimers + self.maxStreams = maxStreams + self.usedStreams = usedStreams + self.wasRunningKeepAlive = wasRunningKeepAlive + } + } + + @inlinable + mutating func closed() -> ClosedAction { + switch self.state { + case .starting, .backingOff, .closed: + preconditionFailure("Invalid state: \(self.state)") + + case .idle(_, let maxStreams, var keepAlive, let idleTimer): + self.state = .closed + return ClosedAction( + previousConnectionState: .idle, + cancelTimers: .init(keepAlive.cancelTimerIfScheduled(), idleTimer?.cancellationContinuation), + maxStreams: maxStreams, + usedStreams: keepAlive.usedStreams, + wasRunningKeepAlive: keepAlive.isRunning + ) + + case .leased(_, let usedStreams, let maxStreams, let keepAlive): + self.state = .closed + return ClosedAction( + previousConnectionState: .leased, + cancelTimers: .init(), + maxStreams: maxStreams, + usedStreams: usedStreams + keepAlive.usedStreams, + wasRunningKeepAlive: keepAlive.isRunning + ) + + case .closing: + self.state = .closed + return ClosedAction( + previousConnectionState: .closing, + cancelTimers: .init(), + maxStreams: 0, + usedStreams: 0, + wasRunningKeepAlive: false + ) + } + } + + // MARK: - Private Methods - + + @inlinable + mutating /*private*/ func _nextTimer() -> State.Timer { + defer { self.nextTimerID += 1 } + return State.Timer(id: self.nextTimerID) + } + } + + @usableFromInline + enum ConnectionAvailableInfo: Equatable { + case leased(availableStreams: UInt16) + case idle(availableStreams: UInt16, newIdle: Bool) + + @usableFromInline + var availableStreams: UInt16 { + switch self { + case .leased(let availableStreams): + return availableStreams + case .idle(let availableStreams, newIdle: _): + return availableStreams + } + } + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.KeepAliveAction: Equatable where TimerCancellationToken: Equatable { + @inlinable + static func == (lhs: Self, rhs: Self) -> Bool { + lhs.connection === rhs.connection && lhs.keepAliveTimerCancellationContinuation == rhs.keepAliveTimerCancellationContinuation + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.ConnectionState.LeaseAction: Equatable where TimerCancellationToken: Equatable { + @inlinable + static func == (lhs: Self, rhs: Self) -> Bool { + lhs.wasIdle == rhs.wasIdle && lhs.connection === rhs.connection && lhs.timersToCancel == rhs.timersToCancel + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.ConnectionState.CloseAction: Equatable where TimerCancellationToken: Equatable { + @inlinable + static func == (lhs: Self, rhs: Self) -> Bool { + lhs.cancelTimers == rhs.cancelTimers && lhs.connection === rhs.connection && lhs.maxStreams == rhs.maxStreams + } +} diff --git a/Sources/ConnectionPoolModule/PoolStateMachine+RequestQueue.swift b/Sources/ConnectionPoolModule/PoolStateMachine+RequestQueue.swift new file mode 100644 index 00000000..99ec4896 --- /dev/null +++ b/Sources/ConnectionPoolModule/PoolStateMachine+RequestQueue.swift @@ -0,0 +1,71 @@ +import DequeModule + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine { + + /// A request queue, which can enqueue requests in O(1), dequeue requests in O(1) and even cancel requests in O(1). + /// + /// While enqueueing and dequeueing on O(1) is trivial, cancellation is hard, as it normally requires a removal within the + /// underlying Deque. However thanks to having an additional `requests` dictionary, we can remove the cancelled + /// request from the dictionary and keep it inside the queue. Whenever we pop a request from the deque, we validate + /// that it hasn't been cancelled in the meantime by checking if the popped request is still in the `requests` dictionary. + @usableFromInline + struct RequestQueue: Sendable { + @usableFromInline + private(set) var queue: Deque + + @usableFromInline + private(set) var requests: [RequestID: Request] + + @inlinable + var count: Int { + self.requests.count + } + + @inlinable + var isEmpty: Bool { + self.count == 0 + } + + @usableFromInline + init() { + self.queue = .init(minimumCapacity: 256) + self.requests = .init(minimumCapacity: 256) + } + + @inlinable + mutating func queue(_ request: Request) { + self.requests[request.id] = request + self.queue.append(request.id) + } + + @inlinable + mutating func pop(max: UInt16) -> TinyFastSequence { + var result = TinyFastSequence() + result.reserveCapacity(Int(max)) + var popped = 0 + while popped < max, let requestID = self.queue.popFirst() { + if let requestIndex = self.requests.index(forKey: requestID) { + popped += 1 + result.append(self.requests.remove(at: requestIndex).value) + } + } + + assert(result.count <= max) + return result + } + + @inlinable + mutating func remove(_ requestID: RequestID) -> Request? { + self.requests.removeValue(forKey: requestID) + } + + @inlinable + mutating func removeAll() -> TinyFastSequence { + let result = TinyFastSequence(self.requests.values) + self.requests.removeAll() + self.queue.removeAll() + return result + } + } +} diff --git a/Sources/ConnectionPoolModule/PoolStateMachine.swift b/Sources/ConnectionPoolModule/PoolStateMachine.swift new file mode 100644 index 00000000..6e41f730 --- /dev/null +++ b/Sources/ConnectionPoolModule/PoolStateMachine.swift @@ -0,0 +1,635 @@ +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif + +@usableFromInline +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +struct PoolConfiguration: Sendable { + /// The minimum number of connections to preserve in the pool. + /// + /// If the pool is mostly idle and the remote servers closes idle connections, + /// the `ConnectionPool` will initiate new outbound connections proactively + /// to avoid the number of available connections dropping below this number. + @usableFromInline + var minimumConnectionCount: Int = 0 + + /// The maximum number of connections to for this pool, to be preserved. + @usableFromInline + var maximumConnectionSoftLimit: Int = 10 + + @usableFromInline + var maximumConnectionHardLimit: Int = 10 + + @usableFromInline + var keepAliveDuration: Duration? + + @usableFromInline + var idleTimeoutDuration: Duration = .seconds(30) +} + +@usableFromInline +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +struct PoolStateMachine< + Connection: PooledConnection, + ConnectionIDGenerator: ConnectionIDGeneratorProtocol, + ConnectionID: Hashable & Sendable, + Request: ConnectionRequestProtocol, + RequestID, + TimerCancellationToken: Sendable +>: Sendable where Connection.ID == ConnectionID, ConnectionIDGenerator.ID == ConnectionID, RequestID == Request.ID { + + @usableFromInline + struct ConnectionRequest: Hashable, Sendable { + @usableFromInline var connectionID: ConnectionID + + @inlinable + init(connectionID: ConnectionID) { + self.connectionID = connectionID + } + } + + @usableFromInline + struct Action { + @usableFromInline let request: RequestAction + @usableFromInline let connection: ConnectionAction + + @inlinable + init(request: RequestAction, connection: ConnectionAction) { + self.request = request + self.connection = connection + } + + @inlinable + static func none() -> Action { Action(request: .none, connection: .none) } + } + + @usableFromInline + enum ConnectionAction { + @usableFromInline + struct Shutdown { + @usableFromInline + var connections: [Connection] + @usableFromInline + var timersToCancel: [TimerCancellationToken] + + @inlinable + init() { + self.connections = [] + self.timersToCancel = [] + } + } + + case scheduleTimers(Max2Sequence) + case makeConnection(ConnectionRequest, TinyFastSequence) + case runKeepAlive(Connection, TimerCancellationToken?) + case cancelTimers(TinyFastSequence) + case closeConnection(Connection, Max2Sequence) + case shutdown(Shutdown) + + case none + } + + @usableFromInline + enum RequestAction { + case leaseConnection(TinyFastSequence, Connection) + + case failRequest(Request, ConnectionPoolError) + case failRequests(TinyFastSequence, ConnectionPoolError) + + case none + } + + @usableFromInline + enum PoolState: Sendable { + case running + case shuttingDown(graceful: Bool) + case shutDown + } + + @usableFromInline + struct Timer: Hashable, Sendable { + @usableFromInline + var underlying: ConnectionTimer + + @usableFromInline + var duration: Duration + + @inlinable + var connectionID: ConnectionID { + self.underlying.connectionID + } + + @inlinable + init(_ connectionTimer: ConnectionTimer, duration: Duration) { + self.underlying = connectionTimer + self.duration = duration + } + } + + @usableFromInline let configuration: PoolConfiguration + @usableFromInline let generator: ConnectionIDGenerator + + @usableFromInline + private(set) var connections: ConnectionGroup + @usableFromInline + private(set) var requestQueue: RequestQueue + @usableFromInline + private(set) var poolState: PoolState = .running + @usableFromInline + private(set) var cacheNoMoreConnectionsAllowed: Bool = false + + @usableFromInline + private(set) var failedConsecutiveConnectionAttempts: Int = 0 + + @inlinable + init( + configuration: PoolConfiguration, + generator: ConnectionIDGenerator, + timerCancellationTokenType: TimerCancellationToken.Type + ) { + self.configuration = configuration + self.generator = generator + self.connections = ConnectionGroup( + generator: generator, + minimumConcurrentConnections: configuration.minimumConnectionCount, + maximumConcurrentConnectionSoftLimit: configuration.maximumConnectionSoftLimit, + maximumConcurrentConnectionHardLimit: configuration.maximumConnectionHardLimit, + keepAlive: configuration.keepAliveDuration != nil, + keepAliveReducesAvailableStreams: true + ) + self.requestQueue = RequestQueue() + } + + mutating func refillConnections() -> [ConnectionRequest] { + return self.connections.refillConnections() + } + + @inlinable + mutating func leaseConnection(_ request: Request) -> Action { + switch self.poolState { + case .running: + break + + case .shuttingDown, .shutDown: + return .init( + request: .failRequest(request, ConnectionPoolError.poolShutdown), + connection: .none + ) + } + + if !self.requestQueue.isEmpty && self.cacheNoMoreConnectionsAllowed { + self.requestQueue.queue(request) + return .none() + } + + var soonAvailable: UInt16 = 0 + + // check if any other EL has an idle connection + switch self.connections.leaseConnectionOrSoonAvailableConnectionCount() { + case .leasedConnection(let leaseResult): + return .init( + request: .leaseConnection(TinyFastSequence(element: request), leaseResult.connection), + connection: .cancelTimers(.init(leaseResult.timersToCancel)) + ) + + case .startingCount(let count): + soonAvailable += count + } + + // we tried everything. there is no connection available. now we must check, if and where we + // can create further connections. but first we must enqueue the new request + + self.requestQueue.queue(request) + + let requestAction = RequestAction.none + + if soonAvailable >= self.requestQueue.count { + // if more connections will be soon available then we have waiters, we don't need to + // create further new connections. + return .init( + request: requestAction, + connection: .none + ) + } else if let request = self.connections.createNewDemandConnectionIfPossible() { + // Can we create a demand connection + return .init( + request: requestAction, + connection: .makeConnection(request, .init()) + ) + } else if let request = self.connections.createNewOverflowConnectionIfPossible() { + // Can we create an overflow connection + return .init( + request: requestAction, + connection: .makeConnection(request, .init()) + ) + } else { + self.cacheNoMoreConnectionsAllowed = true + + // no new connections allowed: + return .init(request: requestAction, connection: .none) + } + } + + @inlinable + mutating func releaseConnection(_ connection: Connection, streams: UInt16) -> Action { + guard let (index, context) = self.connections.releaseConnection(connection.id, streams: streams) else { + return .none() + } + return self.handleAvailableConnection(index: index, availableContext: context) + } + + mutating func cancelRequest(id: RequestID) -> Action { + guard let request = self.requestQueue.remove(id) else { + return .none() + } + + return .init( + request: .failRequest(request, ConnectionPoolError.requestCancelled), + connection: .none + ) + } + + @inlinable + mutating func connectionEstablished(_ connection: Connection, maxStreams: UInt16) -> Action { + switch self.poolState { + case .running, .shuttingDown(graceful: true): + let (index, context) = self.connections.newConnectionEstablished(connection, maxStreams: maxStreams) + return self.handleAvailableConnection(index: index, availableContext: context) + case .shuttingDown(graceful: false), .shutDown: + return .init(request: .none, connection: .closeConnection(connection, [])) + } + } + + @inlinable + mutating func connectionReceivedNewMaxStreamSetting( + _ connection: ConnectionID, + newMaxStreamSetting maxStreams: UInt16 + ) -> Action { + guard let info = self.connections.connectionReceivedNewMaxStreamSetting(connection, newMaxStreamSetting: maxStreams) else { + return .none() + } + + let waitingRequests = self.requestQueue.count + + guard waitingRequests > 0 else { + return .none() + } + + // the only thing we can do if we receive a new max stream setting is check if the new stream + // setting is higher and then dequeue some waiting requests + + guard info.newMaxStreams > info.oldMaxStreams && info.newMaxStreams > info.usedStreams else { + return .none() + } + + let leaseStreams = min(info.newMaxStreams - info.oldMaxStreams, info.newMaxStreams - info.usedStreams, UInt16(clamping: waitingRequests)) + let requests = self.requestQueue.pop(max: leaseStreams) + precondition(Int(leaseStreams) == requests.count) + let leaseResult = self.connections.leaseConnection(at: info.index, streams: leaseStreams) + + return .init( + request: .leaseConnection(requests, leaseResult.connection), + connection: .cancelTimers(.init(leaseResult.timersToCancel)) + ) + } + + @inlinable + mutating func timerScheduled(_ timer: Timer, cancelContinuation: TimerCancellationToken) -> TimerCancellationToken? { + self.connections.timerScheduled(timer.underlying, cancelContinuation: cancelContinuation) + } + + @inlinable + mutating func timerTriggered(_ timer: Timer) -> Action { + switch timer.underlying.usecase { + case .backoff: + return self.connectionCreationBackoffDone(timer.connectionID) + case .keepAlive: + return self.connectionKeepAliveTimerTriggered(timer.connectionID) + case .idleTimeout: + return self.connectionIdleTimerTriggered(timer.connectionID) + } + } + + @inlinable + mutating func connectionEstablishFailed(_ error: Error, for request: ConnectionRequest) -> Action { + switch self.poolState { + case .running, .shuttingDown(graceful: true): + self.failedConsecutiveConnectionAttempts += 1 + + let connectionTimer = self.connections.backoffNextConnectionAttempt(request.connectionID) + let backoff = Self.calculateBackoff(failedAttempt: self.failedConsecutiveConnectionAttempts) + let timer = Timer(connectionTimer, duration: backoff) + return .init(request: .none, connection: .scheduleTimers(.init(timer))) + + case .shuttingDown(graceful: false), .shutDown: + return .none() + } + } + + @inlinable + mutating func connectionCreationBackoffDone(_ connectionID: ConnectionID) -> Action { + switch self.poolState { + case .running, .shuttingDown(graceful: true): + let soonAvailable = self.connections.soonAvailableConnections + let retry = (soonAvailable - 1) < self.requestQueue.count + + switch self.connections.backoffDone(connectionID, retry: retry) { + case .createConnection(let request, let continuation): + let timers: TinyFastSequence + if let continuation { + timers = .init(element: continuation) + } else { + timers = .init() + } + return .init(request: .none, connection: .makeConnection(request, timers)) + + case .cancelTimers(let timers): + return .init(request: .none, connection: .cancelTimers(.init(timers))) + } + + case .shuttingDown(graceful: false), .shutDown: + return .none() + } + } + + @inlinable + mutating func connectionKeepAliveTimerTriggered(_ connectionID: ConnectionID) -> Action { + precondition(self.configuration.keepAliveDuration != nil) + precondition(self.requestQueue.isEmpty) + + guard let keepAliveAction = self.connections.keepAliveIfIdle(connectionID) else { + return .none() + } + return .init(request: .none, connection: .runKeepAlive(keepAliveAction.connection, keepAliveAction.keepAliveTimerCancellationContinuation)) + } + + @inlinable + mutating func connectionKeepAliveDone(_ connection: Connection) -> Action { + precondition(self.configuration.keepAliveDuration != nil) + guard let (index, context) = self.connections.keepAliveSucceeded(connection.id) else { + return .none() + } + return self.handleAvailableConnection(index: index, availableContext: context) + } + + @inlinable + mutating func connectionKeepAliveFailed(_ connectionID: ConnectionID) -> Action { + guard let closeAction = self.connections.keepAliveFailed(connectionID) else { + return .none() + } + + return .init(request: .none, connection: .closeConnection(closeAction.connection, closeAction.timersToCancel)) + } + + @inlinable + mutating func connectionIdleTimerTriggered(_ connectionID: ConnectionID) -> Action { + precondition(self.requestQueue.isEmpty) + + guard let closeAction = self.connections.closeConnectionIfIdle(connectionID) else { + return .none() + } + + self.cacheNoMoreConnectionsAllowed = false + return .init(request: .none, connection: .closeConnection(closeAction.connection, closeAction.timersToCancel)) + } + + @inlinable + mutating func connectionClosed(_ connection: Connection) -> Action { + switch self.poolState { + case .running, .shuttingDown(graceful: true): + self.cacheNoMoreConnectionsAllowed = false + + let closedConnectionAction = self.connections.connectionClosed(connection.id) + + let connectionAction: ConnectionAction + if let newRequest = closedConnectionAction.newConnectionRequest { + connectionAction = .makeConnection(newRequest, closedConnectionAction.timersToCancel) + } else { + connectionAction = .cancelTimers(closedConnectionAction.timersToCancel) + } + + return .init(request: .none, connection: connectionAction) + + case .shuttingDown(graceful: false), .shutDown: + return .none() + } + } + + struct CleanupAction { + struct ConnectionToDrop { + var connection: Connection + var keepAliveTimer: Bool + var idleTimer: Bool + } + + var connections: [ConnectionToDrop] + var requests: [Request] + } + + mutating func triggerGracefulShutdown() -> Action { + fatalError("Unimplemented") + } + + mutating func triggerForceShutdown() -> Action { + switch self.poolState { + case .running: + self.poolState = .shuttingDown(graceful: false) + var shutdown = ConnectionAction.Shutdown() + self.connections.triggerForceShutdown(&shutdown) + + if shutdown.connections.isEmpty { + self.poolState = .shutDown + } + + return .init( + request: .failRequests(self.requestQueue.removeAll(), ConnectionPoolError.poolShutdown), + connection: .shutdown(shutdown) + ) + + case .shuttingDown: + return .none() + + case .shutDown: + return .init(request: .none, connection: .none) + } + } + + @inlinable + /*private*/ mutating func handleAvailableConnection( + index: Int, + availableContext: ConnectionGroup.AvailableConnectionContext + ) -> Action { + // this connection was busy before + let requests = self.requestQueue.pop(max: availableContext.info.availableStreams) + if !requests.isEmpty { + let leaseResult = self.connections.leaseConnection(at: index, streams: UInt16(requests.count)) + return .init( + request: .leaseConnection(requests, leaseResult.connection), + connection: .cancelTimers(.init(leaseResult.timersToCancel)) + ) + } + + switch availableContext.use { + case .persisted, .demand: + switch availableContext.info { + case .leased: + return .none() + + case .idle(_, let newIdle): + let timers = self.connections.parkConnection(at: index, hasBecomeIdle: newIdle).map(self.mapTimers) + + return .init( + request: .none, + connection: .scheduleTimers(timers) + ) + } + + case .overflow: + if let closeAction = self.connections.closeConnectionIfIdle(at: index) { + return .init( + request: .none, + connection: .closeConnection(closeAction.connection, closeAction.timersToCancel) + ) + } else { + return .none() + } + } + + } + + @inlinable + /* private */ func mapTimers(_ connectionTimer: ConnectionTimer) -> Timer { + switch connectionTimer.usecase { + case .backoff: + return Timer( + connectionTimer, + duration: Self.calculateBackoff(failedAttempt: self.failedConsecutiveConnectionAttempts) + ) + + case .keepAlive: + return Timer(connectionTimer, duration: self.configuration.keepAliveDuration!) + + case .idleTimeout: + return Timer(connectionTimer, duration: self.configuration.idleTimeoutDuration) + + } + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine { + /// Calculates the delay for the next connection attempt after the given number of failed `attempts`. + /// + /// Our backoff formula is: 100ms * 1.25^(attempts - 1) with 3% jitter that is capped of at 1 minute. + /// This means for: + /// - 1 failed attempt : 100ms + /// - 5 failed attempts: ~300ms + /// - 10 failed attempts: ~930ms + /// - 15 failed attempts: ~2.84s + /// - 20 failed attempts: ~8.67s + /// - 25 failed attempts: ~26s + /// - 29 failed attempts: ~60s (max out) + /// + /// - Parameter attempts: number of failed attempts in a row + /// - Returns: time to wait until trying to establishing a new connection + @usableFromInline + static func calculateBackoff(failedAttempt attempts: Int) -> Duration { + // Our backoff formula is: 100ms * 1.25^(attempts - 1) that is capped of at 1minute + // This means for: + // - 1 failed attempt : 100ms + // - 5 failed attempts: ~300ms + // - 10 failed attempts: ~930ms + // - 15 failed attempts: ~2.84s + // - 20 failed attempts: ~8.67s + // - 25 failed attempts: ~26s + // - 29 failed attempts: ~60s (max out) + + let start = Double(100_000_000) + let backoffNanosecondsDouble = start * pow(1.25, Double(attempts - 1)) + + // Cap to 60s _before_ we convert to Int64, to avoid trapping in the Int64 initializer. + let backoffNanoseconds = Int64(min(backoffNanosecondsDouble, Double(60_000_000_000))) + + let backoff = Duration.nanoseconds(backoffNanoseconds) + + // Calculate a 3% jitter range + let jitterRange = (backoffNanoseconds / 100) * 3 + // Pick a random element from the range +/- jitter range. + let jitter: Duration = .nanoseconds((-jitterRange...jitterRange).randomElement()!) + let jitteredBackoff = backoff + jitter + return jitteredBackoff + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.Action: Equatable where TimerCancellationToken: Equatable, Request: Equatable {} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.ConnectionAction: Equatable where TimerCancellationToken: Equatable { + @usableFromInline + static func ==(lhs: Self, rhs: Self) -> Bool { + switch (lhs, rhs) { + case (.scheduleTimers(let lhs), .scheduleTimers(let rhs)): + return lhs == rhs + case (.makeConnection(let lhsRequest, let lhsToken), .makeConnection(let rhsRequest, let rhsToken)): + return lhsRequest == rhsRequest && lhsToken == rhsToken + case (.runKeepAlive(let lhsConn, let lhsToken), .runKeepAlive(let rhsConn, let rhsToken)): + return lhsConn === rhsConn && lhsToken == rhsToken + case (.closeConnection(let lhsConn, let lhsTimers), .closeConnection(let rhsConn, let rhsTimers)): + return lhsConn === rhsConn && lhsTimers == rhsTimers + case (.shutdown(let lhs), .shutdown(let rhs)): + return lhs == rhs + case (.cancelTimers(let lhs), .cancelTimers(let rhs)): + return lhs == rhs + case (.none, .none), + (.cancelTimers([]), .none), (.none, .cancelTimers([])), + (.scheduleTimers([]), .none), (.none, .scheduleTimers([])): + return true + default: + return false + } + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.ConnectionAction.Shutdown: Equatable where TimerCancellationToken: Equatable { + @usableFromInline + static func ==(lhs: Self, rhs: Self) -> Bool { + Set(lhs.connections.lazy.map(\.id)) == Set(rhs.connections.lazy.map(\.id)) && lhs.timersToCancel == rhs.timersToCancel + } +} + + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PoolStateMachine.RequestAction: Equatable where Request: Equatable { + + @usableFromInline + static func ==(lhs: Self, rhs: Self) -> Bool { + switch (lhs, rhs) { + case (.leaseConnection(let lhsRequests, let lhsConn), .leaseConnection(let rhsRequests, let rhsConn)): + guard lhsRequests.count == rhsRequests.count else { return false } + var lhsIterator = lhsRequests.makeIterator() + var rhsIterator = rhsRequests.makeIterator() + while let lhsNext = lhsIterator.next(), let rhsNext = rhsIterator.next() { + guard lhsNext.id == rhsNext.id else { return false } + } + return lhsConn === rhsConn + + case (.failRequest(let lhsRequest, let lhsError), .failRequest(let rhsRequest, let rhsError)): + return lhsRequest.id == rhsRequest.id && lhsError == rhsError + + case (.failRequests(let lhsRequests, let lhsError), .failRequests(let rhsRequests, let rhsError)): + return Set(lhsRequests.lazy.map(\.id)) == Set(rhsRequests.lazy.map(\.id)) && lhsError == rhsError + + case (.none, .none): + return true + + default: + return false + } + } +} diff --git a/Sources/ConnectionPoolModule/TinyFastSequence.swift b/Sources/ConnectionPoolModule/TinyFastSequence.swift new file mode 100644 index 00000000..dff8a30b --- /dev/null +++ b/Sources/ConnectionPoolModule/TinyFastSequence.swift @@ -0,0 +1,205 @@ +/// A `Sequence` that does not heap allocate, if it only carries a single element +@usableFromInline +struct TinyFastSequence: Sequence { + @usableFromInline + enum Base { + case none(reserveCapacity: Int) + case one(Element, reserveCapacity: Int) + case two(Element, Element, reserveCapacity: Int) + case n([Element]) + } + + @usableFromInline + private(set) var base: Base + + @inlinable + init() { + self.base = .none(reserveCapacity: 0) + } + + @inlinable + init(element: Element) { + self.base = .one(element, reserveCapacity: 1) + } + + @inlinable + init(_ collection: some Collection) { + switch collection.count { + case 0: + self.base = .none(reserveCapacity: 0) + case 1: + self.base = .one(collection.first!, reserveCapacity: 0) + default: + if let collection = collection as? Array { + self.base = .n(collection) + } else { + self.base = .n(Array(collection)) + } + } + } + + @inlinable + init(_ max2Sequence: Max2Sequence) { + switch max2Sequence.count { + case 0: + self.base = .none(reserveCapacity: 0) + case 1: + self.base = .one(max2Sequence.first!, reserveCapacity: 0) + case 2: + self.base = .n(Array(max2Sequence)) + default: + fatalError() + } + } + + @usableFromInline + var count: Int { + switch self.base { + case .none: + return 0 + case .one: + return 1 + case .two: + return 2 + case .n(let array): + return array.count + } + } + + @inlinable + var first: Element? { + switch self.base { + case .none: + return nil + case .one(let element, _): + return element + case .two(let first, _, _): + return first + case .n(let array): + return array.first + } + } + + @usableFromInline + var isEmpty: Bool { + switch self.base { + case .none: + return true + case .one, .two, .n: + return false + } + } + + @inlinable + mutating func reserveCapacity(_ minimumCapacity: Int) { + switch self.base { + case .none(let reservedCapacity): + self.base = .none(reserveCapacity: Swift.max(reservedCapacity, minimumCapacity)) + case .one(let element, let reservedCapacity): + self.base = .one(element, reserveCapacity: Swift.max(reservedCapacity, minimumCapacity)) + case .two(let first, let second, let reservedCapacity): + self.base = .two(first, second, reserveCapacity: Swift.max(reservedCapacity, minimumCapacity)) + case .n(var array): + self.base = .none(reserveCapacity: 0) // prevent CoW + array.reserveCapacity(minimumCapacity) + self.base = .n(array) + } + } + + @inlinable + mutating func append(_ element: Element) { + switch self.base { + case .none(let reserveCapacity): + self.base = .one(element, reserveCapacity: reserveCapacity) + case .one(let first, let reserveCapacity): + self.base = .two(first, element, reserveCapacity: reserveCapacity) + + case .two(let first, let second, let reserveCapacity): + var new = [Element]() + new.reserveCapacity(Swift.max(4, reserveCapacity)) + new.append(first) + new.append(second) + new.append(element) + self.base = .n(new) + + case .n(var existing): + self.base = .none(reserveCapacity: 0) // prevent CoW + existing.append(element) + self.base = .n(existing) + } + } + + @inlinable + func makeIterator() -> Iterator { + Iterator(self) + } + + @usableFromInline + struct Iterator: IteratorProtocol { + @usableFromInline private(set) var index: Int = 0 + @usableFromInline private(set) var backing: TinyFastSequence + + @inlinable + init(_ backing: TinyFastSequence) { + self.backing = backing + } + + @inlinable + mutating func next() -> Element? { + switch self.backing.base { + case .none: + return nil + case .one(let element, _): + if self.index == 0 { + self.index += 1 + return element + } + return nil + + case .two(let first, let second, _): + defer { self.index += 1 } + switch self.index { + case 0: + return first + case 1: + return second + default: + return nil + } + + case .n(let array): + if self.index < array.endIndex { + defer { self.index += 1} + return array[self.index] + } + return nil + } + } + } +} + +extension TinyFastSequence: Equatable where Element: Equatable {} +extension TinyFastSequence.Base: Equatable where Element: Equatable {} + +extension TinyFastSequence: Hashable where Element: Hashable {} +extension TinyFastSequence.Base: Hashable where Element: Hashable {} + +extension TinyFastSequence: Sendable where Element: Sendable {} +extension TinyFastSequence.Base: Sendable where Element: Sendable {} + +extension TinyFastSequence: ExpressibleByArrayLiteral { + @inlinable + init(arrayLiteral elements: Element...) { + var iterator = elements.makeIterator() + switch elements.count { + case 0: + self.base = .none(reserveCapacity: 0) + case 1: + self.base = .one(iterator.next()!, reserveCapacity: 0) + case 2: + self.base = .two(iterator.next()!, iterator.next()!, reserveCapacity: 0) + default: + self.base = .n(elements) + } + } +} diff --git a/Sources/ConnectionPoolTestUtils/MockClock.swift b/Sources/ConnectionPoolTestUtils/MockClock.swift new file mode 100644 index 00000000..34bf17e3 --- /dev/null +++ b/Sources/ConnectionPoolTestUtils/MockClock.swift @@ -0,0 +1,176 @@ +import _ConnectionPoolModule +import Atomics +import DequeModule +import NIOConcurrencyHelpers + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public final class MockClock: Clock { + public struct Instant: InstantProtocol, Comparable { + public typealias Duration = Swift.Duration + + public func advanced(by duration: Self.Duration) -> Self { + .init(self.base + duration) + } + + public func duration(to other: Self) -> Self.Duration { + self.base - other.base + } + + private var base: Swift.Duration + + public init(_ base: Duration) { + self.base = base + } + + public static func < (lhs: Self, rhs: Self) -> Bool { + lhs.base < rhs.base + } + + public static func == (lhs: Self, rhs: Self) -> Bool { + lhs.base == rhs.base + } + } + + private struct State: Sendable { + var now: Instant + + var sleepersHeap: Array + + var waiters: Deque + var nextDeadlines: Deque + + init() { + self.now = .init(.seconds(0)) + self.sleepersHeap = Array() + self.waiters = Deque() + self.nextDeadlines = Deque() + } + } + + private struct Waiter { + var continuation: CheckedContinuation + } + + private struct Sleeper { + var id: Int + + var deadline: Instant + + var continuation: CheckedContinuation + } + + public typealias Duration = Swift.Duration + + public var minimumResolution: Duration { .nanoseconds(1) } + + public var now: Instant { self.stateBox.withLockedValue { $0.now } } + + private let stateBox = NIOLockedValueBox(State()) + private let waiterIDGenerator = ManagedAtomic(0) + + public init() {} + + public func sleep(until deadline: Instant, tolerance: Duration?) async throws { + let waiterID = self.waiterIDGenerator.loadThenWrappingIncrement(ordering: .relaxed) + + return try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + enum SleepAction { + case none + case resume + case cancel + } + + let action = self.stateBox.withLockedValue { state -> (SleepAction, Waiter?) in + let waiter: Waiter? + if let next = state.waiters.popFirst() { + waiter = next + } else { + state.nextDeadlines.append(deadline) + waiter = nil + } + + if Task.isCancelled { + return (.cancel, waiter) + } + + if state.now >= deadline { + return (.resume, waiter) + } + + let newSleeper = Sleeper(id: waiterID, deadline: deadline, continuation: continuation) + + if let index = state.sleepersHeap.lastIndex(where: { $0.deadline < deadline }) { + state.sleepersHeap.insert(newSleeper, at: index + 1) + } else if let first = state.sleepersHeap.first, first.deadline > deadline { + state.sleepersHeap.insert(newSleeper, at: 0) + } else { + state.sleepersHeap.append(newSleeper) + } + + return (.none, waiter) + } + + switch action.0 { + case .cancel: + continuation.resume(throwing: CancellationError()) + case .resume: + continuation.resume() + case .none: + break + } + + action.1?.continuation.resume(returning: deadline) + } + } onCancel: { + let continuation = self.stateBox.withLockedValue { state -> CheckedContinuation? in + if let index = state.sleepersHeap.firstIndex(where: { $0.id == waiterID }) { + return state.sleepersHeap.remove(at: index).continuation + } + return nil + } + continuation?.resume(throwing: CancellationError()) + } + } + + @discardableResult + public func nextTimerScheduled() async -> Instant { + await withCheckedContinuation { (continuation: CheckedContinuation) in + let instant = self.stateBox.withLockedValue { state -> Instant? in + if let scheduled = state.nextDeadlines.popFirst() { + return scheduled + } else { + let waiter = Waiter(continuation: continuation) + state.waiters.append(waiter) + return nil + } + } + + if let instant { + continuation.resume(returning: instant) + } + } + } + + public func advance(to deadline: Instant) { + let waiters = self.stateBox.withLockedValue { state -> ArraySlice in + precondition(deadline > state.now, "Time can only move forward") + state.now = deadline + + if let newFirstIndex = state.sleepersHeap.firstIndex(where: { $0.deadline > deadline }) { + defer { state.sleepersHeap.removeFirst(newFirstIndex) } + return state.sleepersHeap[0..], [@Sendable ((any Error)?) -> ()]) + case closing([@Sendable ((any Error)?) -> ()]) + case closed + } + + private let lock: NIOLockedValueBox = NIOLockedValueBox(.running([], [])) + + public init(id: Int) { + self.id = id + } + + public var signalToClose: Void { + get async throws { + try await withCheckedThrowingContinuation { continuation in + let runRightAway = self.lock.withLockedValue { state -> Bool in + switch state { + case .running(var continuations, let callbacks): + continuations.append(continuation) + state = .running(continuations, callbacks) + return false + + case .closing, .closed: + return true + } + } + + if runRightAway { + continuation.resume() + } + } + } + } + + public func onClose(_ closure: @escaping @Sendable ((any Error)?) -> ()) { + let enqueued = self.lock.withLockedValue { state -> Bool in + switch state { + case .closed: + return false + + case .running(let continuations, var callbacks): + callbacks.append(closure) + state = .running(continuations, callbacks) + return true + + case .closing(var callbacks): + callbacks.append(closure) + state = .closing(callbacks) + return true + } + } + + if !enqueued { + closure(nil) + } + } + + public func close() { + let continuations = self.lock.withLockedValue { state -> [CheckedContinuation] in + switch state { + case .running(let continuations, let callbacks): + state = .closing(callbacks) + return continuations + + case .closing, .closed: + return [] + } + } + + for continuation in continuations { + continuation.resume() + } + } + + public func closeIfClosing() { + let callbacks = self.lock.withLockedValue { state -> [@Sendable ((any Error)?) -> ()] in + switch state { + case .running, .closed: + return [] + + case .closing(let callbacks): + state = .closed + return callbacks + } + } + + for callback in callbacks { + callback(nil) + } + } +} + +extension MockConnection: CustomStringConvertible { + public var description: String { + let state = self.lock.withLockedValue { $0 } + return "MockConnection(id: \(self.id), state: \(state))" + } +} diff --git a/Sources/ConnectionPoolTestUtils/MockConnectionFactory.swift b/Sources/ConnectionPoolTestUtils/MockConnectionFactory.swift new file mode 100644 index 00000000..936b47cc --- /dev/null +++ b/Sources/ConnectionPoolTestUtils/MockConnectionFactory.swift @@ -0,0 +1,108 @@ +import _ConnectionPoolModule +import DequeModule +import NIOConcurrencyHelpers + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public final class MockConnectionFactory: Sendable where Clock.Duration == Duration { + public typealias ConnectionIDGenerator = _ConnectionPoolModule.ConnectionIDGenerator + public typealias Request = ConnectionRequest + public typealias KeepAliveBehavior = MockPingPongBehavior + public typealias MetricsDelegate = NoOpConnectionPoolMetrics + public typealias ConnectionID = Int + public typealias Connection = MockConnection + + let stateBox = NIOLockedValueBox(State()) + + struct State { + var attempts = Deque<(ConnectionID, CheckedContinuation<(MockConnection, UInt16), any Error>)>() + + var waiter = Deque), Never>>() + + var runningConnections = [ConnectionID: Connection]() + } + + let autoMaxStreams: UInt16? + + public init(autoMaxStreams: UInt16? = nil) { + self.autoMaxStreams = autoMaxStreams + } + + public var pendingConnectionAttemptsCount: Int { + self.stateBox.withLockedValue { $0.attempts.count } + } + + public var runningConnections: [Connection] { + self.stateBox.withLockedValue { Array($0.runningConnections.values) } + } + + public func makeConnection( + id: Int, + for pool: ConnectionPool, NoOpConnectionPoolMetrics, Clock> + ) async throws -> ConnectionAndMetadata { + if let autoMaxStreams = self.autoMaxStreams { + let connection = MockConnection(id: id) + Task { + try? await connection.signalToClose + connection.closeIfClosing() + } + return .init(connection: connection, maximalStreamsOnConnection: autoMaxStreams) + } + + // we currently don't support cancellation when creating a connection + let result = try await withCheckedThrowingContinuation { (checkedContinuation: CheckedContinuation<(MockConnection, UInt16), any Error>) in + let waiter = self.stateBox.withLockedValue { state -> (CheckedContinuation<(ConnectionID, CheckedContinuation<(MockConnection, UInt16), any Error>), Never>)? in + if let waiter = state.waiter.popFirst() { + return waiter + } else { + state.attempts.append((id, checkedContinuation)) + return nil + } + } + + if let waiter { + waiter.resume(returning: (id, checkedContinuation)) + } + } + + return .init(connection: result.0, maximalStreamsOnConnection: result.1) + } + + @discardableResult + public func nextConnectAttempt(_ closure: (ConnectionID) async throws -> UInt16) async rethrows -> Connection { + let (connectionID, continuation) = await withCheckedContinuation { (continuation: CheckedContinuation<(ConnectionID, CheckedContinuation<(MockConnection, UInt16), any Error>), Never>) in + let attempt = self.stateBox.withLockedValue { state -> (ConnectionID, CheckedContinuation<(MockConnection, UInt16), any Error>)? in + if let attempt = state.attempts.popFirst() { + return attempt + } else { + state.waiter.append(continuation) + return nil + } + } + + if let attempt { + continuation.resume(returning: attempt) + } + } + + do { + let streamCount = try await closure(connectionID) + let connection = MockConnection(id: connectionID) + + connection.onClose { _ in + self.stateBox.withLockedValue { state in + _ = state.runningConnections.removeValue(forKey: connectionID) + } + } + + self.stateBox.withLockedValue { state in + _ = state.runningConnections[connectionID] = connection + } + + continuation.resume(returning: (connection, streamCount)) + return connection + } catch { + continuation.resume(throwing: error) + throw error + } + } +} diff --git a/Sources/ConnectionPoolTestUtils/MockPingPongBehaviour.swift b/Sources/ConnectionPoolTestUtils/MockPingPongBehaviour.swift new file mode 100644 index 00000000..de1a7275 --- /dev/null +++ b/Sources/ConnectionPoolTestUtils/MockPingPongBehaviour.swift @@ -0,0 +1,70 @@ +import _ConnectionPoolModule +import DequeModule +import NIOConcurrencyHelpers + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public final class MockPingPongBehavior: ConnectionKeepAliveBehavior { + public let keepAliveFrequency: Duration? + + let stateBox = NIOLockedValueBox(State()) + + struct State { + var runs = Deque<(Connection, CheckedContinuation)>() + + var waiter = Deque), Never>>() + } + + public init(keepAliveFrequency: Duration?, connectionType: Connection.Type) { + self.keepAliveFrequency = keepAliveFrequency + } + + public func runKeepAlive(for connection: Connection) async throws { + precondition(self.keepAliveFrequency != nil) + + // we currently don't support cancellation when creating a connection + let success = try await withCheckedThrowingContinuation { (checkedContinuation: CheckedContinuation) -> () in + let waiter = self.stateBox.withLockedValue { state -> (CheckedContinuation<(Connection, CheckedContinuation), Never>)? in + if let waiter = state.waiter.popFirst() { + return waiter + } else { + state.runs.append((connection, checkedContinuation)) + return nil + } + } + + if let waiter { + waiter.resume(returning: (connection, checkedContinuation)) + } + } + + precondition(success) + } + + @discardableResult + public func nextKeepAlive(_ closure: (Connection) async throws -> Bool) async rethrows -> Connection { + let (connection, continuation) = await withCheckedContinuation { (continuation: CheckedContinuation<(Connection, CheckedContinuation), Never>) in + let run = self.stateBox.withLockedValue { state -> (Connection, CheckedContinuation)? in + if let run = state.runs.popFirst() { + return run + } else { + state.waiter.append(continuation) + return nil + } + } + + if let run { + continuation.resume(returning: run) + } + } + + do { + let success = try await closure(connection) + + continuation.resume(returning: success) + return connection + } catch { + continuation.resume(throwing: error) + throw error + } + } +} diff --git a/Sources/ConnectionPoolTestUtils/MockRequest.swift b/Sources/ConnectionPoolTestUtils/MockRequest.swift new file mode 100644 index 00000000..5e4e2fc0 --- /dev/null +++ b/Sources/ConnectionPoolTestUtils/MockRequest.swift @@ -0,0 +1,29 @@ +import _ConnectionPoolModule + +public final class MockRequest: ConnectionRequestProtocol, Hashable, Sendable { + public typealias Connection = MockConnection + + public struct ID: Hashable, Sendable { + var objectID: ObjectIdentifier + + init(_ request: MockRequest) { + self.objectID = ObjectIdentifier(request) + } + } + + public init() {} + + public var id: ID { ID(self) } + + public static func ==(lhs: MockRequest, rhs: MockRequest) -> Bool { + lhs.id == rhs.id + } + + public func hash(into hasher: inout Hasher) { + hasher.combine(self.id) + } + + public func complete(with: Result) { + + } +} diff --git a/Sources/PostgresNIO/Connection/PostgresConnection+Configuration.swift b/Sources/PostgresNIO/Connection/PostgresConnection+Configuration.swift index 54eefc90..b260723a 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection+Configuration.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection+Configuration.swift @@ -4,16 +4,16 @@ import NIOSSL extension PostgresConnection { /// A configuration object for a connection - public struct Configuration { - + public struct Configuration: Sendable { + // MARK: - TLS /// The possible modes of operation for TLS encapsulation of a connection. - public struct TLS { + public struct TLS: Sendable { // MARK: Initializers /// Do not try to create a TLS connection to the server. - public static var disable: Self = .init(base: .disable) + public static var disable: Self { .init(base: .disable) } /// Try to create a TLS connection to the server. If the server supports TLS, create a TLS connection. /// If the server does not support TLS, create an insecure connection. @@ -63,7 +63,7 @@ extension PostgresConnection { // MARK: - Connection options /// Describes options affecting how the underlying connection is made. - public struct Options { + public struct Options: Sendable { /// A timeout for connection attempts. Defaults to ten seconds. /// /// Ignored when using a preexisting communcation channel. (See @@ -85,7 +85,11 @@ extension PostgresConnection { /// This property is provided for compatibility with Amazon RDS Proxy, which requires it to be `false`. /// If you are not using Amazon RDS Proxy, you should leave this set to `true` (the default). public var requireBackendKeyData: Bool - + + /// Additional parameters to send to the server on startup. The name value pairs are added to the initial + /// startup message that the client sends to the server. + public var additionalStartupParameters: [(String, String)] + /// Create an options structure with default values. /// /// Most users should not need to adjust the defaults. @@ -93,6 +97,7 @@ extension PostgresConnection { self.connectTimeout = .seconds(10) self.tlsServerName = nil self.requireBackendKeyData = true + self.additionalStartupParameters = [] } } @@ -187,9 +192,22 @@ extension PostgresConnection { /// - Parameters: /// - channel: The `NIOCore/Channel` to use. The channel must already be active and connected to an /// endpoint (i.e. `NIOCore/Channel/isActive` must be `true`). - /// - tls: The TLS mode to use. Defaults to ``TLS-swift.struct/disable``. + /// - tls: The TLS mode to use. + public init(establishedChannel channel: Channel, tls: PostgresConnection.Configuration.TLS, username: String, password: String?, database: String?) { + self.init(endpointInfo: .configureChannel(channel), tls: tls, username: username, password: password, database: database) + } + + /// Create a configuration for establishing a connection to a Postgres server over a preestablished + /// `NIOCore/Channel`. + /// + /// This is provided for calling code which wants to manage the underlying connection transport on its + /// own, such as when tunneling a connection through SSH. + /// + /// - Parameters: + /// - channel: The `NIOCore/Channel` to use. The channel must already be active and connected to an + /// endpoint (i.e. `NIOCore/Channel/isActive` must be `true`). public init(establishedChannel channel: Channel, username: String, password: String?, database: String?) { - self.init(endpointInfo: .configureChannel(channel), tls: .disable, username: username, password: password, database: database) + self.init(establishedChannel: channel, tls: .disable, username: username, password: password, database: database) } // MARK: - Implementation details @@ -219,7 +237,7 @@ extension PostgresConnection { /// the deprecated configuration. /// /// TODO: Drop with next major release - struct InternalConfiguration { + struct InternalConfiguration: Sendable { enum Connection { case unresolvedTCP(host: String, port: Int) case unresolvedUDS(path: String) diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift index c24041c9..e267d8f9 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection.swift @@ -38,21 +38,13 @@ public final class PostgresConnection: @unchecked Sendable { } } - /// A dictionary to store notification callbacks in - /// - /// Those are used when `PostgresConnection.addListener` is invoked. This only lives here since properties - /// can not be added in extensions. All relevant code lives in `PostgresConnection+Notifications` - var notificationListeners: [String: [(PostgresListenContext, (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void)]] = [:] { - willSet { - self.channel.eventLoop.preconditionInEventLoop() - } - } + private let internalListenID = ManagedAtomic(0) public var isClosed: Bool { return !self.channel.isActive } - let id: ID + public let id: ID private var _logger: Logger @@ -68,18 +60,18 @@ public final class PostgresConnection: @unchecked Sendable { func start(configuration: InternalConfiguration) -> EventLoopFuture { // 1. configure handlers - let configureSSLCallback: ((Channel) throws -> ())? + let configureSSLCallback: ((Channel, PostgresChannelHandler) throws -> ())? switch configuration.tls.base { case .prefer(let context), .require(let context): - configureSSLCallback = { channel in + configureSSLCallback = { channel, postgresChannelHandler in channel.eventLoop.assertInEventLoop() let sslHandler = try NIOSSLClientHandler( context: context, serverHostname: configuration.serverNameForTLS ) - try channel.pipeline.syncOperations.addHandler(sslHandler, position: .first) + try channel.pipeline.syncOperations.addHandler(sslHandler, position: .before(postgresChannelHandler)) } case .disable: configureSSLCallback = nil @@ -87,10 +79,10 @@ public final class PostgresConnection: @unchecked Sendable { let channelHandler = PostgresChannelHandler( configuration: configuration, + eventLoop: channel.eventLoop, logger: logger, configureSSLCallback: configureSSLCallback ) - channelHandler.notificationDelegate = self let eventHandler = PSQLEventsHandler(logger: logger) @@ -152,8 +144,9 @@ public final class PostgresConnection: @unchecked Sendable { on eventLoop: EventLoop ) -> EventLoopFuture { - var logger = logger - logger[postgresMetadataKey: .connectionID] = "\(connectionID)" + var mlogger = logger + mlogger[postgresMetadataKey: .connectionID] = "\(connectionID)" + let logger = mlogger // Here we dispatch to the `eventLoop` first before we setup the EventLoopFuture chain, to // ensure all `flatMap`s are executed on the EventLoop (this means the enqueuing of the @@ -164,14 +157,16 @@ public final class PostgresConnection: @unchecked Sendable { // thread and the EventLoop. return eventLoop.flatSubmit { () -> EventLoopFuture in let connectFuture: EventLoopFuture - let bootstrap = self.makeBootstrap(on: eventLoop, configuration: configuration) switch configuration.connection { case .resolved(let address): + let bootstrap = self.makeBootstrap(on: eventLoop, configuration: configuration) connectFuture = bootstrap.connect(to: address) case .unresolvedTCP(let host, let port): + let bootstrap = self.makeBootstrap(on: eventLoop, configuration: configuration) connectFuture = bootstrap.connect(host: host, port: port) case .unresolvedUDS(let path): + let bootstrap = self.makeBootstrap(on: eventLoop, configuration: configuration) connectFuture = bootstrap.connect(unixDomainSocketPath: path) case .bootstrapped(let channel): guard channel.isActive else { @@ -224,9 +219,10 @@ public final class PostgresConnection: @unchecked Sendable { let context = ExtendedQueryContext( query: query, logger: logger, - promise: promise) + promise: promise + ) - self.channel.write(PSQLTask.extendedQuery(context), promise: nil) + self.channel.write(HandlerTask.extendedQuery(context), promise: nil) return promise.futureResult } @@ -235,13 +231,15 @@ public final class PostgresConnection: @unchecked Sendable { func prepareStatement(_ query: String, with name: String, logger: Logger) -> EventLoopFuture { let promise = self.channel.eventLoop.makePromise(of: RowDescription?.self) - let context = PrepareStatementContext( + let context = ExtendedQueryContext( name: name, query: query, + bindingDataTypes: [], logger: logger, - promise: promise) + promise: promise + ) - self.channel.write(PSQLTask.preparedStatement(context), promise: nil) + self.channel.write(HandlerTask.extendedQuery(context), promise: nil) return promise.futureResult.map { rowDescription in PSQLPreparedStatement(name: name, query: query, connection: self, rowDescription: rowDescription) } @@ -257,7 +255,7 @@ public final class PostgresConnection: @unchecked Sendable { logger: logger, promise: promise) - self.channel.write(PSQLTask.extendedQuery(context), promise: nil) + self.channel.write(HandlerTask.extendedQuery(context), promise: nil) return promise.futureResult } @@ -265,7 +263,7 @@ public final class PostgresConnection: @unchecked Sendable { let promise = self.channel.eventLoop.makePromise(of: Void.self) let context = CloseCommandContext(target: target, logger: logger, promise: promise) - self.channel.write(PSQLTask.closeCommand(context), promise: nil) + self.channel.write(HandlerTask.closeCommand(context), promise: nil) return promise.futureResult } @@ -364,13 +362,13 @@ extension PostgresConnection { /// Creates a new connection to a Postgres server. /// /// - Parameters: - /// - eventLoop: The `EventLoop` the request shall be created on + /// - eventLoop: The `EventLoop` the connection shall be created on. /// - configuration: A ``Configuration`` that shall be used for the connection /// - connectionID: An `Int` id, used for metadata logging /// - logger: A logger to log background events into /// - Returns: An established ``PostgresConnection`` asynchronously that can be used to run queries. public static func connect( - on eventLoop: EventLoop, + on eventLoop: EventLoop = PostgresConnection.defaultEventLoopGroup.any(), configuration: PostgresConnection.Configuration, id connectionID: ID, logger: Logger @@ -388,6 +386,17 @@ extension PostgresConnection { try await self.close().get() } + /// Closes the connection to the server, _after all queries_ that have been created on this connection have been run. + public func closeGracefully() async throws { + try await withTaskCancellationHandler { () async throws -> () in + let promise = self.eventLoop.makePromise(of: Void.self) + self.channel.triggerUserOutboundEvent(PSQLOutgoingEvent.gracefulShutdown, promise: promise) + return try await promise.futureResult.get() + } onCancel: { + self.close() + } + } + /// Run a query on the Postgres server the connection is connected to. /// /// - Parameters: @@ -417,7 +426,7 @@ extension PostgresConnection { promise: promise ) - self.channel.write(PSQLTask.extendedQuery(context), promise: nil) + self.channel.write(HandlerTask.extendedQuery(context), promise: nil) do { return try await promise.futureResult.map({ $0.asyncSequence() }).get() @@ -428,6 +437,203 @@ extension PostgresConnection { throw error // rethrow with more metadata } } + + /// Start listening for a channel + public func listen(_ channel: String) async throws -> PostgresNotificationSequence { + let id = self.internalListenID.loadThenWrappingIncrement(ordering: .relaxed) + + return try await withTaskCancellationHandler { + try Task.checkCancellation() + + return try await withCheckedThrowingContinuation { continuation in + let listener = NotificationListener( + channel: channel, + id: id, + eventLoop: self.eventLoop, + checkedContinuation: continuation + ) + + let task = HandlerTask.startListening(listener) + + self.channel.write(task, promise: nil) + } + } onCancel: { + let task = HandlerTask.cancelListening(channel, id) + self.channel.write(task, promise: nil) + } + } + + /// Execute a prepared statement, taking care of the preparation when necessary + public func execute( + _ preparedStatement: Statement, + logger: Logger, + file: String = #fileID, + line: Int = #line + ) async throws -> AsyncThrowingMapSequence where Row == Statement.Row { + let bindings = try preparedStatement.makeBindings() + let promise = self.channel.eventLoop.makePromise(of: PSQLRowStream.self) + let task = HandlerTask.executePreparedStatement(.init( + name: Statement.name, + sql: Statement.sql, + bindings: bindings, + bindingDataTypes: Statement.bindingDataTypes, + logger: logger, + promise: promise + )) + self.channel.write(task, promise: nil) + do { + return try await promise.futureResult + .map { $0.asyncSequence() } + .get() + .map { try preparedStatement.decodeRow($0) } + } catch var error as PSQLError { + error.file = file + error.line = line + error.query = .init( + unsafeSQL: Statement.sql, + binds: bindings + ) + throw error // rethrow with more metadata + } + } + + /// Execute a prepared statement, taking care of the preparation when necessary + @_disfavoredOverload + public func execute( + _ preparedStatement: Statement, + logger: Logger, + file: String = #fileID, + line: Int = #line + ) async throws -> String where Statement.Row == () { + let bindings = try preparedStatement.makeBindings() + let promise = self.channel.eventLoop.makePromise(of: PSQLRowStream.self) + let task = HandlerTask.executePreparedStatement(.init( + name: Statement.name, + sql: Statement.sql, + bindings: bindings, + bindingDataTypes: Statement.bindingDataTypes, + logger: logger, + promise: promise + )) + self.channel.write(task, promise: nil) + do { + return try await promise.futureResult + .map { $0.commandTag } + .get() + } catch var error as PSQLError { + error.file = file + error.line = line + error.query = .init( + unsafeSQL: Statement.sql, + binds: bindings + ) + throw error // rethrow with more metadata + } + } + + #if compiler(>=6.0) + /// Puts the connection into an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function starts a transaction by running a `BEGIN` query on the connection against the database. It then + /// lends the connection to the user provided closure. The user can then modify the database as they wish. If the user + /// provided closure returns successfully, the function will attempt to commit the changes by running a `COMMIT` + /// query against the database. If the user provided closure throws an error, the function will attempt to rollback the + /// changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + isolation: isolated (any Actor)? = #isolation, + // DO NOT FIX THE WHITESPACE IN THE NEXT LINE UNTIL 5.10 IS UNSUPPORTED + // https://github.com/swiftlang/swift/issues/79285 + _ process: (PostgresConnection) async throws -> sending Result) async throws -> sending Result { + do { + try await self.query("BEGIN;", logger: logger) + } catch { + throw PostgresTransactionError(file: file, line: line, beginError: error) + } + + var closureHasFinished: Bool = false + do { + let value = try await process(self) + closureHasFinished = true + try await self.query("COMMIT;", logger: logger) + return value + } catch { + var transactionError = PostgresTransactionError(file: file, line: line) + if !closureHasFinished { + transactionError.closureError = error + do { + try await self.query("ROLLBACK;", logger: logger) + } catch { + transactionError.rollbackError = error + } + } else { + transactionError.commitError = error + } + + throw transactionError + } + } + #else + /// Puts the connection into an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function starts a transaction by running a `BEGIN` query on the connection against the database. It then + /// lends the connection to the user provided closure. The user can then modify the database as they wish. If the user + /// provided closure returns successfully, the function will attempt to commit the changes by running a `COMMIT` + /// query against the database. If the user provided closure throws an error, the function will attempt to rollback the + /// changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + _ process: (PostgresConnection) async throws -> Result + ) async throws -> Result { + do { + try await self.query("BEGIN;", logger: logger) + } catch { + throw PostgresTransactionError(file: file, line: line, beginError: error) + } + + var closureHasFinished: Bool = false + do { + let value = try await process(self) + closureHasFinished = true + try await self.query("COMMIT;", logger: logger) + return value + } catch { + var transactionError = PostgresTransactionError(file: file, line: line) + if !closureHasFinished { + transactionError.closureError = error + do { + try await self.query("ROLLBACK;", logger: logger) + } catch { + transactionError.rollbackError = error + } + } else { + transactionError.commitError = error + } + + throw transactionError + } + } + #endif } // MARK: EventLoopFuture interface @@ -469,12 +675,13 @@ extension PostgresConnection { /// - line: The line, the query was started in. Used for better error reporting. /// - onRow: A closure that is invoked for every row. /// - Returns: An EventLoopFuture, that allows access to the future ``PostgresQueryMetadata``. + @preconcurrency public func query( _ query: PostgresQuery, logger: Logger, file: String = #fileID, line: Int = #line, - _ onRow: @escaping (PostgresRow) throws -> () + _ onRow: @escaping @Sendable (PostgresRow) throws -> () ) -> EventLoopFuture { self.queryStream(query, logger: logger).flatMap { rowStream in rowStream.onRow(onRow).flatMapThrowing { () -> PostgresQueryMetadata in @@ -540,6 +747,7 @@ extension PostgresConnection: PostgresDatabase { } } + @preconcurrency public func withConnection(_ closure: (PostgresConnection) -> EventLoopFuture) -> EventLoopFuture { closure(self) } @@ -547,11 +755,11 @@ extension PostgresConnection: PostgresDatabase { internal enum PostgresCommands: PostgresRequest { case query(PostgresQuery, - onMetadata: (PostgresQueryMetadata) -> () = { _ in }, - onRow: (PostgresRow) throws -> ()) - case queryAll(PostgresQuery, onResult: (PostgresQueryResult) -> ()) + onMetadata: @Sendable (PostgresQueryMetadata) -> () = { _ in }, + onRow: @Sendable (PostgresRow) throws -> ()) + case queryAll(PostgresQuery, onResult: @Sendable (PostgresQueryResult) -> ()) case prepareQuery(request: PrepareQueryRequest) - case executePreparedStatement(query: PreparedQuery, binds: [PostgresData], onRow: (PostgresRow) throws -> ()) + case executePreparedStatement(query: PreparedQuery, binds: [PostgresData], onRow: @Sendable (PostgresRow) throws -> ()) func respond(to message: PostgresMessage) throws -> [PostgresMessage]? { fatalError("This function must not be called") @@ -569,73 +777,58 @@ internal enum PostgresCommands: PostgresRequest { // MARK: Notifications /// Context for receiving NotificationResponse messages on a connection, used for PostgreSQL's `LISTEN`/`NOTIFY` support. -public final class PostgresListenContext { - var stopper: (() -> Void)? +public final class PostgresListenContext: Sendable { + private let promise: EventLoopPromise + + var future: EventLoopFuture { + self.promise.futureResult + } + + init(promise: EventLoopPromise) { + self.promise = promise + } + + func cancel() { + self.promise.succeed() + } /// Detach this listener so it no longer receives notifications. Other listeners, including those for the same channel, are unaffected. `UNLISTEN` is not sent; you are responsible for issuing an `UNLISTEN` query yourself if it is appropriate for your application. public func stop() { - stopper?() - stopper = nil + self.promise.succeed() } } extension PostgresConnection { /// Add a handler for NotificationResponse messages on a certain channel. This is used in conjunction with PostgreSQL's `LISTEN`/`NOTIFY` support: to listen on a channel, you add a listener using this method to handle the NotificationResponse messages, then issue a `LISTEN` query to instruct PostgreSQL to begin sending NotificationResponse messages. @discardableResult - public func addListener(channel: String, handler notificationHandler: @escaping (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void) -> PostgresListenContext { - - let listenContext = PostgresListenContext() - - self.channel.pipeline.handler(type: PostgresChannelHandler.self).whenSuccess { handler in - if self.notificationListeners[channel] != nil { - self.notificationListeners[channel]!.append((listenContext, notificationHandler)) - } - else { - self.notificationListeners[channel] = [(listenContext, notificationHandler)] - } - } - - listenContext.stopper = { [weak self, weak listenContext] in - // self is weak, since the connection can long be gone, when the listeners stop is - // triggered. listenContext must be weak to prevent a retain cycle + @preconcurrency + public func addListener( + channel: String, + handler notificationHandler: @Sendable @escaping (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void + ) -> PostgresListenContext { + let listenContext = PostgresListenContext(promise: self.eventLoop.makePromise(of: Void.self)) + let id = self.internalListenID.loadThenWrappingIncrement(ordering: .relaxed) + + let listener = NotificationListener( + channel: channel, + id: id, + eventLoop: self.eventLoop, + context: listenContext, + closure: notificationHandler + ) - self?.channel.eventLoop.execute { - guard - let self = self, // the connection is already gone - var listeners = self.notificationListeners[channel] // we don't have the listeners for this topic ¯\_(ツ)_/¯ - else { - return - } + let task = HandlerTask.startListening(listener) + self.channel.write(task, promise: nil) - assert(listeners.filter { $0.0 === listenContext }.count <= 1, "Listeners can not appear twice in a channel!") - listeners.removeAll(where: { $0.0 === listenContext }) // just in case a listener shows up more than once in a release build, remove all, not just first - self.notificationListeners[channel] = listeners.isEmpty ? nil : listeners - } + listenContext.future.whenComplete { _ in + let task = HandlerTask.cancelListening(channel, id) + self.channel.write(task, promise: nil) } return listenContext } } -extension PostgresConnection: PSQLChannelHandlerNotificationDelegate { - func notificationReceived(_ notification: PostgresBackendMessage.NotificationResponse) { - self.eventLoop.assertInEventLoop() - - guard let listeners = self.notificationListeners[notification.channel] else { - return - } - - let postgresNotification = PostgresMessage.NotificationResponse( - backendPID: notification.backendPID, - channel: notification.channel, - payload: notification.payload) - - listeners.forEach { (listenContext, handler) in - handler(listenContext, postgresNotification) - } - } -} - enum CloseTarget { case preparedStatement(String) case portal(String) @@ -655,3 +848,20 @@ extension EventLoopFuture { } } } + +extension PostgresConnection { + /// Returns the default `EventLoopGroup` singleton, automatically selecting the best for the platform. + /// + /// This will select the concrete `EventLoopGroup` depending which platform this is running on. + public static var defaultEventLoopGroup: EventLoopGroup { +#if canImport(Network) + if #available(OSX 10.14, iOS 12.0, tvOS 12.0, watchOS 6.0, *) { + return NIOTSEventLoopGroup.singleton + } else { + return MultiThreadedEventLoopGroup.singleton + } +#else + return MultiThreadedEventLoopGroup.singleton +#endif + } +} diff --git a/Sources/PostgresNIO/Connection/PostgresDatabase+PreparedQuery.swift b/Sources/PostgresNIO/Connection/PostgresDatabase+PreparedQuery.swift index 074ba6de..56496172 100644 --- a/Sources/PostgresNIO/Connection/PostgresDatabase+PreparedQuery.swift +++ b/Sources/PostgresNIO/Connection/PostgresDatabase+PreparedQuery.swift @@ -1,4 +1,5 @@ import NIOCore +import NIOConcurrencyHelpers import struct Foundation.UUID extension PostgresDatabase { @@ -14,7 +15,8 @@ extension PostgresDatabase { } } - public func prepare(query: String, handler: @escaping (PreparedQuery) -> EventLoopFuture<[[PostgresRow]]>) -> EventLoopFuture<[[PostgresRow]]> { + @preconcurrency + public func prepare(query: String, handler: @Sendable @escaping (PreparedQuery) -> EventLoopFuture<[[PostgresRow]]>) -> EventLoopFuture<[[PostgresRow]]> { prepare(query: query) .flatMap { preparedQuery in handler(preparedQuery) @@ -26,7 +28,7 @@ extension PostgresDatabase { } -public struct PreparedQuery { +public struct PreparedQuery: Sendable { let underlying: PSQLPreparedStatement let database: PostgresDatabase @@ -36,11 +38,16 @@ public struct PreparedQuery { } public func execute(_ binds: [PostgresData] = []) -> EventLoopFuture<[PostgresRow]> { - var rows: [PostgresRow] = [] - return self.execute(binds) { rows.append($0) }.map { rows } + let rowsBoxed = NIOLockedValueBox([PostgresRow]()) + return self.execute(binds) { row in + rowsBoxed.withLockedValue { + $0.append(row) + } + }.map { rowsBoxed.withLockedValue { $0 } } } - public func execute(_ binds: [PostgresData] = [], _ onRow: @escaping (PostgresRow) throws -> ()) -> EventLoopFuture { + @preconcurrency + public func execute(_ binds: [PostgresData] = [], _ onRow: @Sendable @escaping (PostgresRow) throws -> ()) -> EventLoopFuture { let command = PostgresCommands.executePreparedStatement(query: self, binds: binds, onRow: onRow) return self.database.send(command, logger: self.database.logger) } @@ -50,15 +57,23 @@ public struct PreparedQuery { } } -final class PrepareQueryRequest { +final class PrepareQueryRequest: Sendable { let query: String let name: String - var prepared: PreparedQuery? = nil - - + var prepared: PreparedQuery? { + get { + self._prepared.withLockedValue { $0 } + } + set { + self._prepared.withLockedValue { + $0 = newValue + } + } + } + let _prepared: NIOLockedValueBox = .init(nil) + init(_ query: String, as name: String) { self.query = query self.name = name } - } diff --git a/Sources/PostgresNIO/Data/PostgresData+Numeric.swift b/Sources/PostgresNIO/Data/PostgresData+Numeric.swift index 5e564d6d..e736a61c 100644 --- a/Sources/PostgresNIO/Data/PostgresData+Numeric.swift +++ b/Sources/PostgresNIO/Data/PostgresData+Numeric.swift @@ -268,16 +268,10 @@ private extension Collection { // splits the collection into chunks of the supplied size // if the collection is not evenly divisible, the first chunk will be smaller func reverseChunked(by maxSize: Int) -> [SubSequence] { - var lastDistance = 0 var chunkStartIndex = self.startIndex return stride(from: 0, to: self.count, by: maxSize).reversed().map { current in - let distance = (self.count - current) - lastDistance - lastDistance = distance - let chunkEndOffset = Swift.min( - self.distance(from: chunkStartIndex, to: self.endIndex), - distance - ) - let chunkEndIndex = self.index(chunkStartIndex, offsetBy: chunkEndOffset) + let distance = self.count - current + let chunkEndIndex = self.index(self.startIndex, offsetBy: distance) defer { chunkStartIndex = chunkEndIndex } return self[chunkStartIndex..;` to lookup more information. +/// Data types and their raw OIDs. +/// +/// Use `select * from pg_type where oid = ` to look up more information for a given type. +/// +/// This list was generated by running `select oid, typname from pg_type where oid < 10000 order by oid` +/// and manually trimming Postgres-internal types. public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStringConvertible { /// `0` public static let null = PostgresDataType(0) @@ -41,6 +45,8 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri public static let int8 = PostgresDataType(20) /// `21` public static let int2 = PostgresDataType(21) + /// `22` + public static let int2vector = PostgresDataType(22) /// `23` public static let int4 = PostgresDataType(23) /// `24` @@ -49,18 +55,77 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri public static let text = PostgresDataType(25) /// `26` public static let oid = PostgresDataType(26) + /// `27` + public static let tid = PostgresDataType(27) + /// `28` + public static let xid = PostgresDataType(28) + /// `29` + public static let cid = PostgresDataType(29) + /// `30` + public static let oidvector = PostgresDataType(30) + /// `32` + public static let pgDDLCommand = PostgresDataType(32) /// `114` public static let json = PostgresDataType(114) + /// `142` + public static let xml = PostgresDataType(142) + /// `143` + public static let xmlArray = PostgresDataType(143) /// `194` pg_node_tree + @available(*, deprecated, message: "This is internal to Postgres and should not be used.") public static let pgNodeTree = PostgresDataType(194) + /// `199` + public static let jsonArray = PostgresDataType(199) + /// `269` + public static let tableAMHandler = PostgresDataType(269) + /// `271` + public static let xid8Array = PostgresDataType(271) + /// `325` + public static let indexAMHandler = PostgresDataType(325) /// `600` public static let point = PostgresDataType(600) + /// `601` + public static let lseg = PostgresDataType(601) + /// `602` + public static let path = PostgresDataType(602) + /// `603` + public static let box = PostgresDataType(603) + /// `604` + public static let polygon = PostgresDataType(604) + /// `628` + public static let line = PostgresDataType(628) + /// `629` + public static let lineArray = PostgresDataType(629) + /// `650` + public static let cidr = PostgresDataType(650) + /// `651` + public static let cidrArray = PostgresDataType(651) /// `700` public static let float4 = PostgresDataType(700) /// `701` public static let float8 = PostgresDataType(701) + /// `705` + public static let unknown = PostgresDataType(705) + /// `718` + public static let circle = PostgresDataType(718) + /// `719` + public static let circleArray = PostgresDataType(719) + /// `774` + public static let macaddr8 = PostgresDataType(774) + /// `775` + @available(*, deprecated, renamed: "macaddr8Array") + public static let macaddr8Aray = Self.macaddr8Array + public static let macaddr8Array = PostgresDataType(775) /// `790` public static let money = PostgresDataType(790) + /// `791` + @available(*, deprecated, renamed: "moneyArray") + public static let _money = Self.moneyArray + public static let moneyArray = PostgresDataType(791) + /// `829` + public static let macaddr = PostgresDataType(829) + /// `869` + public static let inet = PostgresDataType(869) /// `1000` _bool public static let boolArray = PostgresDataType(1000) /// `1001` _bytea @@ -71,22 +136,52 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri public static let nameArray = PostgresDataType(1003) /// `1005` _int2 public static let int2Array = PostgresDataType(1005) + /// `1006` + public static let int2vectorArray = PostgresDataType(1006) /// `1007` _int4 public static let int4Array = PostgresDataType(1007) + /// `1008` + public static let regprocArray = PostgresDataType(1008) /// `1009` _text public static let textArray = PostgresDataType(1009) + /// `1010` + public static let tidArray = PostgresDataType(1010) + /// `1011` + public static let xidArray = PostgresDataType(1011) + /// `1012` + public static let cidArray = PostgresDataType(1012) + /// `1013` + public static let oidvectorArray = PostgresDataType(1013) + /// `1014` + public static let bpcharArray = PostgresDataType(1014) /// `1015` _varchar public static let varcharArray = PostgresDataType(1015) /// `1016` _int8 public static let int8Array = PostgresDataType(1016) /// `1017` _point public static let pointArray = PostgresDataType(1017) + /// `1018` + public static let lsegArray = PostgresDataType(1018) + /// `1019` + public static let pathArray = PostgresDataType(1019) + /// `1020` + public static let boxArray = PostgresDataType(1020) /// `1021` _float4 public static let float4Array = PostgresDataType(1021) /// `1022` _float8 public static let float8Array = PostgresDataType(1022) + /// `1027` + public static let polygonArray = PostgresDataType(1027) + /// `1028` + public static let oidArray = PostgresDataType(1018) + /// `1033` + public static let aclitem = PostgresDataType(1033) /// `1034` _aclitem public static let aclitemArray = PostgresDataType(1034) + /// `1040` + public static let macaddrArray = PostgresDataType(1040) + /// `1041` + public static let inetArray = PostgresDataType(1041) /// `1042` public static let bpchar = PostgresDataType(1042) /// `1043` @@ -99,22 +194,202 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri public static let timestamp = PostgresDataType(1114) /// `1115` _timestamp public static let timestampArray = PostgresDataType(1115) + /// `1182` + public static let dateArray = PostgresDataType(1182) + /// `1183` + public static let timeArray = PostgresDataType(1183) /// `1184` public static let timestamptz = PostgresDataType(1184) + /// `1185` + public static let timestamptzArray = PostgresDataType(1185) + /// `1186` + public static let interval = PostgresDataType(1186) + /// `1187` + public static let intervalArray = PostgresDataType(1187) + /// `1231` + public static let numericArray = PostgresDataType(1231) + /// `1263` + public static let cstringArray = PostgresDataType(1263) /// `1266` public static let timetz = PostgresDataType(1266) + /// `1270` + public static let timetzArray = PostgresDataType(1270) + /// `1560` + public static let bit = PostgresDataType(1560) + /// `1561` + public static let bitArray = PostgresDataType(1561) + /// `1562` + public static let varbit = PostgresDataType(1562) + /// `1563` + public static let varbitArray = PostgresDataType(1563) /// `1700` public static let numeric = PostgresDataType(1700) + /// `1790` + public static let refcursor = PostgresDataType(1790) + /// `2201` + public static let refcursorArray = PostgresDataType(2201) + /// `2202` + public static let regprocedure = PostgresDataType(2202) + /// `2203` + public static let regoper = PostgresDataType(2203) + /// `2204` + public static let regoperator = PostgresDataType(2204) + /// `2205` + public static let regclass = PostgresDataType(2205) + /// `2206` + public static let regtype = PostgresDataType(2206) + /// `2207` + public static let regprocedureArray = PostgresDataType(2207) + /// `2208` + public static let regoperArray = PostgresDataType(2208) + /// `2209` + public static let regoperatorArray = PostgresDataType(2209) + /// `2210` + public static let regclassArray = PostgresDataType(2210) + /// `2211` + public static let regtypeArray = PostgresDataType(2211) + /// `2249` + public static let record = PostgresDataType(2249) + /// `2275` + public static let cstring = PostgresDataType(2275) + /// `2276` + public static let any = PostgresDataType(2276) + /// `2277` + public static let anyarray = PostgresDataType(2277) /// `2278` public static let void = PostgresDataType(2278) + /// `2279` + public static let trigger = PostgresDataType(2279) + /// `2280` + public static let languageHandler = PostgresDataType(2280) + /// `2281` + public static let `internal` = PostgresDataType(2281) + /// `2283` + public static let anyelement = PostgresDataType(2283) + /// `2287` + public static let recordArray = PostgresDataType(2287) + /// `2776` + public static let anynonarray = PostgresDataType(2776) /// `2950` public static let uuid = PostgresDataType(2950) /// `2951` _uuid public static let uuidArray = PostgresDataType(2951) + /// `3115` + public static let fdwHandler = PostgresDataType(3115) + /// `3220` + public static let pgLSN = PostgresDataType(3220) + /// `3221` + public static let pgLSNArray = PostgresDataType(3221) + /// `3310` + public static let tsmHandler = PostgresDataType(3310) + /// `3500` + public static let anyenum = PostgresDataType(3500) + /// `3614` + public static let tsvector = PostgresDataType(3614) + /// `3615` + public static let tsquery = PostgresDataType(3615) + /// `3642` + public static let gtsvector = PostgresDataType(3642) + /// `3643` + public static let tsvectorArray = PostgresDataType(3643) + /// `3644` + public static let gtsvectorArray = PostgresDataType(3644) + /// `3645` + public static let tsqueryArray = PostgresDataType(3645) + /// `3734` + public static let regconfig = PostgresDataType(3734) + /// `3735` + public static let regconfigArray = PostgresDataType(3735) + /// `3769` + public static let regdictionary = PostgresDataType(3769) + /// `3770` + public static let regdictionaryArray = PostgresDataType(3770) /// `3802` public static let jsonb = PostgresDataType(3802) /// `3807` _jsonb public static let jsonbArray = PostgresDataType(3807) + /// `3831` + public static let anyrange = PostgresDataType(3831) + /// `3838` + public static let eventTrigger = PostgresDataType(3838) + /// `3904` + public static let int4Range = PostgresDataType(3904) + /// `3905` _int4range + public static let int4RangeArray = PostgresDataType(3905) + /// `3906` + public static let numrange = PostgresDataType(3906) + /// `3907` + public static let numrangeArray = PostgresDataType(3907) + /// `3908` + public static let tsrange = PostgresDataType(3908) + /// `3909` + public static let tsrangeArray = PostgresDataType(3909) + /// `3910` + public static let tstzrange = PostgresDataType(3910) + /// `3911` + public static let tstzrangeArray = PostgresDataType(3911) + /// `3912` + public static let daterange = PostgresDataType(3912) + /// `3913` + public static let daterangeArray = PostgresDataType(3913) + /// `3926` + public static let int8Range = PostgresDataType(3926) + /// `3927` _int8range + public static let int8RangeArray = PostgresDataType(3927) + /// `4072` + public static let jsonpath = PostgresDataType(4072) + /// `4073` + public static let jsonpathArray = PostgresDataType(4073) + /// `4089` + public static let regnamespace = PostgresDataType(4089) + /// `4090` + public static let regnamespaceArray = PostgresDataType(4090) + /// `4096` + public static let regrole = PostgresDataType(4096) + /// `4097` + public static let regroleArray = PostgresDataType(4097) + /// `4191` + public static let regcollation = PostgresDataType(4191) + /// `4192` + public static let regcollationArray = PostgresDataType(4192) + /// `4451` + public static let int4multirange = PostgresDataType(4451) + /// `4532` + public static let nummultirange = PostgresDataType(4532) + /// `4533` + public static let tsmultirange = PostgresDataType(4533) + /// `4534` + public static let tstzmultirange = PostgresDataType(4534) + /// `4535` + public static let datemultirange = PostgresDataType(4535) + /// `4536` + public static let int8multirange = PostgresDataType(4536) + /// `4537` + public static let anymultirange = PostgresDataType(4537) + /// `4538` + public static let anycompatiblemultirange = PostgresDataType(4538) + /// `5069` + public static let xid8 = PostgresDataType(5069) + /// `5077` + public static let anycompatible = PostgresDataType(5077) + /// `5078` + public static let anycompatiblearray = PostgresDataType(5078) + /// `5079` + public static let anycompatiblenonarray = PostgresDataType(5079) + /// `5080` + public static let anycompatiblerange = PostgresDataType(5080) + /// `6150` + public static let int4multirangeArray = PostgresDataType(6150) + /// `6151` + public static let nummultirangeArray = PostgresDataType(6151) + /// `6152` + public static let tsmultirangeArray = PostgresDataType(6152) + /// `6153` + public static let tstzmultirangeArray = PostgresDataType(6153) + /// `6155` + public static let datemultirangeArray = PostgresDataType(6155) + /// `6157` + public static let int8multirangeArray = PostgresDataType(6157) /// The raw data type code recognized by PostgreSQL. public var rawValue: UInt32 @@ -136,57 +411,252 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri /// Returns the known SQL name, if one exists. /// Note: This only supports a limited subset of all PSQL types and is meant for convenience only. + /// This list was manually generated. public var knownSQLName: String? { switch self { + case .null: return "NULL" case .bool: return "BOOLEAN" case .bytea: return "BYTEA" case .char: return "CHAR" case .name: return "NAME" case .int8: return "BIGINT" case .int2: return "SMALLINT" + case .int2vector: return "INT2VECTOR" case .int4: return "INTEGER" case .regproc: return "REGPROC" case .text: return "TEXT" case .oid: return "OID" + case .tid: return "TID" + case .xid: return "XID" + case .cid: return "CID" + case .oidvector: return "OIDVECTOR" + case .pgDDLCommand: return "PG_DDL_COMMAND" case .json: return "JSON" - case .pgNodeTree: return "PGNODETREE" + case .xml: return "XML" + case .xmlArray: return "XML[]" + case .jsonArray: return "JSON[]" + case .tableAMHandler: return "TABLE_AM_HANDLER" + case .xid8Array: return "XID8[]" + case .indexAMHandler: return "INDEX_AM_HANDLER" case .point: return "POINT" + case .lseg: return "LSEG" + case .path: return "PATH" + case .box: return "BOX" + case .polygon: return "POLYGON" + case .line: return "LINE" + case .lineArray: return "LINE[]" + case .cidr: return "CIDR" + case .cidrArray: return "CIDR[]" case .float4: return "REAL" case .float8: return "DOUBLE PRECISION" + case .circle: return "CIRCLE" + case .circleArray: return "CIRCLE[]" + case .macaddr8: return "MACADDR8" + case .macaddr8Array: return "MACADDR8[]" case .money: return "MONEY" + case .moneyArray: return "MONEY[]" + case .macaddr: return "MACADDR" + case .inet: return "INET" case .boolArray: return "BOOLEAN[]" case .byteaArray: return "BYTEA[]" case .charArray: return "CHAR[]" case .nameArray: return "NAME[]" case .int2Array: return "SMALLINT[]" + case .int2vectorArray: return "INT2VECTOR[]" case .int4Array: return "INTEGER[]" + case .regprocArray: return "REGPROC[]" case .textArray: return "TEXT[]" + case .tidArray: return "TID[]" + case .xidArray: return "XID[]" + case .cidArray: return "CID[]" + case .oidvectorArray: return "OIDVECTOR[]" + case .bpcharArray: return "CHARACTER[]" case .varcharArray: return "VARCHAR[]" case .int8Array: return "BIGINT[]" case .pointArray: return "POINT[]" + case .lsegArray: return "LSEG[]" + case .pathArray: return "PATH[]" + case .boxArray: return "BOX[]" case .float4Array: return "REAL[]" case .float8Array: return "DOUBLE PRECISION[]" + case .polygonArray: return "POLYGON[]" + case .oidArray: return "OID[]" + case .aclitem: return "ACLITEM" case .aclitemArray: return "ACLITEM[]" - case .bpchar: return "BPCHAR" + case .macaddrArray: return "MACADDR[]" + case .inetArray: return "INET[]" + case .bpchar: return "CHARACTER" case .varchar: return "VARCHAR" case .date: return "DATE" case .time: return "TIME" case .timestamp: return "TIMESTAMP" - case .timestamptz: return "TIMESTAMPTZ" case .timestampArray: return "TIMESTAMP[]" + case .dateArray: return "DATE[]" + case .timeArray: return "TIME[]" + case .timestamptz: return "TIMESTAMPTZ" + case .timestamptzArray: return "TIMESTAMPTZ[]" + case .interval: return "INTERVAL" + case .intervalArray: return "INTERVAL[]" + case .numericArray: return "NUMERIC[]" + case .cstringArray: return "CSTRING[]" + case .timetz: return "TIMETZ" + case .timetzArray: return "TIMETZ[]" + case .bit: return "BIT" + case .bitArray: return "BIT[]" + case .varbit: return "VARBIT" + case .varbitArray: return "VARBIT[]" case .numeric: return "NUMERIC" + case .refcursor: return "REFCURSOR" + case .refcursorArray: return "REFCURSOR[]" + case .regprocedure: return "REGPROCEDURE" + case .regoper: return "REGOPER" + case .regoperator: return "REGOPERATOR" + case .regclass: return "REGCLASS" + case .regtype: return "REGTYPE" + case .regprocedureArray: return "REGPROCEDURE[]" + case .regoperArray: return "REGOPER[]" + case .regoperatorArray: return "REGOPERATOR[]" + case .regclassArray: return "REGCLASS[]" + case .regtypeArray: return "REGTYPE[]" + case .record: return "RECORD" + case .cstring: return "CSTRING" + case .any: return "ANY" + case .anyarray: return "ANYARRAY" case .void: return "VOID" + case .trigger: return "TRIGGER" + case .languageHandler: return "LANGUAGE_HANDLER" + case .`internal`: return "INTERNAL" + case .anyelement: return "ANYELEMENT" + case .recordArray: return "RECORD[]" + case .anynonarray: return "ANYNONARRAY" case .uuid: return "UUID" case .uuidArray: return "UUID[]" + case .fdwHandler: return "FDW_HANDLER" + case .pgLSN: return "PG_LSN" + case .pgLSNArray: return "PG_LSN[]" + case .tsmHandler: return "TSM_HANDLER" + case .anyenum: return "ANYENUM" + case .tsvector: return "TSVECTOR" + case .tsquery: return "TSQUERY" + case .gtsvector: return "GTSVECTOR" + case .tsvectorArray: return "TSVECTOR[]" + case .gtsvectorArray: return "GTSVECTOR[]" + case .tsqueryArray: return "TSQUERY[]" + case .regconfig: return "REGCONFIG" + case .regconfigArray: return "REGCONFIG[]" + case .regdictionary: return "REGDICTIONARY" + case .regdictionaryArray: return "REGDICTIONARY[]" case .jsonb: return "JSONB" case .jsonbArray: return "JSONB[]" + case .anyrange: return "ANYRANGE" + case .eventTrigger: return "EVENT_TRIGGER" + case .int4Range: return "INT4RANGE" + case .int4RangeArray: return "INT4RANGE[]" + case .numrange: return "NUMRANGE" + case .numrangeArray: return "NUMRANGE[]" + case .tsrange: return "TSRANGE" + case .tsrangeArray: return "TSRANGE[]" + case .tstzrange: return "TSTZRANGE" + case .tstzrangeArray: return "TSTZRANGE[]" + case .daterange: return "DATERANGE" + case .daterangeArray: return "DATERANGE[]" + case .int8Range: return "INT8RANGE" + case .int8RangeArray: return "INT8RANGE[]" + case .jsonpath: return "JSONPATH" + case .jsonpathArray: return "JSONPATH[]" + case .regnamespace: return "REGNAMESPACE" + case .regnamespaceArray: return "REGNAMESPACE[]" + case .regrole: return "REGROLE" + case .regroleArray: return "REGROLE[]" + case .regcollation: return "REGCOLLATION" + case .regcollationArray: return "REGCOLLATION[]" + case .int4multirange: return "INT4MULTIRANGE" + case .nummultirange: return "NUMMULTIRANGE" + case .tsmultirange: return "TSMULTIRANGE" + case .tstzmultirange: return "TSTZMULTIRANGE" + case .datemultirange: return "DATEMULTIRANGE" + case .int8multirange: return "INT8MULTIRANGE" + case .anymultirange: return "ANYMULTIRANGE" + case .anycompatiblemultirange: return "ANYCOMPATIBLEMULTIRANGE" + case .xid8: return "XID8" + case .anycompatible: return "ANYCOMPATIBLE" + case .anycompatiblearray: return "ANYCOMPATIBLEARRAY" + case .anycompatiblenonarray: return "ANYCOMPATIBLENONARRAY" + case .anycompatiblerange: return "ANYCOMPATIBLERANG" + case .int4multirangeArray: return "INT4MULTIRANGE[]" + case .nummultirangeArray: return "NUMMULTIRANGE[]" + case .tsmultirangeArray: return "TSMULTIRANGE[]" + case .tstzmultirangeArray: return "TSTZMULTIRANGE[]" + case .datemultirangeArray: return "DATEMULTIRANGE[]" + case .int8multirangeArray: return "INT8MULTIRANGE[]" default: return nil } } /// Returns the array type for this type if one is known. + /// + /// This list was manually generated. internal var arrayType: PostgresDataType? { switch self { + case .xml: return .xmlArray + case .json: return .jsonArray + case .xid8: return .xid8Array + case .line: return .lineArray + case .cidr: return .cidrArray + case .circle: return .circleArray + case .macaddr8: return .macaddr8Array + case .money: return .moneyArray + case .int2vector: return .int2vectorArray + case .regproc: return .regprocArray + case .tid: return .tidArray + case .xid: return .xidArray + case .cid: return .cidArray + case .oidvector: return .oidvectorArray + case .bpchar: return .bpcharArray + case .lseg: return .lsegArray + case .path: return .pathArray + case .box: return .boxArray + case .polygon: return .polygonArray + case .oid: return .oidArray + case .aclitem: return .aclitemArray + case .macaddr: return .macaddrArray + case .inet: return .inetArray + case .timestamp: return .timestampArray + case .date: return .dateArray + case .time: return .timeArray + case .timestamptz: return .timestamptzArray + case .interval: return .intervalArray + case .numeric: return .numericArray + case .cstring: return .cstringArray + case .timetz: return .timetzArray + case .bit: return .bitArray + case .varbit: return .varbitArray + case .refcursor: return .refcursorArray + case .regprocedure: return .regprocedureArray + case .regoper: return .regoperArray + case .regoperator: return .regoperatorArray + case .regclass: return .regclassArray + case .regtype: return .regtypeArray + case .record: return .recordArray + case .pgLSN: return .pgLSNArray + case .tsvector: return .tsvectorArray + case .gtsvector: return .gtsvectorArray + case .tsquery: return .tsqueryArray + case .regconfig: return .regconfigArray + case .regdictionary: return .regdictionaryArray + case .numrange: return .numrangeArray + case .tsrange: return .tsrangeArray + case .tstzrange: return .tstzrangeArray + case .daterange: return .daterangeArray + case .jsonpath: return .jsonpathArray + case .regnamespace: return .regnamespaceArray + case .regrole: return .regroleArray + case .regcollation: return .regcollationArray + case .int4multirange: return .int4multirangeArray + case .tsmultirange: return .tsmultirangeArray + case .tstzmultirange: return .tstzmultirangeArray + case .datemultirange: return .datemultirangeArray + case .int8multirange: return .int8multirangeArray case .bool: return .boolArray case .bytea: return .byteaArray case .char: return .charArray @@ -201,14 +671,77 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri case .jsonb: return .jsonbArray case .text: return .textArray case .varchar: return .varcharArray + case .int4Range: return .int4RangeArray + case .int8Range: return .int8RangeArray default: return nil } } /// Returns the element type for this type if one is known. /// Returns nil if this is not an array type. + /// + /// This list was manually generated. internal var elementType: PostgresDataType? { switch self { + case .xmlArray: return .xml + case .jsonArray: return .json + case .xid8Array: return .xid8 + case .lineArray: return .line + case .cidrArray: return .cidr + case .circleArray: return .circle + case .macaddr8Array: return .macaddr8 + case .moneyArray: return .money + case .int2vectorArray: return .int2vector + case .regprocArray: return .regproc + case .tidArray: return .tid + case .xidArray: return .xid + case .cidArray: return .cid + case .oidvectorArray: return .oidvector + case .bpcharArray: return .bpchar + case .lsegArray: return .lseg + case .pathArray: return .path + case .boxArray: return .box + case .polygonArray: return .polygon + case .oidArray: return .oid + case .aclitemArray: return .aclitem + case .macaddrArray: return .macaddr + case .inetArray: return .inet + case .timestampArray: return .timestamp + case .dateArray: return .date + case .timeArray: return .time + case .timestamptzArray: return .timestamptz + case .intervalArray: return .interval + case .numericArray: return .numeric + case .cstringArray: return .cstring + case .timetzArray: return .timetz + case .bitArray: return .bit + case .varbitArray: return .varbit + case .refcursorArray: return .refcursor + case .regprocedureArray: return .regprocedure + case .regoperArray: return .regoper + case .regoperatorArray: return .regoperator + case .regclassArray: return .regclass + case .regtypeArray: return .regtype + case .recordArray: return .record + case .pgLSNArray: return .pgLSN + case .tsvectorArray: return .tsvector + case .gtsvectorArray: return .gtsvector + case .tsqueryArray: return .tsquery + case .regconfigArray: return .regconfig + case .regdictionaryArray: return .regdictionary + case .numrangeArray: return .numrange + case .tsrangeArray: return .tsrange + case .tstzrangeArray: return .tstzrange + case .daterangeArray: return .daterange + case .jsonpathArray: return .jsonpath + case .regnamespaceArray: return .regnamespace + case .regroleArray: return .regrole + case .regcollationArray: return .regcollation + case .int4multirangeArray: return .int4multirange + case .tsmultirangeArray: return .tsmultirange + case .tstzmultirangeArray: return .tstzmultirange + case .datemultirangeArray: return .datemultirange + case .int8multirangeArray: return .int8multirange case .boolArray: return .bool case .byteaArray: return .bytea case .charArray: return .char @@ -223,11 +756,30 @@ public struct PostgresDataType: RawRepresentable, Sendable, Hashable, CustomStri case .jsonbArray: return .jsonb case .textArray: return .text case .varcharArray: return .varchar + case .int4RangeArray: return .int4Range + case .int8RangeArray: return .int8Range + default: return nil + } + } + + /// Returns the bound type for this type if one is known. + /// Returns nil if this is not a range type. + /// + /// This list was manually generated. + @usableFromInline + internal var boundType: PostgresDataType? { + switch self { + case .int4Range: return .int4 + case .int8Range: return .int8 + case .numrange: return .numeric + case .tsrange: return .timestamp + case .tstzrange: return .timestamptz + case .daterange: return .date default: return nil } } - /// See `CustomStringConvertible`. + // See `CustomStringConvertible.description`. public var description: String { return self.knownSQLName ?? "UNKNOWN \(self.rawValue)" } diff --git a/Sources/PostgresNIO/Docs.docc/coding.md b/Sources/PostgresNIO/Docs.docc/coding.md new file mode 100644 index 00000000..3bcc4a7e --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/coding.md @@ -0,0 +1,39 @@ +# PostgreSQL data types + +Translate Swift data types to Postgres data types and vica versa. Learn how to write translations +for your own custom Swift types. + +## Topics + +### Essentials + +- ``PostgresCodable`` +- ``PostgresDataType`` +- ``PostgresFormat`` +- ``PostgresNumeric`` + +### Encoding + +- ``PostgresEncodable`` +- ``PostgresNonThrowingEncodable`` +- ``PostgresDynamicTypeEncodable`` +- ``PostgresThrowingDynamicTypeEncodable`` +- ``PostgresArrayEncodable`` +- ``PostgresRangeEncodable`` +- ``PostgresRangeArrayEncodable`` +- ``PostgresEncodingContext`` + +### Decoding + +- ``PostgresDecodable`` +- ``PostgresArrayDecodable`` +- ``PostgresRangeDecodable`` +- ``PostgresRangeArrayDecodable`` +- ``PostgresDecodingContext`` + +### JSON + +- ``PostgresJSONEncoder`` +- ``PostgresJSONDecoder`` + + diff --git a/Sources/PostgresNIO/Docs.docc/deprecated.md b/Sources/PostgresNIO/Docs.docc/deprecated.md new file mode 100644 index 00000000..a29465f6 --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/deprecated.md @@ -0,0 +1,43 @@ +# Deprecations + +`PostgresNIO` follows SemVer 2.0.0. Learn which APIs are considered deprecated and how to migrate to +their replacements. + +``PostgresNIO`` reached 1.0 in April 2020. Since then the maintainers have been hard at work to +guarantee API stability. However as the Swift and Swift on server ecosystem have matured approaches +have changed. The introduction of structured concurrency changed what developers expect from a +modern Swift library. Because of this ``PostgresNIO`` added various APIs that embrace the new Swift +patterns. This means however, that PostgresNIO still offers APIs that have fallen out of favor. +Those are documented here. All those APIs will be removed once the maintainers release the next +major version. The maintainers recommend all adopters to move of those APIs sooner rather than +later. + +## Topics + +### Migrate of deprecated APIs + +- + +### Deprecated APIs + +These types are already deprecated or will be deprecated in the near future. All of them will be +removed from the public API with the next major release. + +- ``PostgresDatabase`` +- ``PostgresData`` +- ``PostgresDataConvertible`` +- ``PostgresQueryResult`` +- ``PostgresJSONCodable`` +- ``PostgresJSONBCodable`` +- ``PostgresMessageEncoder`` +- ``PostgresMessageDecoder`` +- ``PostgresRequest`` +- ``PostgresMessage`` +- ``PostgresMessageType`` +- ``PostgresFormatCode`` +- ``PostgresListenContext`` +- ``PreparedQuery`` +- ``SASLAuthenticationManager`` +- ``SASLAuthenticationMechanism`` +- ``SASLAuthenticationError`` +- ``SASLAuthenticationStepResult`` diff --git a/Sources/PostgresNIO/Docs.docc/images/vapor-postgresnio-logo.svg b/Sources/PostgresNIO/Docs.docc/images/vapor-postgresnio-logo.svg new file mode 100644 index 00000000..a831189d --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/images/vapor-postgresnio-logo.svg @@ -0,0 +1,21 @@ + + + + + + + + + + + + + + + + + + diff --git a/Sources/PostgresNIO/Docs.docc/index.md b/Sources/PostgresNIO/Docs.docc/index.md index e7363054..6355a7a4 100644 --- a/Sources/PostgresNIO/Docs.docc/index.md +++ b/Sources/PostgresNIO/Docs.docc/index.md @@ -1,83 +1,58 @@ # ``PostgresNIO`` -🐘 Non-blocking, event-driven Swift client for PostgreSQL built on [SwiftNIO]. +@Metadata { + @TitleHeading(Package) +} -## Overview - -Features: +🐘 Non-blocking, event-driven Swift client for PostgreSQL built on SwiftNIO. -- A ``PostgresConnection`` which allows you to connect to, authorize with, query, and retrieve results from a PostgreSQL server -- An async/await interface that supports backpressure -- Automatic conversions between Swift primitive types and the Postgres wire format -- Integrated with the Swift server ecosystem, including use of [SwiftLog]. -- Designed to run efficiently on all supported platforms (tested extensively on Linux and Darwin systems) -- Support for `Network.framework` when available (e.g. on Apple platforms) - -## Topics +## Overview -### Articles +``PostgresNIO`` allows you to connect to, authorize with, query, and retrieve results from a +PostgreSQL server. PostgreSQL is an open source relational database. -- +Use a ``PostgresConnection`` to create a connection to the PostgreSQL server. You can then use it to +run queries and prepared statements against the server. ``PostgresConnection`` also supports +PostgreSQL's Listen & Notify API. -### Connections +Developers, who don't want to manage connections themselves, can use the ``PostgresClient``, which +offers the same functionality as ``PostgresConnection``. ``PostgresClient`` +pools connections for rapid connection reuse and hides the complexities of connection +management from the user, allowing developers to focus on their SQL queries. ``PostgresClient`` +implements the `Service` protocol from Service Lifecycle allowing an easy adoption in Swift server +applications. -- ``PostgresConnection`` +``PostgresNIO`` embraces Swift structured concurrency, offering async/await APIs which handle +task cancellation. The query interface makes use of backpressure to ensure that memory can not grow +unbounded for queries that return thousands of rows. -### Querying - -- ``PostgresQuery`` -- ``PostgresBindings`` -- ``PostgresRow`` -- ``PostgresRowSequence`` -- ``PostgresRandomAccessRow`` -- ``PostgresCell`` -- ``PreparedQuery`` -- ``PostgresQueryMetadata`` +``PostgresNIO`` runs efficiently on Linux and Apple platforms. On Apple platforms developers can +configure ``PostgresConnection`` to use `Network.framework` as the underlying transport framework. + +## Topics -### Encoding and Decoding +### Essentials -- ``PostgresEncodable`` -- ``PostgresEncodingContext`` -- ``PostgresDecodable`` -- ``PostgresDecodingContext`` -- ``PostgresArrayEncodable`` -- ``PostgresArrayDecodable`` -- ``PostgresJSONEncoder`` -- ``PostgresJSONDecoder`` -- ``PostgresDataType`` -- ``PostgresFormat`` -- ``PostgresNumeric`` +- ``PostgresClient`` +- ``PostgresClient/Configuration`` +- ``PostgresConnection`` +- -### Notifications +### Advanced -- ``PostgresListenContext`` +- +- +- ### Errors - ``PostgresError`` - ``PostgresDecodingError`` +- ``PSQLError`` + +### Deprecations -### Deprecated - -These types are already deprecated or will be deprecated in the near future. All of them will be -removed from the public API with the next major release. - -- ``PostgresDatabase`` -- ``PostgresData`` -- ``PostgresDataConvertible`` -- ``PostgresQueryResult`` -- ``PostgresJSONCodable`` -- ``PostgresJSONBCodable`` -- ``PostgresMessageEncoder`` -- ``PostgresMessageDecoder`` -- ``PostgresRequest`` -- ``PostgresMessage`` -- ``PostgresMessageType`` -- ``PostgresFormatCode`` -- ``SASLAuthenticationManager`` -- ``SASLAuthenticationMechanism`` -- ``SASLAuthenticationError`` -- ``SASLAuthenticationStepResult`` +- [SwiftNIO]: https://github.com/apple/swift-nio [SwiftLog]: https://github.com/apple/swift-log diff --git a/Sources/PostgresNIO/Docs.docc/listen.md b/Sources/PostgresNIO/Docs.docc/listen.md new file mode 100644 index 00000000..10c5d8bf --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/listen.md @@ -0,0 +1,9 @@ +# Listen & Notify + +``PostgresNIO`` supports PostgreSQL's listen and notify API. Learn how to listen for changes and +notify other listeners. + +## Topics + +- ``PostgresNotification`` +- ``PostgresNotificationSequence`` diff --git a/Sources/PostgresNIO/Docs.docc/migrations.md b/Sources/PostgresNIO/Docs.docc/migrations.md index 33c8afd4..3a7c634a 100644 --- a/Sources/PostgresNIO/Docs.docc/migrations.md +++ b/Sources/PostgresNIO/Docs.docc/migrations.md @@ -6,7 +6,7 @@ which use the ``PostgresRow/column(_:)`` API today. ## TLDR 1. Map your sequence of ``PostgresRow``s to ``PostgresRandomAccessRow``s. -2. Use the ``PostgresRandomAccessRow/subscript(name:)`` API to receive a ``PostgresCell`` +2. Use the ``PostgresRandomAccessRow/subscript(_:)-3facl`` API to receive a ``PostgresCell`` 3. Decode the ``PostgresCell`` into a Swift type using the ``PostgresCell/decode(_:file:line:)`` method. ```swift @@ -87,16 +87,4 @@ connection.query("SELECT id, name, email, age FROM users").whenComplete { } ``` -## Topics - -### Relevant types - -- ``PostgresConnection`` -- ``PostgresQuery`` -- ``PostgresBindings`` -- ``PostgresRow`` -- ``PostgresRandomAccessRow`` -- ``PostgresEncodable`` -- ``PostgresDecodable`` - [`1.9.0`]: https://github.com/vapor/postgres-nio/releases/tag/1.9.0 diff --git a/Sources/PostgresNIO/Docs.docc/prepared-statement.md b/Sources/PostgresNIO/Docs.docc/prepared-statement.md new file mode 100644 index 00000000..ff4b1c62 --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/prepared-statement.md @@ -0,0 +1,7 @@ +# Boosting Performance with Prepared Statements + +Improve performance by leveraging PostgreSQL's prepared statements. + +## Topics + +- ``PostgresPreparedStatement`` diff --git a/Sources/PostgresNIO/Docs.docc/running-queries.md b/Sources/PostgresNIO/Docs.docc/running-queries.md new file mode 100644 index 00000000..b2c4586f --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/running-queries.md @@ -0,0 +1,27 @@ +# Running Queries + +Interact with the PostgreSQL database by running Queries. + +## Overview + + + +You interact with the Postgres database by running SQL [Queries]. + + + +``PostgresQuery`` conforms to + + +## Topics + +- ``PostgresQuery`` +- ``PostgresBindings`` +- ``PostgresRow`` +- ``PostgresRowSequence`` +- ``PostgresRandomAccessRow`` +- ``PostgresCell`` +- ``PostgresQueryMetadata`` + +[Queries]: doc:PostgresQuery +[`ExpressibleByStringInterpolation`]: https://developer.apple.com/documentation/swift/expressiblebystringinterpolation diff --git a/Sources/PostgresNIO/Docs.docc/theme-settings.json b/Sources/PostgresNIO/Docs.docc/theme-settings.json new file mode 100644 index 00000000..38914a04 --- /dev/null +++ b/Sources/PostgresNIO/Docs.docc/theme-settings.json @@ -0,0 +1,24 @@ +{ + "theme": { + "aside": { "border-radius": "16px", "border-style": "double", "border-width": "3px" }, + "border-radius": "0", + "button": { "border-radius": "16px", "border-width": "1px", "border-style": "solid" }, + "code": { "border-radius": "16px", "border-width": "1px", "border-style": "solid" }, + "color": { + "fill": { "dark": "#000", "light": "#fff" }, + "psqlnio": "#336791", + "documentation-intro-fill": "radial-gradient(circle at top, var(--color-psqlnio) 30%, #000 100%)", + "documentation-intro-accent": "var(--color-psqlnio)", + "documentation-intro-eyebrow": "white", + "documentation-intro-figure": "white", + "documentation-intro-title": "white", + "logo-base": { "dark": "#fff", "light": "#000" }, + "logo-shape": { "dark": "#000", "light": "#fff" } + }, + "icons": { "technology": "/postgresnio/images/vapor-postgresnio-logo.svg" } + }, + "features": { + "quickNavigation": { "enable": true }, + "i18n": { "enable": true } + } +} diff --git a/Sources/PostgresNIO/Message/PostgresMessage+Error.swift b/Sources/PostgresNIO/Message/PostgresMessage+Error.swift index 44f9e6bf..45cda21f 100644 --- a/Sources/PostgresNIO/Message/PostgresMessage+Error.swift +++ b/Sources/PostgresNIO/Message/PostgresMessage+Error.swift @@ -2,8 +2,8 @@ import NIOCore extension PostgresMessage { /// First message sent from the frontend during startup. - public struct Error: CustomStringConvertible { - public enum Field: UInt8, Hashable { + public struct Error: CustomStringConvertible, Sendable { + public enum Field: UInt8, Hashable, Sendable { /// Severity: the field contents are ERROR, FATAL, or PANIC (in an error message), /// or WARNING, NOTICE, DEBUG, INFO, or LOG (in a notice message), or a //// localized translation of one of these. Always present. diff --git a/Sources/PostgresNIO/Message/PostgresMessage+Identifier.swift b/Sources/PostgresNIO/Message/PostgresMessage+Identifier.swift index 786b91ef..5d111e3b 100644 --- a/Sources/PostgresNIO/Message/PostgresMessage+Identifier.swift +++ b/Sources/PostgresNIO/Message/PostgresMessage+Identifier.swift @@ -4,7 +4,7 @@ extension PostgresMessage { /// Identifies an incoming or outgoing postgres message. Sent as the first byte, before the message size. /// Values are not unique across all identifiers, meaning some messages will require keeping state to identify. @available(*, deprecated, message: "Will be removed from public API.") - public struct Identifier: ExpressibleByIntegerLiteral, Equatable, CustomStringConvertible { + public struct Identifier: Sendable, ExpressibleByIntegerLiteral, Equatable, CustomStringConvertible { // special public static let none: Identifier = 0x00 // special diff --git a/Sources/PostgresNIO/New/BufferedMessageEncoder.swift b/Sources/PostgresNIO/New/BufferedMessageEncoder.swift deleted file mode 100644 index f202fcff..00000000 --- a/Sources/PostgresNIO/New/BufferedMessageEncoder.swift +++ /dev/null @@ -1,35 +0,0 @@ -import NIOCore - -struct BufferedMessageEncoder { - private enum State { - case flushed - case writable - } - - private var buffer: ByteBuffer - private var state: State = .writable - private var encoder: PSQLFrontendMessageEncoder - - init(buffer: ByteBuffer, encoder: PSQLFrontendMessageEncoder) { - self.buffer = buffer - self.encoder = encoder - } - - mutating func encode(_ message: PostgresFrontendMessage) { - switch self.state { - case .flushed: - self.state = .writable - self.buffer.clear() - - case .writable: - break - } - - self.encoder.encode(data: message, out: &self.buffer) - } - - mutating func flush() -> ByteBuffer { - self.state = .flushed - return self.buffer - } -} diff --git a/Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift index 563bb026..9d264bcc 100644 --- a/Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift +++ b/Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift @@ -31,13 +31,11 @@ struct ConnectionStateMachine { case readyForQuery(ConnectionContext) case extendedQuery(ExtendedQueryStateMachine, ConnectionContext) - case prepareStatement(PrepareStatementStateMachine, ConnectionContext) case closeCommand(CloseStateMachine, ConnectionContext) - - case error(PSQLError) - case closing - case closed - + + case closing(PSQLError?) + case closed(clientInitiated: Bool, error: PSQLError?) + case modifying } @@ -89,10 +87,9 @@ struct ConnectionStateMachine { // --- general actions case sendParseDescribeBindExecuteSync(PostgresQuery) case sendBindExecuteSync(PSQLExecuteStatement) - case failQuery(ExtendedQueryContext, with: PSQLError, cleanupContext: CleanUpContext?) - case succeedQuery(ExtendedQueryContext, columns: [RowDescription.Column]) - case succeedQueryNoRowsComming(ExtendedQueryContext, commandTag: String) - + case failQuery(EventLoopPromise, with: PSQLError, cleanupContext: CleanUpContext?) + case succeedQuery(EventLoopPromise, with: QueryResult) + // --- streaming actions // actions if query has requested next row but we are waiting for backend case forwardRows([DataRow]) @@ -100,10 +97,10 @@ struct ConnectionStateMachine { case forwardStreamError(PSQLError, read: Bool, cleanupContext: CleanUpContext?) // Prepare statement actions - case sendParseDescribeSync(name: String, query: String) - case succeedPreparedStatementCreation(PrepareStatementContext, with: RowDescription?) - case failPreparedStatementCreation(PrepareStatementContext, with: PSQLError, cleanupContext: CleanUpContext?) - + case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType]) + case succeedPreparedStatementCreation(EventLoopPromise, with: RowDescription?) + case failPreparedStatementCreation(EventLoopPromise, with: PSQLError, cleanupContext: CleanUpContext?) + // Close actions case sendCloseSync(CloseTarget) case succeedClose(CloseCommandContext) @@ -159,9 +156,7 @@ struct ConnectionStateMachine { .authenticated, .readyForQuery, .extendedQuery, - .prepareStatement, .closeCommand, - .error, .closing, .closed, .modifying: @@ -173,9 +168,9 @@ struct ConnectionStateMachine { self.startAuthentication(authContext) } - mutating func close(_ promise: EventLoopPromise?) -> ConnectionAction { + mutating func gracefulClose(_ promise: EventLoopPromise?) -> ConnectionAction { switch self.state { - case .closing, .closed, .error: + case .closing, .closed: // we are already closed, but sometimes an upstream handler might want to close the // connection, though it has already been closed by the remote. Typical race condition. return .closeConnection(promise) @@ -183,7 +178,7 @@ struct ConnectionStateMachine { precondition(self.taskQueue.isEmpty, """ The state should only be .readyForQuery if there are no more tasks in the queue """) - self.state = .closing + self.state = .closing(nil) return .closeConnection(promise) default: switch self.quiescingState { @@ -197,14 +192,18 @@ struct ConnectionStateMachine { return .wait } } - + + mutating func close(promise: EventLoopPromise?) -> ConnectionAction { + return self.closeConnectionAndCleanup(.clientClosedConnection(underlying: nil), closePromise: promise) + } + mutating func closed() -> ConnectionAction { switch self.state { case .initialized: preconditionFailure("How can a connection be closed, if it was never connected.") case .closed: - preconditionFailure("How can a connection be closed, if it is already closed.") + return .wait case .authenticated, .sslRequestSent, @@ -214,12 +213,11 @@ struct ConnectionStateMachine { .authenticating, .readyForQuery, .extendedQuery, - .prepareStatement, .closeCommand: - return self.errorHappened(.uncleanShutdown) - - case .error, .closing: - self.state = .closed + return self.errorHappened(.serverClosedConnection(underlying: nil)) + + case .closing(let error): + self.state = .closed(clientInitiated: true, error: error) self.quiescingState = .notQuiescing return .fireChannelInactive @@ -245,9 +243,7 @@ struct ConnectionStateMachine { .authenticated, .readyForQuery, .extendedQuery, - .prepareStatement, .closeCommand, - .error, .closing, .closed: return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.sslSupported)) @@ -274,9 +270,7 @@ struct ConnectionStateMachine { .authenticated, .readyForQuery, .extendedQuery, - .prepareStatement, .closeCommand, - .error, .closing, .closed: return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.sslSupported)) @@ -296,9 +290,7 @@ struct ConnectionStateMachine { .authenticated, .readyForQuery, .extendedQuery, - .prepareStatement, .closeCommand, - .error, .closing, .closed: preconditionFailure("Can only add a ssl handler after negotiation: \(self.state)") @@ -322,9 +314,7 @@ struct ConnectionStateMachine { .authenticated, .readyForQuery, .extendedQuery, - .prepareStatement, .closeCommand, - .error, .closing, .closed: preconditionFailure("Can only establish a ssl connection after adding a ssl handler: \(self.state)") @@ -343,11 +333,10 @@ struct ConnectionStateMachine { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.authentication(message))) } - return self.avoidingStateMachineCoW { machine in - let action = authState.authenticationMessageReceived(message) - machine.state = .authenticating(authState) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = authState.authenticationMessageReceived(message) + self.state = .authenticating(authState) + return self.modify(with: action) } mutating func backendKeyDataReceived(_ keyData: PostgresBackendMessage.BackendKeyData) -> ConnectionAction { @@ -371,40 +360,31 @@ struct ConnectionStateMachine { .waitingToStartAuthentication, .authenticating, .closing: - self.state = .error(.unexpectedBackendMessage(.parameterStatus(status))) - return .wait + return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.parameterStatus(status))) case .authenticated(let keyData, var parameters): - return self.avoidingStateMachineCoW { machine in - parameters[status.parameter] = status.value - machine.state = .authenticated(keyData, parameters) - return .wait - } + self.state = .modifying // avoid CoW + parameters[status.parameter] = status.value + self.state = .authenticated(keyData, parameters) + return .wait + case .readyForQuery(var connectionContext): - return self.avoidingStateMachineCoW { machine in - connectionContext.parameters[status.parameter] = status.value - machine.state = .readyForQuery(connectionContext) - return .wait - } + self.state = .modifying // avoid CoW + connectionContext.parameters[status.parameter] = status.value + self.state = .readyForQuery(connectionContext) + return .wait + case .extendedQuery(let query, var connectionContext): - return self.avoidingStateMachineCoW { machine in - connectionContext.parameters[status.parameter] = status.value - machine.state = .extendedQuery(query, connectionContext) - return .wait - } - case .prepareStatement(let prepareState, var connectionContext): - return self.avoidingStateMachineCoW { machine in - connectionContext.parameters[status.parameter] = status.value - machine.state = .prepareStatement(prepareState, connectionContext) - return .wait - } + self.state = .modifying // avoid CoW + connectionContext.parameters[status.parameter] = status.value + self.state = .extendedQuery(query, connectionContext) + return .wait + case .closeCommand(let closeState, var connectionContext): - return self.avoidingStateMachineCoW { machine in - connectionContext.parameters[status.parameter] = status.value - machine.state = .closeCommand(closeState, connectionContext) - return .wait - } - case .error(_): + self.state = .modifying // avoid CoW + connectionContext.parameters[status.parameter] = status.value + self.state = .closeCommand(closeState, connectionContext) return .wait + case .initialized, .closed: preconditionFailure("We shouldn't receive messages if we are not connected") @@ -420,45 +400,35 @@ struct ConnectionStateMachine { .sslHandlerAdded, .waitingToStartAuthentication, .authenticated, - .readyForQuery, - .error: + .readyForQuery: return self.closeConnectionAndCleanup(.server(errorMessage)) case .authenticating(var authState): if authState.isComplete { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.error(errorMessage))) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = authState.errorReceived(errorMessage) - machine.state = .authenticating(authState) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = authState.errorReceived(errorMessage) + self.state = .authenticating(authState) + return self.modify(with: action) + case .closeCommand(var closeStateMachine, let connectionContext): if closeStateMachine.isComplete { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.error(errorMessage))) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = closeStateMachine.errorReceived(errorMessage) - machine.state = .closeCommand(closeStateMachine, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = closeStateMachine.errorReceived(errorMessage) + self.state = .closeCommand(closeStateMachine, connectionContext) + return self.modify(with: action) + case .extendedQuery(var extendedQueryState, let connectionContext): if extendedQueryState.isComplete { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.error(errorMessage))) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = extendedQueryState.errorReceived(errorMessage) - machine.state = .extendedQuery(extendedQueryState, connectionContext) - return machine.modify(with: action) - } - case .prepareStatement(var preparedState, let connectionContext): - if preparedState.isComplete { - return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.error(errorMessage))) - } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = preparedState.errorReceived(errorMessage) - machine.state = .prepareStatement(preparedState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = extendedQueryState.errorReceived(errorMessage) + self.state = .extendedQuery(extendedQueryState, connectionContext) + return self.modify(with: action) + case .closing: // If the state machine is in state `.closing`, the connection shutdown was initiated // by the client. This means a `TERMINATE` message has already been sent and the @@ -493,13 +463,6 @@ struct ConnectionStateMachine { let action = queryState.errorHappened(error) return self.modify(with: action) } - case .prepareStatement(var prepareState, _): - if prepareState.isComplete { - return self.closeConnectionAndCleanup(error) - } else { - let action = prepareState.errorHappened(error) - return self.modify(with: action) - } case .closeCommand(var closeState, _): if closeState.isComplete { return self.closeConnectionAndCleanup(error) @@ -507,8 +470,6 @@ struct ConnectionStateMachine { let action = closeState.errorHappened(error) return self.modify(with: action) } - case .error: - return .wait case .closing: // If the state machine is in state `.closing`, the connection shutdown was initiated // by the client. This means a `TERMINATE` message has already been sent and the @@ -530,11 +491,11 @@ struct ConnectionStateMachine { mutating func noticeReceived(_ notice: PostgresBackendMessage.NoticeResponse) -> ConnectionAction { switch self.state { case .extendedQuery(var extendedQuery, let connectionContext): - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = extendedQuery.noticeReceived(notice) - machine.state = .extendedQuery(extendedQuery, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = extendedQuery.noticeReceived(notice) + self.state = .extendedQuery(extendedQuery, connectionContext) + return self.modify(with: action) + default: return .wait } @@ -567,16 +528,6 @@ struct ConnectionStateMachine { self.state = .readyForQuery(connectionContext) return self.executeNextQueryFromQueue() - case .prepareStatement(let preparedStateMachine, var connectionContext): - guard preparedStateMachine.isComplete else { - return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.readyForQuery(transactionState))) - } - - connectionContext.transactionState = transactionState - - self.state = .readyForQuery(connectionContext) - return self.executeNextQueryFromQueue() - case .closeCommand(let closeStateMachine, var connectionContext): guard closeStateMachine.isComplete else { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.readyForQuery(transactionState))) @@ -593,33 +544,54 @@ struct ConnectionStateMachine { } mutating func enqueue(task: PSQLTask) -> ConnectionAction { + let psqlErrror: PSQLError + // check if we are quiescing. if so fail task immidiatly - if case .quiescing = self.quiescingState { - switch task { - case .extendedQuery(let queryContext): - return .failQuery(queryContext, with: .connectionQuiescing, cleanupContext: nil) - case .preparedStatement(let prepareContext): - return .failPreparedStatementCreation(prepareContext, with: .connectionQuiescing, cleanupContext: nil) - case .closeCommand(let closeContext): - return .failClose(closeContext, with: .connectionQuiescing, cleanupContext: nil) + switch self.quiescingState { + case .quiescing: + psqlErrror = PSQLError.clientClosedConnection(underlying: nil) + + case .notQuiescing: + switch self.state { + case .initialized, + .authenticated, + .authenticating, + .closeCommand, + .extendedQuery, + .sslNegotiated, + .sslHandlerAdded, + .sslRequestSent, + .waitingToStartAuthentication: + self.taskQueue.append(task) + return .wait + + case .readyForQuery: + return self.executeTask(task) + + case .closing(let error): + psqlErrror = PSQLError.clientClosedConnection(underlying: error) + + case .closed(clientInitiated: true, error: let error): + psqlErrror = PSQLError.clientClosedConnection(underlying: error) + + case .closed(clientInitiated: false, error: let error): + psqlErrror = PSQLError.serverClosedConnection(underlying: error) + + case .modifying: + preconditionFailure("Invalid state: \(self.state)") } } - switch self.state { - case .readyForQuery: - return self.executeTask(task) - case .closed: - switch task { - case .extendedQuery(let queryContext): - return .failQuery(queryContext, with: .connectionClosed, cleanupContext: nil) - case .preparedStatement(let prepareContext): - return .failPreparedStatementCreation(prepareContext, with: .connectionClosed, cleanupContext: nil) - case .closeCommand(let closeContext): - return .failClose(closeContext, with: .connectionClosed, cleanupContext: nil) + switch task { + case .extendedQuery(let queryContext): + switch queryContext.query { + case .executeStatement(_, let promise), .unnamed(_, let promise): + return .failQuery(promise, with: psqlErrror, cleanupContext: nil) + case .prepareStatement(_, _, _, let promise): + return .failPreparedStatementCreation(promise, with: psqlErrror, cleanupContext: nil) } - default: - self.taskQueue.append(task) - return .wait + case .closeCommand(let closeContext): + return .failClose(closeContext, with: psqlErrror, cleanupContext: nil) } } @@ -633,19 +605,16 @@ struct ConnectionStateMachine { .authenticating, .authenticated, .readyForQuery, - .prepareStatement, .closeCommand, - .error, .closing, .closed: return .wait case .extendedQuery(var extendedQuery, let connectionContext): - return self.avoidingStateMachineCoW { machine in - let action = extendedQuery.channelReadComplete() - machine.state = .extendedQuery(extendedQuery, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = extendedQuery.channelReadComplete() + self.state = .extendedQuery(extendedQuery, connectionContext) + return self.modify(with: action) case .modifying: preconditionFailure("Invalid state") @@ -655,47 +624,40 @@ struct ConnectionStateMachine { mutating func readEventCaught() -> ConnectionAction { switch self.state { case .initialized: - preconditionFailure("Received a read event on a connection that was never opened.") - case .sslRequestSent: - return .read - case .sslNegotiated: - return .read - case .sslHandlerAdded: - return .read - case .waitingToStartAuthentication: - return .read - case .authenticating: - return .read - case .authenticated: - return .read - case .readyForQuery: + preconditionFailure("Invalid state: \(self.state). Read event before connection established?") + + case .sslRequestSent, + .sslNegotiated, + .sslHandlerAdded, + .waitingToStartAuthentication, + .authenticating, + .authenticated, + .readyForQuery, + .closing: + // all states in which we definitely want to make further forward progress... return .read + case .extendedQuery(var extendedQuery, let connectionContext): - return self.avoidingStateMachineCoW { machine in - let action = extendedQuery.readEventCaught() - machine.state = .extendedQuery(extendedQuery, connectionContext) - return machine.modify(with: action) - } - case .prepareStatement(var preparedStatement, let connectionContext): - return self.avoidingStateMachineCoW { machine in - let action = preparedStatement.readEventCaught() - machine.state = .prepareStatement(preparedStatement, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = extendedQuery.readEventCaught() + self.state = .extendedQuery(extendedQuery, connectionContext) + return self.modify(with: action) + case .closeCommand(var closeState, let connectionContext): - return self.avoidingStateMachineCoW { machine in - let action = closeState.readEventCaught() - machine.state = .closeCommand(closeState, connectionContext) - return machine.modify(with: action) - } - case .error: - return .read - case .closing: - return .read + self.state = .modifying // avoid CoW + let action = closeState.readEventCaught() + self.state = .closeCommand(closeState, connectionContext) + return self.modify(with: action) + case .closed: - preconditionFailure("How can we receive a read, if the connection is closed") + // Generally we shouldn't see this event (read after connection closed?!). + // But truth is, adopters run into this, again and again. So preconditioning here leads + // to unnecessary crashes. So let's be resilient and just make more forward progress. + // If we really care, we probably need to dive deep into PostgresNIO and SwiftNIO. + return .read + case .modifying: - preconditionFailure("Invalid state") + preconditionFailure("Invalid state: \(self.state)") } } @@ -704,17 +666,11 @@ struct ConnectionStateMachine { mutating func parseCompleteReceived() -> ConnectionAction { switch self.state { case .extendedQuery(var queryState, let connectionContext) where !queryState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.parseCompletedReceived() - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } - case .prepareStatement(var preparedState, let connectionContext) where !preparedState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = preparedState.parseCompletedReceived() - machine.state = .prepareStatement(preparedState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.parseCompletedReceived() + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) + default: return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.parseComplete)) } @@ -725,27 +681,20 @@ struct ConnectionStateMachine { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.bindComplete)) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.bindCompleteReceived() - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.bindCompleteReceived() + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) } mutating func parameterDescriptionReceived(_ description: PostgresBackendMessage.ParameterDescription) -> ConnectionAction { switch self.state { case .extendedQuery(var queryState, let connectionContext) where !queryState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.parameterDescriptionReceived(description) - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } - case .prepareStatement(var preparedState, let connectionContext) where !preparedState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = preparedState.parameterDescriptionReceived(description) - machine.state = .prepareStatement(preparedState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.parameterDescriptionReceived(description) + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) + default: return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.parameterDescription(description))) } @@ -754,17 +703,11 @@ struct ConnectionStateMachine { mutating func rowDescriptionReceived(_ description: RowDescription) -> ConnectionAction { switch self.state { case .extendedQuery(var queryState, let connectionContext) where !queryState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.rowDescriptionReceived(description) - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } - case .prepareStatement(var preparedState, let connectionContext) where !preparedState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = preparedState.rowDescriptionReceived(description) - machine.state = .prepareStatement(preparedState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.rowDescriptionReceived(description) + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) + default: return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.rowDescription(description))) } @@ -773,17 +716,11 @@ struct ConnectionStateMachine { mutating func noDataReceived() -> ConnectionAction { switch self.state { case .extendedQuery(var queryState, let connectionContext) where !queryState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.noDataReceived() - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } - case .prepareStatement(var preparedState, let connectionContext) where !preparedState.isComplete: - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = preparedState.noDataReceived() - machine.state = .prepareStatement(preparedState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.noDataReceived() + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) + default: return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.noData)) } @@ -798,11 +735,10 @@ struct ConnectionStateMachine { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.closeComplete)) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = closeState.closeCompletedReceived() - machine.state = .closeCommand(closeState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = closeState.closeCompletedReceived() + self.state = .closeCommand(closeState, connectionContext) + return self.modify(with: action) } mutating func commandCompletedReceived(_ commandTag: String) -> ConnectionAction { @@ -810,11 +746,10 @@ struct ConnectionStateMachine { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.commandComplete(commandTag))) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.commandCompletedReceived(commandTag) - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.commandCompletedReceived(commandTag) + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) } mutating func emptyQueryResponseReceived() -> ConnectionAction { @@ -822,11 +757,10 @@ struct ConnectionStateMachine { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.emptyQueryResponse)) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.emptyQueryResponseReceived() - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.emptyQueryResponseReceived() + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) } mutating func dataRowReceived(_ dataRow: DataRow) -> ConnectionAction { @@ -834,11 +768,10 @@ struct ConnectionStateMachine { return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.dataRow(dataRow))) } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.dataRowReceived(dataRow) - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.dataRowReceived(dataRow) + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) } // MARK: Consumer @@ -848,11 +781,10 @@ struct ConnectionStateMachine { preconditionFailure("Tried to cancel stream without active query") } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.cancel() - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.cancel() + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) } mutating func requestQueryRows() -> ConnectionAction { @@ -860,11 +792,10 @@ struct ConnectionStateMachine { preconditionFailure("Tried to consume next row, without active query") } - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - let action = queryState.requestQueryRows() - machine.state = .extendedQuery(queryState, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + let action = queryState.requestQueryRows() + self.state = .extendedQuery(queryState, connectionContext) + return self.modify(with: action) } // MARK: - Private Methods - @@ -874,15 +805,14 @@ struct ConnectionStateMachine { preconditionFailure("Can only start authentication after connect or ssl establish") } - return self.avoidingStateMachineCoW { machine in - var authState = AuthenticationStateMachine(authContext: authContext) - let action = authState.start() - machine.state = .authenticating(authState) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + var authState = AuthenticationStateMachine(authContext: authContext) + let action = authState.start() + self.state = .authenticating(authState) + return self.modify(with: action) } - private mutating func closeConnectionAndCleanup(_ error: PSQLError) -> ConnectionAction { + private mutating func closeConnectionAndCleanup(_ error: PSQLError, closePromise: EventLoopPromise? = nil) -> ConnectionAction { switch self.state { case .initialized, .sslRequestSent, @@ -891,12 +821,12 @@ struct ConnectionStateMachine { .waitingToStartAuthentication, .authenticated, .readyForQuery: - let cleanupContext = self.setErrorAndCreateCleanupContext(error) + let cleanupContext = self.setErrorAndCreateCleanupContext(error, closePromise: closePromise) return .closeConnectionAndCleanup(cleanupContext) case .authenticating(var authState): - let cleanupContext = self.setErrorAndCreateCleanupContext(error) - + let cleanupContext = self.setErrorAndCreateCleanupContext(error, closePromise: closePromise) + if authState.isComplete { // in case the auth state machine is complete all necessary actions have already // been forwarded to the consumer. We can close and cleanup without caring about the @@ -909,55 +839,46 @@ struct ConnectionStateMachine { preconditionFailure("Expect to fail auth") } return .closeConnectionAndCleanup(cleanupContext) + case .extendedQuery(var queryStateMachine, _): - let cleanupContext = self.setErrorAndCreateCleanupContext(error) - + let cleanupContext = self.setErrorAndCreateCleanupContext(error, closePromise: closePromise) + if queryStateMachine.isComplete { // in case the query state machine is complete all necessary actions have already // been forwarded to the consumer. We can close and cleanup without caring about the // substate machine. return .closeConnectionAndCleanup(cleanupContext) } - - switch queryStateMachine.errorHappened(error) { + + let action = queryStateMachine.errorHappened(error) + switch action { case .sendParseDescribeBindExecuteSync, + .sendParseDescribeSync, .sendBindExecuteSync, .succeedQuery, - .succeedQueryNoRowsComming, + .succeedPreparedStatementCreation, .forwardRows, .forwardStreamComplete, .wait, .read: - preconditionFailure("Expecting only failure actions if an error happened") + preconditionFailure("Invalid query state machine action in state: \(self.state), action: \(action)") + case .evaluateErrorAtConnectionLevel: return .closeConnectionAndCleanup(cleanupContext) + case .failQuery(let queryContext, with: let error): return .failQuery(queryContext, with: error, cleanupContext: cleanupContext) + case .forwardStreamError(let error, let read): return .forwardStreamError(error, read: read, cleanupContext: cleanupContext) + + case .failPreparedStatementCreation(let promise, with: let error): + return .failPreparedStatementCreation(promise, with: error, cleanupContext: cleanupContext) } - case .prepareStatement(var prepareStateMachine, _): - let cleanupContext = self.setErrorAndCreateCleanupContext(error) - - if prepareStateMachine.isComplete { - // in case the prepare state machine is complete all necessary actions have already - // been forwarded to the consumer. We can close and cleanup without caring about the - // substate machine. - return .closeConnectionAndCleanup(cleanupContext) - } - - switch prepareStateMachine.errorHappened(error) { - case .sendParseDescribeSync, - .succeedPreparedStatementCreation, - .read, - .wait: - preconditionFailure("Expecting only failure actions if an error happened") - case .failPreparedStatementCreation(let preparedStatementContext, with: let error): - return .failPreparedStatementCreation(preparedStatementContext, with: error, cleanupContext: cleanupContext) - } + case .closeCommand(var closeStateMachine, _): - let cleanupContext = self.setErrorAndCreateCleanupContext(error) - + let cleanupContext = self.setErrorAndCreateCleanupContext(error, closePromise: closePromise) + if closeStateMachine.isComplete { // in case the close state machine is complete all necessary actions have already // been forwarded to the consumer. We can close and cleanup without caring about the @@ -965,27 +886,27 @@ struct ConnectionStateMachine { return .closeConnectionAndCleanup(cleanupContext) } - switch closeStateMachine.errorHappened(error) { + let action = closeStateMachine.errorHappened(error) + switch action { case .sendCloseSync, .succeedClose, .read, .wait: - preconditionFailure("Expecting only failure actions if an error happened") + preconditionFailure("Invalid close state machine action in state: \(self.state), action: \(action)") case .failClose(let closeCommandContext, with: let error): return .failClose(closeCommandContext, with: error, cleanupContext: cleanupContext) } - case .error: - // TBD: this is an interesting case. why would this case happen? - let cleanupContext = self.setErrorAndCreateCleanupContext(error) - return .closeConnectionAndCleanup(cleanupContext) - - case .closing: - let cleanupContext = self.setErrorAndCreateCleanupContext(error) - return .closeConnectionAndCleanup(cleanupContext) - case .closed: - preconditionFailure("How can an error occur if the connection is already closed?") + + case .closing, .closed: + // We might run into this case because of reentrancy. For example: After we received an + // backend unexpected message, that we read of the wire, we bring this connection into + // the error state and will try to close the connection. However the server might have + // send further follow up messages. In those cases we will run into this method again + // and again. We should just ignore those events. + return .closeConnection(closePromise) + case .modifying: - preconditionFailure("Invalid state") + preconditionFailure("Invalid state: \(self.state)") } } @@ -1000,7 +921,7 @@ struct ConnectionStateMachine { // if we don't have anything left to do and we are quiescing, next we should close if case .quiescing(let promise) = self.quiescingState { - self.state = .closing + self.state = .closing(nil) return .closeConnection(promise) } @@ -1014,26 +935,18 @@ struct ConnectionStateMachine { switch task { case .extendedQuery(let queryContext): - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - var extendedQuery = ExtendedQueryStateMachine(queryContext: queryContext) - let action = extendedQuery.start() - machine.state = .extendedQuery(extendedQuery, connectionContext) - return machine.modify(with: action) - } - case .preparedStatement(let prepareContext): - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - var prepareStatement = PrepareStatementStateMachine(createContext: prepareContext) - let action = prepareStatement.start() - machine.state = .prepareStatement(prepareStatement, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + var extendedQuery = ExtendedQueryStateMachine(queryContext: queryContext) + let action = extendedQuery.start() + self.state = .extendedQuery(extendedQuery, connectionContext) + return self.modify(with: action) + case .closeCommand(let closeContext): - return self.avoidingStateMachineCoW { machine -> ConnectionAction in - var closeStateMachine = CloseStateMachine(closeContext: closeContext) - let action = closeStateMachine.start() - machine.state = .closeCommand(closeStateMachine, connectionContext) - return machine.modify(with: action) - } + self.state = .modifying // avoid CoW + var closeStateMachine = CloseStateMachine(closeContext: closeContext) + let action = closeStateMachine.start() + self.state = .closeCommand(closeStateMachine, connectionContext) + return self.modify(with: action) } } @@ -1042,43 +955,6 @@ struct ConnectionStateMachine { } } -// MARK: CoW helpers - -extension ConnectionStateMachine { - /// So, uh...this function needs some explaining. - /// - /// While the state machine logic above is great, there is a downside to having all of the state machine data in - /// associated data on enumerations: any modification of that data will trigger copy on write for heap-allocated - /// data. That means that for _every operation on the state machine_ we will CoW our underlying state, which is - /// not good. - /// - /// The way we can avoid this is by using this helper function. It will temporarily set state to a value with no - /// associated data, before attempting the body of the function. It will also verify that the state machine never - /// remains in this bad state. - /// - /// A key note here is that all callers must ensure that they return to a good state before they exit. - /// - /// Sadly, because it's generic and has a closure, we need to force it to be inlined at all call sites, which is - /// not ideal. - @inline(__always) - private mutating func avoidingStateMachineCoW(_ body: (inout ConnectionStateMachine) -> ReturnType) -> ReturnType { - self.state = .modifying - defer { - assert(!self.isModifying) - } - - return body(&self) - } - - private var isModifying: Bool { - if case .modifying = self.state { - return true - } else { - return false - } - } -} - extension ConnectionStateMachine { func shouldCloseConnection(reason error: PSQLError) -> Bool { switch error.code.base { @@ -1093,11 +969,12 @@ extension ConnectionStateMachine { .tooManyParameters, .invalidCommandTag, .connectionError, - .uncleanShutdown: + .uncleanShutdown, + .unlistenFailed: return true case .queryCancelled: return false - case .server: + case .server, .listenFailed: guard let sqlState = error.serverInfo?[.sqlState] else { // any error message that doesn't have a sql state field, is unexpected by default. return true @@ -1109,38 +986,45 @@ extension ConnectionStateMachine { } return false - case .connectionQuiescing: - preconditionFailure("Pure client error, that is thrown directly in PostgresConnection") - case .connectionClosed: - preconditionFailure("Pure client error, that is thrown directly and should never ") + case .clientClosedConnection, .poolClosed: + preconditionFailure("A pure client error was thrown directly in PostgresConnection, this shouldn't happen") + case .serverClosedConnection: + return true } } mutating func setErrorAndCreateCleanupContextIfNeeded(_ error: PSQLError) -> ConnectionAction.CleanUpContext? { - guard self.shouldCloseConnection(reason: error) else { - return nil + if self.shouldCloseConnection(reason: error) { + return self.setErrorAndCreateCleanupContext(error) } - return self.setErrorAndCreateCleanupContext(error) + return nil } - mutating func setErrorAndCreateCleanupContext(_ error: PSQLError) -> ConnectionAction.CleanUpContext { + mutating func setErrorAndCreateCleanupContext(_ error: PSQLError, closePromise: EventLoopPromise? = nil) -> ConnectionAction.CleanUpContext { let tasks = Array(self.taskQueue) self.taskQueue.removeAll() - var closePromise: EventLoopPromise? = nil - if case .quiescing(let promise) = self.quiescingState { - closePromise = promise + var forwardedPromise: EventLoopPromise? = nil + if case .quiescing(.some(let quiescePromise)) = self.quiescingState, let closePromise = closePromise { + quiescePromise.futureResult.cascade(to: closePromise) + forwardedPromise = quiescePromise + } else if case .quiescing(.some(let quiescePromise)) = self.quiescingState { + forwardedPromise = quiescePromise + } else { + forwardedPromise = closePromise } - - self.state = .error(error) - - var action = ConnectionAction.CleanUpContext.Action.close - if case .uncleanShutdown = error.code.base { + + let action: ConnectionAction.CleanUpContext.Action + if case .serverClosedConnection = error.code.base { + self.state = .closed(clientInitiated: false, error: error) action = .fireChannelInactive + } else { + self.state = .closing(error) + action = .close } - - return .init(action: action, tasks: tasks, error: error, closePromise: closePromise) + + return .init(action: action, tasks: tasks, error: error, closePromise: forwardedPromise) } } @@ -1154,10 +1038,8 @@ extension ConnectionStateMachine { case .failQuery(let requestContext, with: let error): let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error) return .failQuery(requestContext, with: error, cleanupContext: cleanupContext) - case .succeedQuery(let requestContext, columns: let columns): - return .succeedQuery(requestContext, columns: columns) - case .succeedQueryNoRowsComming(let requestContext, let commandTag): - return .succeedQueryNoRowsComming(requestContext, commandTag: commandTag) + case .succeedQuery(let requestContext, with: let result): + return .succeedQuery(requestContext, with: result) case .forwardRows(let buffer): return .forwardRows(buffer) case .forwardStreamComplete(let buffer, let commandTag): @@ -1175,24 +1057,13 @@ extension ConnectionStateMachine { return .read case .wait: return .wait - } - } -} - -extension ConnectionStateMachine { - mutating func modify(with action: PrepareStatementStateMachine.Action) -> ConnectionStateMachine.ConnectionAction { - switch action { - case .sendParseDescribeSync(let name, let query): - return .sendParseDescribeSync(name: name, query: query) - case .succeedPreparedStatementCreation(let prepareContext, with: let rowDescription): - return .succeedPreparedStatementCreation(prepareContext, with: rowDescription) - case .failPreparedStatementCreation(let prepareContext, with: let error): + case .sendParseDescribeSync(name: let name, query: let query, bindingDataTypes: let bindingDataTypes): + return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes) + case .succeedPreparedStatementCreation(let promise, with: let rowDescription): + return .succeedPreparedStatementCreation(promise, with: rowDescription) + case .failPreparedStatementCreation(let promise, with: let error): let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error) - return .failPreparedStatementCreation(prepareContext, with: error, cleanupContext: cleanupContext) - case .read: - return .read - case .wait: - return .wait + return .failPreparedStatementCreation(promise, with: error, cleanupContext: cleanupContext) } } } @@ -1243,11 +1114,19 @@ struct SendPrepareStatement { let query: String } -struct AuthContext: Equatable, CustomDebugStringConvertible { - let username: String - let password: String? - let database: String? - +struct AuthContext: CustomDebugStringConvertible { + var username: String + var password: String? + var database: String? + var additionalParameters: [(String, String)] + + init(username: String, password: String? = nil, database: String? = nil, additionalParameters: [(String, String)] = []) { + self.username = username + self.password = password + self.database = database + self.additionalParameters = additionalParameters + } + var debugDescription: String { """ AuthContext(username: \(String(reflecting: self.username)), \ @@ -1257,22 +1136,27 @@ struct AuthContext: Equatable, CustomDebugStringConvertible { } } -enum PasswordAuthencationMode: Equatable { - case cleartext - case md5(salt: (UInt8, UInt8, UInt8, UInt8)) - +extension AuthContext: Equatable { static func ==(lhs: Self, rhs: Self) -> Bool { - switch (lhs, rhs) { - case (.cleartext, .cleartext): - return true - case (.md5(let lhs), .md5(let rhs)): - return lhs == rhs - default: + guard lhs.username == rhs.username + && lhs.password == rhs.password + && lhs.database == rhs.database + && lhs.additionalParameters.count == rhs.additionalParameters.count + else { return false } + + return lhs.additionalParameters.elementsEqual(rhs.additionalParameters) { lhs, rhs in + lhs.0 == rhs.0 && lhs.1 == rhs.1 + } } } +enum PasswordAuthencationMode: Equatable { + case cleartext + case md5(salt: UInt32) +} + extension ConnectionStateMachine.State: CustomDebugStringConvertible { var debugDescription: String { switch self { @@ -1294,12 +1178,8 @@ extension ConnectionStateMachine.State: CustomDebugStringConvertible { return ".readyForQuery(connectionContext: \(String(reflecting: connectionContext)))" case .extendedQuery(let subStateMachine, let connectionContext): return ".extendedQuery(\(String(reflecting: subStateMachine)), connectionContext: \(String(reflecting: connectionContext)))" - case .prepareStatement(let subStateMachine, let connectionContext): - return ".prepareStatement(\(String(reflecting: subStateMachine)), connectionContext: \(String(reflecting: connectionContext)))" case .closeCommand(let subStateMachine, let connectionContext): return ".closeCommand(\(String(reflecting: subStateMachine)), connectionContext: \(String(reflecting: connectionContext)))" - case .error(let error): - return ".error(\(String(reflecting: error)))" case .closing: return ".closing" case .closed: diff --git a/Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift index 8b46fd0b..087a6c24 100644 --- a/Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift +++ b/Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift @@ -4,13 +4,14 @@ struct ExtendedQueryStateMachine { private enum State { case initialized(ExtendedQueryContext) - case parseDescribeBindExecuteSyncSent(ExtendedQueryContext) + case messagesSent(ExtendedQueryContext) case parseCompleteReceived(ExtendedQueryContext) case parameterDescriptionReceived(ExtendedQueryContext) case rowDescriptionReceived(ExtendedQueryContext, [RowDescription.Column]) case noDataMessageReceived(ExtendedQueryContext) - + case emptyQueryResponseReceived + /// A state that is used if a noData message was received before. If a row description was received `bufferingRows` is /// used after receiving a `bindComplete` message case bindCompleteReceived(ExtendedQueryContext) @@ -26,15 +27,18 @@ struct ExtendedQueryStateMachine { enum Action { case sendParseDescribeBindExecuteSync(PostgresQuery) + case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType]) case sendBindExecuteSync(PSQLExecuteStatement) // --- general actions - case failQuery(ExtendedQueryContext, with: PSQLError) - case succeedQuery(ExtendedQueryContext, columns: [RowDescription.Column]) - case succeedQueryNoRowsComming(ExtendedQueryContext, commandTag: String) + case failQuery(EventLoopPromise, with: PSQLError) + case succeedQuery(EventLoopPromise, with: QueryResult) case evaluateErrorAtConnectionLevel(PSQLError) + case succeedPreparedStatementCreation(EventLoopPromise, with: RowDescription?) + case failPreparedStatementCreation(EventLoopPromise, with: PSQLError) + // --- streaming actions // actions if query has requested next row but we are waiting for backend case forwardRows([DataRow]) @@ -59,13 +63,13 @@ struct ExtendedQueryStateMachine { } switch queryContext.query { - case .unnamed(let query): + case .unnamed(let query, _): return self.avoidingStateMachineCoW { state -> Action in - state = .parseDescribeBindExecuteSyncSent(queryContext) + state = .messagesSent(queryContext) return .sendParseDescribeBindExecuteSync(query) } - case .preparedStatement(let prepared): + case .executeStatement(let prepared, _): return self.avoidingStateMachineCoW { state -> Action in switch prepared.rowDescription { case .some(let rowDescription): @@ -75,6 +79,12 @@ struct ExtendedQueryStateMachine { } return .sendBindExecuteSync(prepared) } + + case .prepareStatement(let name, let query, let bindingDataTypes, _): + return self.avoidingStateMachineCoW { state -> Action in + state = .messagesSent(queryContext) + return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes) + } } } @@ -83,7 +93,7 @@ struct ExtendedQueryStateMachine { case .initialized: preconditionFailure("Start must be called immediatly after the query was created") - case .parseDescribeBindExecuteSyncSent(let queryContext), + case .messagesSent(let queryContext), .parseCompleteReceived(let queryContext), .parameterDescriptionReceived(let queryContext), .rowDescriptionReceived(let queryContext, _), @@ -94,7 +104,13 @@ struct ExtendedQueryStateMachine { } self.isCancelled = true - return .failQuery(queryContext, with: .queryCancelled) + switch queryContext.query { + case .unnamed(_, let eventLoopPromise), .executeStatement(_, let eventLoopPromise): + return .failQuery(eventLoopPromise, with: .queryCancelled) + + case .prepareStatement(_, _, _, let eventLoopPromise): + return .failPreparedStatementCreation(eventLoopPromise, with: .queryCancelled) + } case .streaming(let columns, var streamStateMachine): precondition(!self.isCancelled) @@ -107,7 +123,7 @@ struct ExtendedQueryStateMachine { return .forwardStreamError(.queryCancelled, read: true) } - case .commandComplete, .error, .drain: + case .commandComplete, .emptyQueryResponseReceived, .error, .drain: // the stream has already finished. return .wait @@ -117,7 +133,7 @@ struct ExtendedQueryStateMachine { } mutating func parseCompletedReceived() -> Action { - guard case .parseDescribeBindExecuteSyncSent(let queryContext) = self.state else { + guard case .messagesSent(let queryContext) = self.state else { return self.setAndFireError(.unexpectedBackendMessage(.parseComplete)) } @@ -143,9 +159,18 @@ struct ExtendedQueryStateMachine { return self.setAndFireError(.unexpectedBackendMessage(.noData)) } - return self.avoidingStateMachineCoW { state -> Action in - state = .noDataMessageReceived(queryContext) - return .wait + switch queryContext.query { + case .unnamed, .executeStatement: + return self.avoidingStateMachineCoW { state -> Action in + state = .noDataMessageReceived(queryContext) + return .wait + } + + case .prepareStatement(_, _, _, let promise): + return self.avoidingStateMachineCoW { state -> Action in + state = .noDataMessageReceived(queryContext) + return .succeedPreparedStatementCreation(promise, with: nil) + } } } @@ -153,42 +178,59 @@ struct ExtendedQueryStateMachine { guard case .parameterDescriptionReceived(let queryContext) = self.state else { return self.setAndFireError(.unexpectedBackendMessage(.rowDescription(rowDescription))) } - - return self.avoidingStateMachineCoW { state -> Action in - // In Postgres extended queries we receive the `rowDescription` before we send the - // `Bind` message. Well actually it's vice versa, but this is only true since we do - // pipelining during a query. - // - // In the actual protocol description we receive a rowDescription before the Bind - - // In Postgres extended queries we always request the response rows to be returned in - // `.binary` format. - let columns = rowDescription.columns.map { column -> RowDescription.Column in - var column = column - column.format = .binary - return column - } + + // In Postgres extended queries we receive the `rowDescription` before we send the + // `Bind` message. Well actually it's vice versa, but this is only true since we do + // pipelining during a query. + // + // In the actual protocol description we receive a rowDescription before the Bind + + // In Postgres extended queries we always request the response rows to be returned in + // `.binary` format. + let columns = rowDescription.columns.map { column -> RowDescription.Column in + var column = column + column.format = .binary + return column + } + + self.avoidingStateMachineCoW { state in state = .rowDescriptionReceived(queryContext, columns) + } + + switch queryContext.query { + case .unnamed, .executeStatement: return .wait + + case .prepareStatement(_, _, _, let eventLoopPromise): + return .succeedPreparedStatementCreation(eventLoopPromise, with: rowDescription) } } mutating func bindCompleteReceived() -> Action { switch self.state { - case .rowDescriptionReceived(let context, let columns): - return self.avoidingStateMachineCoW { state -> Action in - state = .streaming(columns, .init()) - return .succeedQuery(context, columns: columns) + case .rowDescriptionReceived(let queryContext, let columns): + switch queryContext.query { + case .unnamed(_, let eventLoopPromise), .executeStatement(_, let eventLoopPromise): + return self.avoidingStateMachineCoW { state -> Action in + state = .streaming(columns, .init()) + let result = QueryResult(value: .rowDescription(columns), logger: queryContext.logger) + return .succeedQuery(eventLoopPromise, with: result) + } + + case .prepareStatement: + return .evaluateErrorAtConnectionLevel(.unexpectedBackendMessage(.bindComplete)) } + case .noDataMessageReceived(let queryContext): return self.avoidingStateMachineCoW { state -> Action in state = .bindCompleteReceived(queryContext) return .wait } case .initialized, - .parseDescribeBindExecuteSyncSent, + .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, + .emptyQueryResponseReceived, .bindCompleteReceived, .streaming, .drain, @@ -224,10 +266,11 @@ struct ExtendedQueryStateMachine { return .wait case .initialized, - .parseDescribeBindExecuteSyncSent, + .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .noDataMessageReceived, + .emptyQueryResponseReceived, .rowDescriptionReceived, .bindCompleteReceived, .commandComplete, @@ -241,9 +284,16 @@ struct ExtendedQueryStateMachine { mutating func commandCompletedReceived(_ commandTag: String) -> Action { switch self.state { case .bindCompleteReceived(let context): - return self.avoidingStateMachineCoW { state -> Action in - state = .commandComplete(commandTag: commandTag) - return .succeedQueryNoRowsComming(context, commandTag: commandTag) + switch context.query { + case .unnamed(_, let eventLoopPromise), .executeStatement(_, let eventLoopPromise): + return self.avoidingStateMachineCoW { state -> Action in + state = .commandComplete(commandTag: commandTag) + let result = QueryResult(value: .noRows(.tag(commandTag)), logger: context.logger) + return .succeedQuery(eventLoopPromise, with: result) + } + + case .prepareStatement: + preconditionFailure("Invalid state: \(self.state)") } case .streaming(_, var demandStateMachine): @@ -258,10 +308,11 @@ struct ExtendedQueryStateMachine { return .wait case .initialized, - .parseDescribeBindExecuteSyncSent, + .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .noDataMessageReceived, + .emptyQueryResponseReceived, .rowDescriptionReceived, .commandComplete, .error: @@ -272,7 +323,22 @@ struct ExtendedQueryStateMachine { } mutating func emptyQueryResponseReceived() -> Action { - preconditionFailure("Unimplemented") + guard case .bindCompleteReceived(let queryContext) = self.state else { + return self.setAndFireError(.unexpectedBackendMessage(.emptyQueryResponse)) + } + + switch queryContext.query { + case .unnamed(_, let eventLoopPromise), + .executeStatement(_, let eventLoopPromise): + return self.avoidingStateMachineCoW { state -> Action in + state = .emptyQueryResponseReceived + let result = QueryResult(value: .noRows(.emptyResponse), logger: queryContext.logger) + return .succeedQuery(eventLoopPromise, with: result) + } + + case .prepareStatement(_, _, _, _): + return self.setAndFireError(.unexpectedBackendMessage(.emptyQueryResponse)) + } } mutating func errorReceived(_ errorMessage: PostgresBackendMessage.ErrorResponse) -> Action { @@ -280,7 +346,7 @@ struct ExtendedQueryStateMachine { switch self.state { case .initialized: return self.setAndFireError(.unexpectedBackendMessage(.error(errorMessage))) - case .parseDescribeBindExecuteSyncSent, + case .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .bindCompleteReceived: @@ -289,7 +355,7 @@ struct ExtendedQueryStateMachine { return self.setAndFireError(error) case .streaming, .drain: return self.setAndFireError(error) - case .commandComplete: + case .commandComplete, .emptyQueryResponseReceived: return self.setAndFireError(.unexpectedBackendMessage(.error(errorMessage))) case .error: preconditionFailure(""" @@ -331,10 +397,11 @@ struct ExtendedQueryStateMachine { return .wait case .initialized, - .parseDescribeBindExecuteSyncSent, + .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .noDataMessageReceived, + .emptyQueryResponseReceived, .rowDescriptionReceived, .bindCompleteReceived: preconditionFailure("Requested to consume next row without anything going on.") @@ -354,10 +421,11 @@ struct ExtendedQueryStateMachine { .commandComplete, .drain, .error, - .parseDescribeBindExecuteSyncSent, + .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .noDataMessageReceived, + .emptyQueryResponseReceived, .rowDescriptionReceived, .bindCompleteReceived: return .wait @@ -381,7 +449,7 @@ struct ExtendedQueryStateMachine { mutating func readEventCaught() -> Action { switch self.state { - case .parseDescribeBindExecuteSyncSent, + case .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .noDataMessageReceived, @@ -402,6 +470,7 @@ struct ExtendedQueryStateMachine { } case .initialized, .commandComplete, + .emptyQueryResponseReceived, .drain, .error: // we already have the complete stream received, now we are waiting for a @@ -417,7 +486,7 @@ struct ExtendedQueryStateMachine { private mutating func setAndFireError(_ error: PSQLError) -> Action { switch self.state { case .initialized(let context), - .parseDescribeBindExecuteSyncSent(let context), + .messagesSent(let context), .parseCompleteReceived(let context), .parameterDescriptionReceived(let context), .rowDescriptionReceived(let context, _), @@ -427,7 +496,12 @@ struct ExtendedQueryStateMachine { if self.isCancelled { return .evaluateErrorAtConnectionLevel(error) } else { - return .failQuery(context, with: error) + switch context.query { + case .unnamed(_, let eventLoopPromise), .executeStatement(_, let eventLoopPromise): + return .failQuery(eventLoopPromise, with: error) + case .prepareStatement(_, _, _, let eventLoopPromise): + return .failPreparedStatementCreation(eventLoopPromise, with: error) + } } case .drain: @@ -443,7 +517,7 @@ struct ExtendedQueryStateMachine { return .forwardStreamError(error, read: true) } - case .commandComplete, .error: + case .commandComplete, .emptyQueryResponseReceived, .error: preconditionFailure(""" This state must not be reached. If the query `.isComplete`, the ConnectionStateMachine must not send any further events to the substate machine. @@ -455,11 +529,22 @@ struct ExtendedQueryStateMachine { var isComplete: Bool { switch self.state { - case .commandComplete, - .error: + case .commandComplete, .emptyQueryResponseReceived, .error: return true - default: + + case .noDataMessageReceived(let context), .rowDescriptionReceived(let context, _): + switch context.query { + case .prepareStatement: + return true + case .unnamed, .executeStatement: + return false + } + + case .initialized, .messagesSent, .parseCompleteReceived, .parameterDescriptionReceived, .bindCompleteReceived, .streaming, .drain: return false + + case .modifying: + preconditionFailure("Invalid state: \(self.state)") } } } diff --git a/Sources/PostgresNIO/New/Connection State Machine/ListenStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/ListenStateMachine.swift new file mode 100644 index 00000000..89f40469 --- /dev/null +++ b/Sources/PostgresNIO/New/Connection State Machine/ListenStateMachine.swift @@ -0,0 +1,254 @@ +import NIOCore + +struct ListenStateMachine { + var channels: [String: ChannelState] + + init() { + self.channels = [:] + } + + enum StartListeningAction { + case none + case startListening(String) + case succeedListenStart(NotificationListener) + } + + mutating func startListening(_ new: NotificationListener) -> StartListeningAction { + return self.channels[new.channel, default: .init()].start(new) + } + + enum StartListeningSuccessAction { + case stopListening + case activateListeners(Dictionary.Values) + } + + mutating func startListeningSucceeded(channel: String) -> StartListeningSuccessAction { + return self.channels[channel]!.startListeningSucceeded() + } + + mutating func startListeningFailed(channel: String, error: Error) -> Dictionary.Values { + return self.channels[channel]!.startListeningFailed(error) + } + + enum StopListeningSuccessAction { + case startListening + case none + } + + mutating func stopListeningSucceeded(channel: String) -> StopListeningSuccessAction { + switch self.channels[channel]!.stopListeningSucceeded() { + case .none: + self.channels.removeValue(forKey: channel) + return .none + + case .startListening: + return .startListening + } + } + + enum CancelAction { + case stopListening(String, cancelListener: NotificationListener) + case cancelListener(NotificationListener) + case none + } + + mutating func cancelNotificationListener(channel: String, id: Int) -> CancelAction { + return self.channels[channel]?.cancelListening(id: id) ?? .none + } + + mutating func fail(_ error: Error) -> [NotificationListener] { + var result = [NotificationListener]() + while var (_, channel) = self.channels.popFirst() { + switch channel.fail(error) { + case .none: + continue + + case .failListeners(let listeners): + result.append(contentsOf: listeners) + } + } + return result + } + + enum ReceivedAction { + case none + case notify(Dictionary.Values) + } + + func notificationReceived(channel: String) -> ReceivedAction { + // TODO: Do we want to close the connection, if we receive a notification on a channel that we don't listen to? + // We can only change this with the next major release, as it would break current functionality. + return self.channels[channel]?.notificationReceived() ?? .none + } +} + +extension ListenStateMachine { + struct ChannelState { + enum State { + case initialized + case starting([Int: NotificationListener]) + case listening([Int: NotificationListener]) + case stopping([Int: NotificationListener]) + case failed(Error) + } + + private var state: State + + init() { + self.state = .initialized + } + + mutating func start(_ new: NotificationListener) -> StartListeningAction { + switch self.state { + case .initialized: + self.state = .starting([new.id: new]) + return .startListening(new.channel) + + case .starting(var listeners): + listeners[new.id] = new + self.state = .starting(listeners) + return .none + + case .listening(var listeners): + listeners[new.id] = new + self.state = .listening(listeners) + return .succeedListenStart(new) + + case .stopping(var listeners): + listeners[new.id] = new + self.state = .stopping(listeners) + return .none + + case .failed: + fatalError("Invalid state: \(self.state)") + } + } + + mutating func startListeningSucceeded() -> StartListeningSuccessAction { + switch self.state { + case .initialized, .listening, .stopping: + fatalError("Invalid state: \(self.state)") + + case .starting(let listeners): + if listeners.isEmpty { + self.state = .stopping(listeners) + return .stopListening + } else { + self.state = .listening(listeners) + return .activateListeners(listeners.values) + } + + case .failed: + fatalError("Invalid state: \(self.state)") + } + } + + mutating func startListeningFailed(_ error: Error) -> Dictionary.Values { + switch self.state { + case .initialized, .listening, .stopping: + fatalError("Invalid state: \(self.state)") + + case .starting(let listeners): + self.state = .initialized + return listeners.values + + case .failed: + fatalError("Invalid state: \(self.state)") + } + } + + mutating func stopListeningSucceeded() -> StopListeningSuccessAction { + switch self.state { + case .initialized, .listening, .starting: + fatalError("Invalid state: \(self.state)") + + case .stopping(let listeners): + if listeners.isEmpty { + self.state = .initialized + return .none + } else { + self.state = .starting(listeners) + return .startListening + } + + case .failed: + return .none + } + } + + mutating func cancelListening(id: Int) -> CancelAction { + switch self.state { + case .initialized: + fatalError("Invalid state: \(self.state)") + + case .starting(var listeners): + let removed = listeners.removeValue(forKey: id) + self.state = .starting(listeners) + if let removed = removed { + return .cancelListener(removed) + } + return .none + + case .listening(var listeners): + precondition(!listeners.isEmpty) + let maybeLast = listeners.removeValue(forKey: id) + if let last = maybeLast, listeners.isEmpty { + self.state = .stopping(listeners) + return .stopListening(last.channel, cancelListener: last) + } else { + self.state = .listening(listeners) + if let notLast = maybeLast { + return .cancelListener(notLast) + } + return .none + } + + case .stopping(var listeners): + let removed = listeners.removeValue(forKey: id) + self.state = .stopping(listeners) + if let removed = removed { + return .cancelListener(removed) + } + return .none + + case .failed: + return .none + } + } + + enum FailAction { + case failListeners(Dictionary.Values) + case none + } + + mutating func fail(_ error: Error) -> FailAction { + switch self.state { + case .initialized: + fatalError("Invalid state: \(self.state)") + + case .starting(let listeners), .listening(let listeners), .stopping(let listeners): + self.state = .failed(error) + return .failListeners(listeners.values) + + case .failed: + return .none + } + } + + func notificationReceived() -> ReceivedAction { + switch self.state { + case .initialized, .starting: + fatalError("Invalid state: \(self.state)") + + case .listening(let listeners): + return .notify(listeners.values) + + case .stopping: + return .none + + default: + preconditionFailure("TODO: Implemented") + } + } + } +} diff --git a/Sources/PostgresNIO/New/Connection State Machine/PrepareStatementStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/PrepareStatementStateMachine.swift deleted file mode 100644 index 5b65fc90..00000000 --- a/Sources/PostgresNIO/New/Connection State Machine/PrepareStatementStateMachine.swift +++ /dev/null @@ -1,147 +0,0 @@ - -struct PrepareStatementStateMachine { - - enum State { - case initialized(PrepareStatementContext) - case parseDescribeSent(PrepareStatementContext) - - case parseCompleteReceived(PrepareStatementContext) - case parameterDescriptionReceived(PrepareStatementContext) - case rowDescriptionReceived - case noDataMessageReceived - - case error(PSQLError) - } - - enum Action { - case sendParseDescribeSync(name: String, query: String) - case succeedPreparedStatementCreation(PrepareStatementContext, with: RowDescription?) - case failPreparedStatementCreation(PrepareStatementContext, with: PSQLError) - - case read - case wait - } - - var state: State - - init(createContext: PrepareStatementContext) { - self.state = .initialized(createContext) - } - - #if DEBUG - /// for testing purposes only - init(_ state: State) { - self.state = state - } - #endif - - mutating func start() -> Action { - guard case .initialized(let createContext) = self.state else { - preconditionFailure("Start must only be called after the query has been initialized") - } - - self.state = .parseDescribeSent(createContext) - - return .sendParseDescribeSync(name: createContext.name, query: createContext.query) - } - - mutating func parseCompletedReceived() -> Action { - guard case .parseDescribeSent(let createContext) = self.state else { - return self.setAndFireError(.unexpectedBackendMessage(.parseComplete)) - } - - self.state = .parseCompleteReceived(createContext) - return .wait - } - - mutating func parameterDescriptionReceived(_ parameterDescription: PostgresBackendMessage.ParameterDescription) -> Action { - guard case .parseCompleteReceived(let createContext) = self.state else { - return self.setAndFireError(.unexpectedBackendMessage(.parameterDescription(parameterDescription))) - } - - self.state = .parameterDescriptionReceived(createContext) - return .wait - } - - mutating func noDataReceived() -> Action { - guard case .parameterDescriptionReceived(let queryContext) = self.state else { - return self.setAndFireError(.unexpectedBackendMessage(.noData)) - } - - self.state = .noDataMessageReceived - return .succeedPreparedStatementCreation(queryContext, with: nil) - } - - mutating func rowDescriptionReceived(_ rowDescription: RowDescription) -> Action { - guard case .parameterDescriptionReceived(let queryContext) = self.state else { - return self.setAndFireError(.unexpectedBackendMessage(.rowDescription(rowDescription))) - } - - self.state = .rowDescriptionReceived - return .succeedPreparedStatementCreation(queryContext, with: rowDescription) - } - - mutating func errorReceived(_ errorMessage: PostgresBackendMessage.ErrorResponse) -> Action { - let error = PSQLError.server(errorMessage) - switch self.state { - case .initialized: - return self.setAndFireError(.unexpectedBackendMessage(.error(errorMessage))) - - case .parseDescribeSent, - .parseCompleteReceived, - .parameterDescriptionReceived: - return self.setAndFireError(error) - - case .rowDescriptionReceived, - .noDataMessageReceived, - .error: - preconditionFailure(""" - This state must not be reached. If the prepared statement `.isComplete`, the - ConnectionStateMachine must not send any further events to the substate machine. - """) - } - } - - mutating func errorHappened(_ error: PSQLError) -> Action { - return self.setAndFireError(error) - } - - private mutating func setAndFireError(_ error: PSQLError) -> Action { - switch self.state { - case .initialized(let context), - .parseDescribeSent(let context), - .parseCompleteReceived(let context), - .parameterDescriptionReceived(let context): - self.state = .error(error) - return .failPreparedStatementCreation(context, with: error) - case .rowDescriptionReceived, - .noDataMessageReceived, - .error: - preconditionFailure(""" - This state must not be reached. If the prepared statement `.isComplete`, the - ConnectionStateMachine must not send any further events to the substate machine. - """) - } - } - - // MARK: Channel actions - - mutating func readEventCaught() -> Action { - return .read - } - - var isComplete: Bool { - switch self.state { - case .rowDescriptionReceived, - .noDataMessageReceived, - .error: - return true - case .initialized, - .parseDescribeSent, - .parseCompleteReceived, - .parameterDescriptionReceived: - return false - } - } - -} diff --git a/Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift new file mode 100644 index 00000000..5afa4d0b --- /dev/null +++ b/Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift @@ -0,0 +1,93 @@ +import NIOCore + +struct PreparedStatementStateMachine { + enum State { + case preparing([PreparedStatementContext]) + case prepared(RowDescription?) + case error(PSQLError) + } + + var preparedStatements: [String: State] = [:] + + enum LookupAction { + case prepareStatement + case waitForAlreadyInFlightPreparation + case executeStatement(RowDescription?) + case returnError(PSQLError) + } + + mutating func lookup(preparedStatement: PreparedStatementContext) -> LookupAction { + if let state = self.preparedStatements[preparedStatement.name] { + switch state { + case .preparing(var statements): + statements.append(preparedStatement) + self.preparedStatements[preparedStatement.name] = .preparing(statements) + return .waitForAlreadyInFlightPreparation + case .prepared(let rowDescription): + return .executeStatement(rowDescription) + case .error(let error): + return .returnError(error) + } + } else { + self.preparedStatements[preparedStatement.name] = .preparing([preparedStatement]) + return .prepareStatement + } + } + + struct PreparationCompleteAction { + var statements: [PreparedStatementContext] + var rowDescription: RowDescription? + } + + mutating func preparationComplete( + name: String, + rowDescription: RowDescription? + ) -> PreparationCompleteAction { + guard let state = self.preparedStatements[name] else { + fatalError("Unknown prepared statement \(name)") + } + switch state { + case .preparing(let statements): + // When sending the bindings we are going to ask for binary data. + if var rowDescription = rowDescription { + for i in 0.. ErrorHappenedAction { + guard let state = self.preparedStatements[name] else { + fatalError("Unknown prepared statement \(name)") + } + switch state { + case .preparing(let statements): + self.preparedStatements[name] = .error(error) + return ErrorHappenedAction( + statements: statements, + error: error + ) + case .prepared, .error: + preconditionFailure("Error happened in an unexpected state \(state)") + } + } +} diff --git a/Sources/PostgresNIO/New/Data/Array+PostgresCodable.swift b/Sources/PostgresNIO/New/Data/Array+PostgresCodable.swift index 2c57b605..ddab0fff 100644 --- a/Sources/PostgresNIO/New/Data/Array+PostgresCodable.swift +++ b/Sources/PostgresNIO/New/Data/Array+PostgresCodable.swift @@ -1,4 +1,5 @@ import NIOCore +import struct Foundation.Date import struct Foundation.UUID // MARK: Protocols @@ -85,6 +86,24 @@ extension UUID: PostgresArrayEncodable { public static var psqlArrayType: PostgresDataType { .uuidArray } } +extension Date: PostgresArrayDecodable {} + +extension Date: PostgresArrayEncodable { + public static var psqlArrayType: PostgresDataType { .timestamptzArray } +} + +extension Range: PostgresArrayDecodable where Bound: PostgresRangeArrayDecodable {} + +extension Range: PostgresArrayEncodable where Bound: PostgresRangeArrayEncodable { + public static var psqlArrayType: PostgresDataType { Bound.psqlRangeArrayType } +} + +extension ClosedRange: PostgresArrayDecodable where Bound: PostgresRangeArrayDecodable {} + +extension ClosedRange: PostgresArrayEncodable where Bound: PostgresRangeArrayEncodable { + public static var psqlArrayType: PostgresDataType { Bound.psqlRangeArrayType } +} + // MARK: Array conformances extension Array: PostgresEncodable where Element: PostgresArrayEncodable { @@ -124,6 +143,10 @@ extension Array: PostgresEncodable where Element: PostgresArrayEncodable { } } +// explicitly conforming to PostgresThrowingDynamicTypeEncodable because of: +// https://github.com/apple/swift/issues/54132 +extension Array: PostgresThrowingDynamicTypeEncodable where Element: PostgresArrayEncodable {} + extension Array: PostgresNonThrowingEncodable where Element: PostgresArrayEncodable & PostgresNonThrowingEncodable { public static var psqlType: PostgresDataType { Element.psqlArrayType @@ -161,6 +184,9 @@ extension Array: PostgresNonThrowingEncodable where Element: PostgresArrayEncoda } } +// explicitly conforming to PostgresDynamicTypeEncodable because of: +// https://github.com/apple/swift/issues/54132 +extension Array: PostgresDynamicTypeEncodable where Element: PostgresArrayEncodable & PostgresNonThrowingEncodable {} extension Array: PostgresDecodable where Element: PostgresArrayDecodable, Element == Element._DecodableType { public init( diff --git a/Sources/PostgresNIO/New/Data/Range+PostgresCodable.swift b/Sources/PostgresNIO/New/Data/Range+PostgresCodable.swift new file mode 100644 index 00000000..6279cf4b --- /dev/null +++ b/Sources/PostgresNIO/New/Data/Range+PostgresCodable.swift @@ -0,0 +1,325 @@ +import NIOCore + +// MARK: Protocols + +/// A type that can be encoded into a Postgres range type where it is the bound type +public protocol PostgresRangeEncodable: PostgresNonThrowingEncodable { + static var psqlRangeType: PostgresDataType { get } +} + +/// A type that can be decoded into a Swift RangeExpression type from a Postgres range where it is the bound type +public protocol PostgresRangeDecodable: PostgresDecodable { + /// If a Postgres range type has a well-defined step, + /// Postgres automatically converts it to a canonical form. + /// Types such as `int4range` get converted to upper-bound-exclusive. + /// This method is needed when converting an upper bound to inclusive. + /// It should throw if the type lacks a well-defined step. + func upperBoundExclusiveToUpperBoundInclusive() throws -> Self + + /// Postgres does not store any bound values for empty ranges, + /// but Swift requires a value to initialize an empty Range. + static var valueForEmptyRange: Self { get } +} + +/// A type that can be encoded into a Postgres range array type where it is the bound type +public protocol PostgresRangeArrayEncodable: PostgresRangeEncodable { + static var psqlRangeArrayType: PostgresDataType { get } +} + +/// A type that can be decoded into a Swift RangeExpression array type from a Postgres range array where it is the bound type +public protocol PostgresRangeArrayDecodable: PostgresRangeDecodable {} + +// MARK: Bound conformances + +extension FixedWidthInteger where Self: PostgresRangeDecodable { + public func upperBoundExclusiveToUpperBoundInclusive() -> Self { + return self - 1 + } + + public static var valueForEmptyRange: Self { + return .zero + } +} + +extension Int32: PostgresRangeEncodable { + public static var psqlRangeType: PostgresDataType { return .int4Range } +} + +extension Int32: PostgresRangeDecodable {} + +extension Int32: PostgresRangeArrayEncodable { + public static var psqlRangeArrayType: PostgresDataType { return .int4RangeArray } +} + +extension Int32: PostgresRangeArrayDecodable {} + +extension Int64: PostgresRangeEncodable { + public static var psqlRangeType: PostgresDataType { return .int8Range } +} + +extension Int64: PostgresRangeDecodable {} + +extension Int64: PostgresRangeArrayEncodable { + public static var psqlRangeArrayType: PostgresDataType { return .int8RangeArray } +} + +extension Int64: PostgresRangeArrayDecodable {} + +// MARK: PostgresRange + +@usableFromInline +struct PostgresRange { + @usableFromInline let lowerBound: Bound? + @usableFromInline let upperBound: Bound? + @usableFromInline let isLowerBoundInclusive: Bool + @usableFromInline let isUpperBoundInclusive: Bool + + @inlinable + init( + lowerBound: Bound?, + upperBound: Bound?, + isLowerBoundInclusive: Bool, + isUpperBoundInclusive: Bool + ) { + self.lowerBound = lowerBound + self.upperBound = upperBound + self.isLowerBoundInclusive = isLowerBoundInclusive + self.isUpperBoundInclusive = isUpperBoundInclusive + } +} + +/// Used by Postgres to represent certain range properties +@usableFromInline +struct PostgresRangeFlag { + @usableFromInline static let isEmpty: UInt8 = 0x01 + @usableFromInline static let isLowerBoundInclusive: UInt8 = 0x02 + @usableFromInline static let isUpperBoundInclusive: UInt8 = 0x04 +} + +extension PostgresRange: PostgresDecodable where Bound: PostgresRangeDecodable { + @inlinable + init( + from byteBuffer: inout ByteBuffer, + type: PostgresDataType, + format: PostgresFormat, + context: PostgresDecodingContext + ) throws { + guard case .binary = format else { + throw PostgresDecodingError.Code.failure + } + + guard let boundType: PostgresDataType = type.boundType else { + throw PostgresDecodingError.Code.failure + } + + // flags byte contains certain properties of the range + guard let flags: UInt8 = byteBuffer.readInteger(as: UInt8.self) else { + throw PostgresDecodingError.Code.failure + } + + let isEmpty: Bool = flags & PostgresRangeFlag.isEmpty != 0 + if isEmpty { + self = PostgresRange( + lowerBound: Bound.valueForEmptyRange, + upperBound: Bound.valueForEmptyRange, + isLowerBoundInclusive: true, + isUpperBoundInclusive: false + ) + return + } + + guard let lowerBoundSize: Int32 = byteBuffer.readInteger(as: Int32.self), + Int(lowerBoundSize) == MemoryLayout.size, + var lowerBoundBytes: ByteBuffer = byteBuffer.readSlice(length: Int(lowerBoundSize)) + else { + throw PostgresDecodingError.Code.failure + } + + let lowerBound = try Bound(from: &lowerBoundBytes, type: boundType, format: format, context: context) + + guard let upperBoundSize = byteBuffer.readInteger(as: Int32.self), + Int(upperBoundSize) == MemoryLayout.size, + var upperBoundBytes: ByteBuffer = byteBuffer.readSlice(length: Int(upperBoundSize)) + else { + throw PostgresDecodingError.Code.failure + } + + let upperBound = try Bound(from: &upperBoundBytes, type: boundType, format: format, context: context) + + let isLowerBoundInclusive: Bool = flags & PostgresRangeFlag.isLowerBoundInclusive != 0 + let isUpperBoundInclusive: Bool = flags & PostgresRangeFlag.isUpperBoundInclusive != 0 + + self = PostgresRange( + lowerBound: lowerBound, + upperBound: upperBound, + isLowerBoundInclusive: isLowerBoundInclusive, + isUpperBoundInclusive: isUpperBoundInclusive + ) + + } +} + +extension PostgresRange: PostgresEncodable & PostgresNonThrowingEncodable where Bound: PostgresRangeEncodable { + @usableFromInline + static var psqlType: PostgresDataType { return Bound.psqlRangeType } + + @usableFromInline + static var psqlFormat: PostgresFormat { return .binary } + + @inlinable + func encode(into byteBuffer: inout ByteBuffer, context: PostgresEncodingContext) { + // flags byte contains certain properties of the range + var flags: UInt8 = 0 + if self.isLowerBoundInclusive { + flags |= PostgresRangeFlag.isLowerBoundInclusive + } + if self.isUpperBoundInclusive { + flags |= PostgresRangeFlag.isUpperBoundInclusive + } + + let boundMemorySize = Int32(MemoryLayout.size) + + byteBuffer.writeInteger(flags) + if let lowerBound = self.lowerBound { + byteBuffer.writeInteger(boundMemorySize) + lowerBound.encode(into: &byteBuffer, context: context) + } + if let upperBound = self.upperBound { + byteBuffer.writeInteger(boundMemorySize) + upperBound.encode(into: &byteBuffer, context: context) + } + } +} + +// explicitly conforming to PostgresDynamicTypeEncodable and PostgresThrowingDynamicTypeEncodable because of: +// https://github.com/apple/swift/issues/54132 +extension PostgresRange: PostgresThrowingDynamicTypeEncodable & PostgresDynamicTypeEncodable + where Bound: PostgresRangeEncodable {} + +extension PostgresRange where Bound: Comparable { + @inlinable + init(range: Range) { + self.lowerBound = range.lowerBound + self.upperBound = range.upperBound + self.isLowerBoundInclusive = true + self.isUpperBoundInclusive = false + } + + @inlinable + init(closedRange: ClosedRange) { + self.lowerBound = closedRange.lowerBound + self.upperBound = closedRange.upperBound + self.isLowerBoundInclusive = true + self.isUpperBoundInclusive = true + } +} + +// MARK: Range + +extension Range: PostgresEncodable where Bound: PostgresRangeEncodable { + public static var psqlType: PostgresDataType { return Bound.psqlRangeType } + public static var psqlFormat: PostgresFormat { return .binary } + + @inlinable + public func encode( + into byteBuffer: inout ByteBuffer, + context: PostgresEncodingContext + ) { + let postgresRange = PostgresRange(range: self) + postgresRange.encode(into: &byteBuffer, context: context) + } +} + +extension Range: PostgresNonThrowingEncodable where Bound: PostgresRangeEncodable {} + +// explicitly conforming to PostgresDynamicTypeEncodable and PostgresThrowingDynamicTypeEncodable because of: +// https://github.com/apple/swift/issues/54132 +extension Range: PostgresDynamicTypeEncodable & PostgresThrowingDynamicTypeEncodable + where Bound: PostgresRangeEncodable {} + +extension Range: PostgresDecodable where Bound: PostgresRangeDecodable { + @inlinable + public init( + from buffer: inout ByteBuffer, + type: PostgresDataType, + format: PostgresFormat, + context: PostgresDecodingContext + ) throws { + let postgresRange = try PostgresRange( + from: &buffer, + type: type, + format: format, + context: context + ) + + guard let lowerBound: Bound = postgresRange.lowerBound, + let upperBound: Bound = postgresRange.upperBound, + postgresRange.isLowerBoundInclusive, + !postgresRange.isUpperBoundInclusive + else { + throw PostgresDecodingError.Code.failure + } + + self = lowerBound..( + into byteBuffer: inout ByteBuffer, + context: PostgresEncodingContext + ) { + let postgresRange = PostgresRange(closedRange: self) + postgresRange.encode(into: &byteBuffer, context: context) + } +} + +// explicitly conforming to PostgresThrowingDynamicTypeEncodable because of: +// https://github.com/apple/swift/issues/54132 +extension ClosedRange: PostgresThrowingDynamicTypeEncodable where Bound: PostgresRangeEncodable {} + +extension ClosedRange: PostgresNonThrowingEncodable where Bound: PostgresRangeEncodable {} + +// explicitly conforming to PostgresDynamicTypeEncodable because of: +// https://github.com/apple/swift/issues/54132 +extension ClosedRange: PostgresDynamicTypeEncodable where Bound: PostgresRangeEncodable {} + +extension ClosedRange: PostgresDecodable where Bound: PostgresRangeDecodable { + @inlinable + public init( + from buffer: inout ByteBuffer, + type: PostgresDataType, + format: PostgresFormat, + context: PostgresDecodingContext + ) throws { + let postgresRange = try PostgresRange( + from: &buffer, + type: type, + format: format, + context: context + ) + + guard let lowerBound: Bound = postgresRange.lowerBound, + var upperBound: Bound = postgresRange.upperBound, + postgresRange.isLowerBoundInclusive + else { + throw PostgresDecodingError.Code.failure + } + + if !postgresRange.isUpperBoundInclusive { + upperBound = try upperBound.upperBoundExclusiveToUpperBoundInclusive() + } + + if lowerBound > upperBound { + throw PostgresDecodingError.Code.failure + } + + self = lowerBound...upperBound + } +} diff --git a/Sources/PostgresNIO/New/Data/RawRepresentable+PostgresCodable.swift b/Sources/PostgresNIO/New/Data/RawRepresentable+PostgresCodable.swift index 4d6c20c4..ea097963 100644 --- a/Sources/PostgresNIO/New/Data/RawRepresentable+PostgresCodable.swift +++ b/Sources/PostgresNIO/New/Data/RawRepresentable+PostgresCodable.swift @@ -19,7 +19,7 @@ extension PostgresEncodable where Self: RawRepresentable, RawValue: PostgresEnco } extension PostgresDecodable where Self: RawRepresentable, RawValue: PostgresDecodable, RawValue._DecodableType == RawValue { - init( + public init( from buffer: inout ByteBuffer, type: PostgresDataType, format: PostgresFormat, diff --git a/Sources/PostgresNIO/New/Data/String+PostgresCodable.swift b/Sources/PostgresNIO/New/Data/String+PostgresCodable.swift index f8e93e94..41091ab3 100644 --- a/Sources/PostgresNIO/New/Data/String+PostgresCodable.swift +++ b/Sources/PostgresNIO/New/Data/String+PostgresCodable.swift @@ -30,6 +30,7 @@ extension String: PostgresDecodable { ) throws { switch (format, type) { case (_, .varchar), + (_, .bpchar), (_, .text), (_, .name): // we can force unwrap here, since this method only fails if there are not enough diff --git a/Sources/PostgresNIO/New/Extensions/ByteBuffer+PSQL.swift b/Sources/PostgresNIO/New/Extensions/ByteBuffer+PSQL.swift index 6d632b6f..838e624d 100644 --- a/Sources/PostgresNIO/New/Extensions/ByteBuffer+PSQL.swift +++ b/Sources/PostgresNIO/New/Extensions/ByteBuffer+PSQL.swift @@ -2,14 +2,6 @@ import NIOCore internal extension ByteBuffer { - mutating func psqlWriteBackendMessageID(_ messageID: PostgresBackendMessage.ID) { - self.writeInteger(messageID.rawValue) - } - - mutating func psqlWriteFrontendMessageID(_ messageID: PostgresFrontendMessage.ID) { - self.writeInteger(messageID.rawValue) - } - @usableFromInline mutating func psqlReadFloat() -> Float? { return self.readInteger(as: UInt32.self).map { Float(bitPattern: $0) } diff --git a/Sources/PostgresNIO/New/Messages/Authentication.swift b/Sources/PostgresNIO/New/Messages/Authentication.swift index bd0d2e57..eff62e91 100644 --- a/Sources/PostgresNIO/New/Messages/Authentication.swift +++ b/Sources/PostgresNIO/New/Messages/Authentication.swift @@ -2,10 +2,10 @@ import NIOCore extension PostgresBackendMessage { - enum Authentication: PayloadDecodable { + enum Authentication: PayloadDecodable, Hashable { case ok case kerberosV5 - case md5(salt: (UInt8, UInt8, UInt8, UInt8)) + case md5(salt: UInt32) case plaintext case scmCredential case gss @@ -26,7 +26,7 @@ extension PostgresBackendMessage { case 3: return .plaintext case 5: - guard let salt = buffer.readMultipleIntegers(endianness: .big, as: (UInt8, UInt8, UInt8, UInt8).self) else { + guard let salt = buffer.readInteger(as: UInt32.self) else { throw PSQLPartialDecodingError.expectedAtLeastNRemainingBytes(4, actual: buffer.readableBytes) } return .md5(salt: salt) @@ -61,37 +61,6 @@ extension PostgresBackendMessage { } } -extension PostgresBackendMessage.Authentication: Equatable { - static func ==(lhs: Self, rhs: Self) -> Bool { - switch (lhs, rhs) { - case (.ok, .ok): - return true - case (.kerberosV5, .kerberosV5): - return true - case (.md5(let lhs), .md5(let rhs)): - return lhs == rhs - case (.plaintext, .plaintext): - return true - case (.scmCredential, .scmCredential): - return true - case (.gss, .gss): - return true - case (.sspi, .sspi): - return true - case (.gssContinue(let lhs), .gssContinue(let rhs)): - return lhs == rhs - case (.sasl(let lhs), .sasl(let rhs)): - return lhs == rhs - case (.saslContinue(let lhs), .saslContinue(let rhs)): - return lhs == rhs - case (.saslFinal(let lhs), .saslFinal(let rhs)): - return lhs == rhs - default: - return false - } - } -} - extension PostgresBackendMessage.Authentication: CustomDebugStringConvertible { var debugDescription: String { switch self { diff --git a/Sources/PostgresNIO/New/Messages/BackendKeyData.swift b/Sources/PostgresNIO/New/Messages/BackendKeyData.swift index 498c5110..31a676d2 100644 --- a/Sources/PostgresNIO/New/Messages/BackendKeyData.swift +++ b/Sources/PostgresNIO/New/Messages/BackendKeyData.swift @@ -2,7 +2,7 @@ import NIOCore extension PostgresBackendMessage { - struct BackendKeyData: PayloadDecodable, Equatable { + struct BackendKeyData: PayloadDecodable, Hashable { let processID: Int32 let secretKey: Int32 diff --git a/Sources/PostgresNIO/New/Messages/Bind.swift b/Sources/PostgresNIO/New/Messages/Bind.swift deleted file mode 100644 index 898018d4..00000000 --- a/Sources/PostgresNIO/New/Messages/Bind.swift +++ /dev/null @@ -1,45 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct Bind: PSQLMessagePayloadEncodable, Equatable { - /// The name of the destination portal (an empty string selects the unnamed portal). - var portalName: String - - /// The name of the source prepared statement (an empty string selects the unnamed prepared statement). - var preparedStatementName: String - - /// The number of parameter values that follow (possibly zero). This must match the number of parameters needed by the query. - var bind: PostgresBindings - - func encode(into buffer: inout ByteBuffer) { - buffer.writeNullTerminatedString(self.portalName) - buffer.writeNullTerminatedString(self.preparedStatementName) - - // The number of parameter format codes that follow (denoted C below). This can be - // zero to indicate that there are no parameters or that the parameters all use the - // default format (text); or one, in which case the specified format code is applied - // to all parameters; or it can equal the actual number of parameters. - buffer.writeInteger(UInt16(self.bind.count)) - - // The parameter format codes. Each must presently be zero (text) or one (binary). - self.bind.metadata.forEach { - buffer.writeInteger($0.format.rawValue) - } - - buffer.writeInteger(UInt16(self.bind.count)) - - var parametersCopy = self.bind.bytes - buffer.writeBuffer(¶metersCopy) - - // The number of result-column format codes that follow (denoted R below). This can be - // zero to indicate that there are no result columns or that the result columns should - // all use the default format (text); or one, in which case the specified format code - // is applied to all result columns (if any); or it can equal the actual number of - // result columns of the query. - buffer.writeInteger(1, as: Int16.self) - // The result-column format codes. Each must presently be zero (text) or one (binary). - buffer.writeInteger(PostgresFormat.binary.rawValue, as: Int16.self) - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/Cancel.swift b/Sources/PostgresNIO/New/Messages/Cancel.swift deleted file mode 100644 index 2f29d239..00000000 --- a/Sources/PostgresNIO/New/Messages/Cancel.swift +++ /dev/null @@ -1,21 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct Cancel: PSQLMessagePayloadEncodable, Equatable { - /// The cancel request code. The value is chosen to contain 1234 in the most significant 16 bits, - /// and 5678 in the least significant 16 bits. (To avoid confusion, this code must not be the same - /// as any protocol version number.) - let cancelRequestCode: Int32 = 80877102 - - /// The process ID of the target backend. - let processID: Int32 - - /// The secret key for the target backend. - let secretKey: Int32 - - func encode(into buffer: inout ByteBuffer) { - buffer.writeMultipleIntegers(self.cancelRequestCode, self.processID, self.secretKey) - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/Close.swift b/Sources/PostgresNIO/New/Messages/Close.swift deleted file mode 100644 index 7f038f94..00000000 --- a/Sources/PostgresNIO/New/Messages/Close.swift +++ /dev/null @@ -1,20 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - enum Close: PSQLMessagePayloadEncodable, Equatable { - case preparedStatement(String) - case portal(String) - - func encode(into buffer: inout ByteBuffer) { - switch self { - case .preparedStatement(let name): - buffer.writeInteger(UInt8(ascii: "S")) - buffer.writeNullTerminatedString(name) - case .portal(let name): - buffer.writeInteger(UInt8(ascii: "P")) - buffer.writeNullTerminatedString(name) - } - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/DataRow.swift b/Sources/PostgresNIO/New/Messages/DataRow.swift index b181e600..491e10dc 100644 --- a/Sources/PostgresNIO/New/Messages/DataRow.swift +++ b/Sources/PostgresNIO/New/Messages/DataRow.swift @@ -9,7 +9,7 @@ import NIOCore /// Not putting `DataRow` in ``PSQLBackendMessage`` is our way to trick /// the Swift compiler @usableFromInline -struct DataRow: Sendable, PostgresBackendMessage.PayloadDecodable, Equatable { +struct DataRow: Sendable, PostgresBackendMessage.PayloadDecodable, Hashable { @usableFromInline var columnCount: Int16 @usableFromInline diff --git a/Sources/PostgresNIO/New/Messages/Describe.swift b/Sources/PostgresNIO/New/Messages/Describe.swift deleted file mode 100644 index 76167d32..00000000 --- a/Sources/PostgresNIO/New/Messages/Describe.swift +++ /dev/null @@ -1,21 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - enum Describe: PSQLMessagePayloadEncodable, Equatable { - - case preparedStatement(String) - case portal(String) - - func encode(into buffer: inout ByteBuffer) { - switch self { - case .preparedStatement(let name): - buffer.writeInteger(UInt8(ascii: "S")) - buffer.writeNullTerminatedString(name) - case .portal(let name): - buffer.writeInteger(UInt8(ascii: "P")) - buffer.writeNullTerminatedString(name) - } - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/ErrorResponse.swift b/Sources/PostgresNIO/New/Messages/ErrorResponse.swift index 818c1ebf..d0bb6044 100644 --- a/Sources/PostgresNIO/New/Messages/ErrorResponse.swift +++ b/Sources/PostgresNIO/New/Messages/ErrorResponse.swift @@ -80,7 +80,7 @@ extension PostgresBackendMessage { case routine = 0x52 /// R } - struct ErrorResponse: PSQLMessageNotice, PayloadDecodable, Equatable { + struct ErrorResponse: PSQLMessageNotice, PayloadDecodable, Hashable { let fields: [PostgresBackendMessage.Field: String] init(fields: [PostgresBackendMessage.Field: String]) { @@ -88,7 +88,7 @@ extension PostgresBackendMessage { } } - struct NoticeResponse: PSQLMessageNotice, PayloadDecodable, Equatable { + struct NoticeResponse: PSQLMessageNotice, PayloadDecodable, Hashable { let fields: [PostgresBackendMessage.Field: String] init(fields: [PostgresBackendMessage.Field: String]) { diff --git a/Sources/PostgresNIO/New/Messages/Execute.swift b/Sources/PostgresNIO/New/Messages/Execute.swift deleted file mode 100644 index 17646484..00000000 --- a/Sources/PostgresNIO/New/Messages/Execute.swift +++ /dev/null @@ -1,23 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct Execute: PSQLMessagePayloadEncodable, Equatable { - /// The name of the portal to execute (an empty string selects the unnamed portal). - let portalName: String - - /// Maximum number of rows to return, if portal contains a query that returns rows (ignored otherwise). Zero denotes “no limit”. - let maxNumberOfRows: Int32 - - init(portalName: String, maxNumberOfRows: Int32 = 0) { - self.portalName = portalName - self.maxNumberOfRows = maxNumberOfRows - } - - func encode(into buffer: inout ByteBuffer) { - buffer.writeNullTerminatedString(self.portalName) - buffer.writeInteger(self.maxNumberOfRows) - } - } - -} diff --git a/Sources/PostgresNIO/New/Messages/NotificationResponse.swift b/Sources/PostgresNIO/New/Messages/NotificationResponse.swift index 5cd9422e..01b9ab4a 100644 --- a/Sources/PostgresNIO/New/Messages/NotificationResponse.swift +++ b/Sources/PostgresNIO/New/Messages/NotificationResponse.swift @@ -2,7 +2,7 @@ import NIOCore extension PostgresBackendMessage { - struct NotificationResponse: PayloadDecodable, Equatable { + struct NotificationResponse: PayloadDecodable, Hashable { let backendPID: Int32 let channel: String let payload: String diff --git a/Sources/PostgresNIO/New/Messages/ParameterDescription.swift b/Sources/PostgresNIO/New/Messages/ParameterDescription.swift index 1ccc91e5..4d12b1b6 100644 --- a/Sources/PostgresNIO/New/Messages/ParameterDescription.swift +++ b/Sources/PostgresNIO/New/Messages/ParameterDescription.swift @@ -2,7 +2,7 @@ import NIOCore extension PostgresBackendMessage { - struct ParameterDescription: PayloadDecodable, Equatable { + struct ParameterDescription: PayloadDecodable, Hashable { /// Specifies the object ID of the parameter data type. var dataTypes: [PostgresDataType] diff --git a/Sources/PostgresNIO/New/Messages/ParameterStatus.swift b/Sources/PostgresNIO/New/Messages/ParameterStatus.swift index 4ffcbe12..52d07e01 100644 --- a/Sources/PostgresNIO/New/Messages/ParameterStatus.swift +++ b/Sources/PostgresNIO/New/Messages/ParameterStatus.swift @@ -2,7 +2,7 @@ import NIOCore extension PostgresBackendMessage { - struct ParameterStatus: PayloadDecodable, Equatable { + struct ParameterStatus: PayloadDecodable, Hashable { /// The name of the run-time parameter being reported. var parameter: String diff --git a/Sources/PostgresNIO/New/Messages/Parse.swift b/Sources/PostgresNIO/New/Messages/Parse.swift deleted file mode 100644 index 9d3cfa0b..00000000 --- a/Sources/PostgresNIO/New/Messages/Parse.swift +++ /dev/null @@ -1,26 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct Parse: PSQLMessagePayloadEncodable, Equatable { - /// The name of the destination prepared statement (an empty string selects the unnamed prepared statement). - let preparedStatementName: String - - /// The query string to be parsed. - let query: String - - /// The number of parameter data types specified (can be zero). Note that this is not an indication of the number of parameters that might appear in the query string, only the number that the frontend wants to prespecify types for. - let parameters: [PostgresDataType] - - func encode(into buffer: inout ByteBuffer) { - buffer.writeNullTerminatedString(self.preparedStatementName) - buffer.writeNullTerminatedString(self.query) - buffer.writeInteger(UInt16(self.parameters.count)) - - self.parameters.forEach { dataType in - buffer.writeInteger(dataType.rawValue) - } - } - } - -} diff --git a/Sources/PostgresNIO/New/Messages/Password.swift b/Sources/PostgresNIO/New/Messages/Password.swift deleted file mode 100644 index 81d7ab30..00000000 --- a/Sources/PostgresNIO/New/Messages/Password.swift +++ /dev/null @@ -1,13 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct Password: PSQLMessagePayloadEncodable, Equatable { - let value: String - - func encode(into buffer: inout ByteBuffer) { - buffer.writeNullTerminatedString(value) - } - } - -} diff --git a/Sources/PostgresNIO/New/Messages/ReadyForQuery.swift b/Sources/PostgresNIO/New/Messages/ReadyForQuery.swift index a300f714..41af1b60 100644 --- a/Sources/PostgresNIO/New/Messages/ReadyForQuery.swift +++ b/Sources/PostgresNIO/New/Messages/ReadyForQuery.swift @@ -1,37 +1,11 @@ import NIOCore extension PostgresBackendMessage { - enum TransactionState: PayloadDecodable, RawRepresentable { - typealias RawValue = UInt8 - - case idle - case inTransaction - case inFailedTransaction - - init?(rawValue: UInt8) { - switch rawValue { - case UInt8(ascii: "I"): - self = .idle - case UInt8(ascii: "T"): - self = .inTransaction - case UInt8(ascii: "E"): - self = .inFailedTransaction - default: - return nil - } - } + enum TransactionState: UInt8, PayloadDecodable, Hashable { + case idle = 73 // ascii: I + case inTransaction = 84 // ascii: T + case inFailedTransaction = 69 // ascii: E - var rawValue: Self.RawValue { - switch self { - case .idle: - return UInt8(ascii: "I") - case .inTransaction: - return UInt8(ascii: "T") - case .inFailedTransaction: - return UInt8(ascii: "E") - } - } - static func decode(from buffer: inout ByteBuffer) throws -> Self { let value = try buffer.throwingReadInteger(as: UInt8.self) guard let state = Self.init(rawValue: value) else { diff --git a/Sources/PostgresNIO/New/Messages/RowDescription.swift b/Sources/PostgresNIO/New/Messages/RowDescription.swift index 66c71215..766d06e9 100644 --- a/Sources/PostgresNIO/New/Messages/RowDescription.swift +++ b/Sources/PostgresNIO/New/Messages/RowDescription.swift @@ -9,13 +9,13 @@ import NIOCore /// Not putting `DataRow` in ``PSQLBackendMessage`` is our way to trick /// the Swift compiler. @usableFromInline -struct RowDescription: PostgresBackendMessage.PayloadDecodable, Sendable, Equatable { +struct RowDescription: PostgresBackendMessage.PayloadDecodable, Sendable, Hashable { /// Specifies the object ID of the parameter data type. @usableFromInline var columns: [Column] @usableFromInline - struct Column: Equatable, Sendable { + struct Column: Hashable, Sendable { /// The field name. @usableFromInline var name: String diff --git a/Sources/PostgresNIO/New/Messages/SASLInitialResponse.swift b/Sources/PostgresNIO/New/Messages/SASLInitialResponse.swift deleted file mode 100644 index 73db9332..00000000 --- a/Sources/PostgresNIO/New/Messages/SASLInitialResponse.swift +++ /dev/null @@ -1,28 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct SASLInitialResponse: PSQLMessagePayloadEncodable, Equatable { - - let saslMechanism: String - let initialData: [UInt8] - - /// Creates a new `SSLRequest`. - init(saslMechanism: String, initialData: [UInt8]) { - self.saslMechanism = saslMechanism - self.initialData = initialData - } - - /// Serializes this message into a byte buffer. - func encode(into buffer: inout ByteBuffer) { - buffer.writeNullTerminatedString(self.saslMechanism) - - if self.initialData.count > 0 { - buffer.writeInteger(Int32(self.initialData.count)) - buffer.writeBytes(self.initialData) - } else { - buffer.writeInteger(Int32(-1)) - } - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/SASLResponse.swift b/Sources/PostgresNIO/New/Messages/SASLResponse.swift deleted file mode 100644 index a6709dcd..00000000 --- a/Sources/PostgresNIO/New/Messages/SASLResponse.swift +++ /dev/null @@ -1,19 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - - struct SASLResponse: PSQLMessagePayloadEncodable, Equatable { - - let data: [UInt8] - - /// Creates a new `SSLRequest`. - init(data: [UInt8]) { - self.data = data - } - - /// Serializes this message into a byte buffer. - func encode(into buffer: inout ByteBuffer) { - buffer.writeBytes(self.data) - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/SSLRequest.swift b/Sources/PostgresNIO/New/Messages/SSLRequest.swift deleted file mode 100644 index 6f9c45a3..00000000 --- a/Sources/PostgresNIO/New/Messages/SSLRequest.swift +++ /dev/null @@ -1,21 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - /// A message asking the PostgreSQL server if TLS is supported - /// For more info, see https://www.postgresql.org/docs/10/static/protocol-flow.html#id-1.10.5.7.11 - struct SSLRequest: PSQLMessagePayloadEncodable, Equatable { - /// The SSL request code. The value is chosen to contain 1234 in the most significant 16 bits, - /// and 5679 in the least significant 16 bits. - let code: Int32 - - /// Creates a new `SSLRequest`. - init() { - self.code = 80877103 - } - - /// Serializes this message into a byte buffer. - func encode(into buffer: inout ByteBuffer) { - buffer.writeInteger(self.code) - } - } -} diff --git a/Sources/PostgresNIO/New/Messages/Startup.swift b/Sources/PostgresNIO/New/Messages/Startup.swift deleted file mode 100644 index f7da2127..00000000 --- a/Sources/PostgresNIO/New/Messages/Startup.swift +++ /dev/null @@ -1,82 +0,0 @@ -import NIOCore - -extension PostgresFrontendMessage { - struct Startup: PSQLMessagePayloadEncodable, Equatable { - - /// Creates a `Startup` with "3.0" as the protocol version. - static func versionThree(parameters: Parameters) -> Startup { - return .init(protocolVersion: 0x00_03_00_00, parameters: parameters) - } - - /// The protocol version number. The most significant 16 bits are the major - /// version number (3 for the protocol described here). The least significant - /// 16 bits are the minor version number (0 for the protocol described here). - var protocolVersion: Int32 - - /// The protocol version number is followed by one or more pairs of parameter - /// name and value strings. A zero byte is required as a terminator after - /// the last name/value pair. `user` is required, others are optional. - struct Parameters: Equatable { - enum Replication { - case `true` - case `false` - case database - } - - /// The database user name to connect as. Required; there is no default. - var user: String - - /// The database to connect to. Defaults to the user name. - var database: String? - - /// Command-line arguments for the backend. (This is deprecated in favor - /// of setting individual run-time parameters.) Spaces within this string are - /// considered to separate arguments, unless escaped with a - /// backslash (\); write \\ to represent a literal backslash. - var options: String? - - /// Used to connect in streaming replication mode, where a small set of - /// replication commands can be issued instead of SQL statements. Value - /// can be true, false, or database, and the default is false. - var replication: Replication - } - var parameters: Parameters - - /// Creates a new `PostgreSQLStartupMessage`. - init(protocolVersion: Int32, parameters: Parameters) { - self.protocolVersion = protocolVersion - self.parameters = parameters - } - - /// Serializes this message into a byte buffer. - func encode(into buffer: inout ByteBuffer) { - buffer.writeInteger(self.protocolVersion) - buffer.writeNullTerminatedString("user") - buffer.writeNullTerminatedString(self.parameters.user) - - if let database = self.parameters.database { - buffer.writeNullTerminatedString("database") - buffer.writeNullTerminatedString(database) - } - - if let options = self.parameters.options { - buffer.writeNullTerminatedString("options") - buffer.writeNullTerminatedString(options) - } - - switch self.parameters.replication { - case .database: - buffer.writeNullTerminatedString("replication") - buffer.writeNullTerminatedString("replication") - case .true: - buffer.writeNullTerminatedString("replication") - buffer.writeNullTerminatedString("true") - case .false: - break - } - - buffer.writeInteger(UInt8(0)) - } - } - -} diff --git a/Sources/PostgresNIO/New/NotificationListener.swift b/Sources/PostgresNIO/New/NotificationListener.swift new file mode 100644 index 00000000..2f784e33 --- /dev/null +++ b/Sources/PostgresNIO/New/NotificationListener.swift @@ -0,0 +1,142 @@ +import NIOCore + +// This object is @unchecked Sendable, since we syncronize state on the EL +final class NotificationListener: @unchecked Sendable { + let eventLoop: EventLoop + + let channel: String + let id: Int + + private var state: State + + enum State { + case streamInitialized(CheckedContinuation) + case streamListening(AsyncThrowingStream.Continuation) + + case closure(PostgresListenContext, (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void) + case done + } + + init( + channel: String, + id: Int, + eventLoop: EventLoop, + checkedContinuation: CheckedContinuation + ) { + self.channel = channel + self.id = id + self.eventLoop = eventLoop + self.state = .streamInitialized(checkedContinuation) + } + + init( + channel: String, + id: Int, + eventLoop: EventLoop, + context: PostgresListenContext, + closure: @Sendable @escaping (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void + ) { + self.channel = channel + self.id = id + self.eventLoop = eventLoop + self.state = .closure(context, closure) + } + + func startListeningSucceeded(handler: PostgresChannelHandler) { + self.eventLoop.preconditionInEventLoop() + let handlerLoopBound = NIOLoopBound(handler, eventLoop: self.eventLoop) + + switch self.state { + case .streamInitialized(let checkedContinuation): + let (stream, continuation) = AsyncThrowingStream.makeStream(of: PostgresNotification.self) + let eventLoop = self.eventLoop + let channel = self.channel + let listenerID = self.id + continuation.onTermination = { reason in + switch reason { + case .cancelled: + eventLoop.execute { + handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID) + } + + case .finished: + break + + @unknown default: + break + } + } + self.state = .streamListening(continuation) + + let notificationSequence = PostgresNotificationSequence(base: stream) + checkedContinuation.resume(returning: notificationSequence) + + case .streamListening, .done: + fatalError("Invalid state: \(self.state)") + + case .closure: + break // ignore + } + } + + func notificationReceived(_ backendMessage: PostgresBackendMessage.NotificationResponse) { + self.eventLoop.preconditionInEventLoop() + + switch self.state { + case .streamInitialized, .done: + fatalError("Invalid state: \(self.state)") + case .streamListening(let continuation): + continuation.yield(.init(payload: backendMessage.payload)) + + case .closure(let postgresListenContext, let closure): + let message = PostgresMessage.NotificationResponse( + backendPID: backendMessage.backendPID, + channel: backendMessage.channel, + payload: backendMessage.payload + ) + closure(postgresListenContext, message) + } + } + + func failed(_ error: Error) { + self.eventLoop.preconditionInEventLoop() + + switch self.state { + case .streamInitialized(let checkedContinuation): + self.state = .done + checkedContinuation.resume(throwing: error) + + case .streamListening(let continuation): + self.state = .done + continuation.finish(throwing: error) + + case .closure(let postgresListenContext, _): + self.state = .done + postgresListenContext.cancel() + + case .done: + break // ignore + } + } + + func cancelled() { + self.eventLoop.preconditionInEventLoop() + + switch self.state { + case .streamInitialized(let checkedContinuation): + self.state = .done + checkedContinuation.resume(throwing: PSQLError(code: .queryCancelled)) + + case .streamListening(let continuation): + self.state = .done + continuation.finish() + + case .closure(let postgresListenContext, _): + self.state = .done + postgresListenContext.cancel() + + case .done: + break // ignore + } + } +} diff --git a/Sources/PostgresNIO/New/PSQLError.swift b/Sources/PostgresNIO/New/PSQLError.swift index 08b6a01e..4a9f9216 100644 --- a/Sources/PostgresNIO/New/PSQLError.swift +++ b/Sources/PostgresNIO/New/PSQLError.swift @@ -1,7 +1,8 @@ import NIOCore /// An error that is thrown from the PostgresClient. -public struct PSQLError: Error { +/// Sendability enforced through Copy on Write semantics +public struct PSQLError: Error, @unchecked Sendable { public struct Code: Sendable, Hashable, CustomStringConvertible { enum Base: Sendable, Hashable { @@ -18,10 +19,14 @@ public struct PSQLError: Error { case queryCancelled case tooManyParameters - case connectionQuiescing - case connectionClosed + case clientClosedConnection + case serverClosedConnection case connectionError case uncleanShutdown + + case listenFailed + case unlistenFailed + case poolClosed } internal var base: Base @@ -30,7 +35,7 @@ public struct PSQLError: Error { self.base = base } - public static let sslUnsupported = Self.init(.sslUnsupported) + public static let sslUnsupported = Self(.sslUnsupported) public static let failedToAddSSLHandler = Self(.failedToAddSSLHandler) public static let receivedUnencryptedDataAfterSSLRequest = Self(.receivedUnencryptedDataAfterSSLRequest) public static let server = Self(.server) @@ -38,14 +43,25 @@ public struct PSQLError: Error { public static let unexpectedBackendMessage = Self(.unexpectedBackendMessage) public static let unsupportedAuthMechanism = Self(.unsupportedAuthMechanism) public static let authMechanismRequiresPassword = Self(.authMechanismRequiresPassword) - public static let saslError = Self.init(.saslError) + public static let saslError = Self(.saslError) public static let invalidCommandTag = Self(.invalidCommandTag) public static let queryCancelled = Self(.queryCancelled) public static let tooManyParameters = Self(.tooManyParameters) - public static let connectionQuiescing = Self(.connectionQuiescing) - public static let connectionClosed = Self(.connectionClosed) + public static let clientClosedConnection = Self(.clientClosedConnection) + public static let serverClosedConnection = Self(.serverClosedConnection) public static let connectionError = Self(.connectionError) - public static let uncleanShutdown = Self.init(.uncleanShutdown) + + public static let uncleanShutdown = Self(.uncleanShutdown) + public static let poolClosed = Self(.poolClosed) + + public static let listenFailed = Self.init(.listenFailed) + public static let unlistenFailed = Self.init(.unlistenFailed) + + @available(*, deprecated, renamed: "clientClosedConnection") + public static let connectionQuiescing = Self.clientClosedConnection + + @available(*, deprecated, message: "Use the more specific `serverClosedConnection` or `clientClosedConnection` instead") + public static let connectionClosed = Self.serverClosedConnection public var description: String { switch self.base { @@ -73,21 +89,27 @@ public struct PSQLError: Error { return "queryCancelled" case .tooManyParameters: return "tooManyParameters" - case .connectionQuiescing: - return "connectionQuiescing" - case .connectionClosed: - return "connectionClosed" + case .clientClosedConnection: + return "clientClosedConnection" + case .serverClosedConnection: + return "serverClosedConnection" case .connectionError: return "connectionError" case .uncleanShutdown: return "uncleanShutdown" + case .poolClosed: + return "poolClosed" + case .listenFailed: + return "listenFailed" + case .unlistenFailed: + return "unlistenFailed" } } } private var backing: Backing - private mutating func copyBackingStoriageIfNecessary() { + private mutating func copyBackingStorageIfNecessary() { if !isKnownUniquelyReferenced(&self.backing) { self.backing = self.backing.copy() } @@ -97,7 +119,7 @@ public struct PSQLError: Error { public internal(set) var code: Code { get { self.backing.code } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.code = newValue } } @@ -106,7 +128,7 @@ public struct PSQLError: Error { public internal(set) var serverInfo: ServerInfo? { get { self.backing.serverInfo } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.serverInfo = newValue } } @@ -115,7 +137,7 @@ public struct PSQLError: Error { public internal(set) var underlying: Error? { get { self.backing.underlying } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.underlying = newValue } } @@ -124,7 +146,7 @@ public struct PSQLError: Error { public internal(set) var file: String? { get { self.backing.file } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.file = newValue } } @@ -133,7 +155,7 @@ public struct PSQLError: Error { public internal(set) var line: Int? { get { self.backing.line } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.line = newValue } } @@ -142,7 +164,7 @@ public struct PSQLError: Error { public internal(set) var query: PostgresQuery? { get { self.backing.query } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.query = newValue } } @@ -152,7 +174,7 @@ public struct PSQLError: Error { var backendMessage: PostgresBackendMessage? { get { self.backing.backendMessage } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.backendMessage = newValue } } @@ -162,7 +184,7 @@ public struct PSQLError: Error { var unsupportedAuthScheme: UnsupportedAuthScheme? { get { self.backing.unsupportedAuthScheme } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.unsupportedAuthScheme = newValue } } @@ -172,7 +194,7 @@ public struct PSQLError: Error { var invalidCommandTag: String? { get { self.backing.invalidCommandTag } set { - self.copyBackingStoriageIfNecessary() + self.copyBackingStorageIfNecessary() self.backing.invalidCommandTag = newValue } } @@ -190,21 +212,13 @@ public struct PSQLError: Error { private final class Backing { fileprivate var code: Code - fileprivate var serverInfo: ServerInfo? - fileprivate var underlying: Error? - fileprivate var file: String? - fileprivate var line: Int? - fileprivate var query: PostgresQuery? - fileprivate var backendMessage: PostgresBackendMessage? - fileprivate var unsupportedAuthScheme: UnsupportedAuthScheme? - fileprivate var invalidCommandTag: String? init(code: Code) { @@ -224,10 +238,10 @@ public struct PSQLError: Error { } public struct ServerInfo { - public struct Field: Hashable, Sendable { + public struct Field: Hashable, Sendable, CustomStringConvertible { fileprivate let backing: PostgresBackendMessage.Field - private init(_ backing: PostgresBackendMessage.Field) { + fileprivate init(_ backing: PostgresBackendMessage.Field) { self.backing = backing } @@ -306,6 +320,47 @@ public struct PSQLError: Error { /// Routine: the name of the source-code routine reporting the error. public static let routine = Self(.routine) + + public var description: String { + switch self.backing { + case .localizedSeverity: + return "localizedSeverity" + case .severity: + return "severity" + case .sqlState: + return "sqlState" + case .message: + return "message" + case .detail: + return "detail" + case .hint: + return "hint" + case .position: + return "position" + case .internalPosition: + return "internalPosition" + case .internalQuery: + return "internalQuery" + case .locationContext: + return "locationContext" + case .schemaName: + return "schemaName" + case .tableName: + return "tableName" + case .columnName: + return "columnName" + case .dataTypeName: + return "dataTypeName" + case .constraintName: + return "constraintName" + case .file: + return "file" + case .line: + return "line" + case .routine: + return "routine" + } + } } let underlying: PostgresBackendMessage.ErrorResponse @@ -335,19 +390,27 @@ public struct PSQLError: Error { return new } - static var connectionQuiescing: PSQLError { PSQLError(code: .connectionQuiescing) } + static func clientClosedConnection(underlying: Error?) -> PSQLError { + var error = PSQLError(code: .clientClosedConnection) + error.underlying = underlying + return error + } - static var connectionClosed: PSQLError { PSQLError(code: .connectionClosed) } + static func serverClosedConnection(underlying: Error?) -> PSQLError { + var error = PSQLError(code: .serverClosedConnection) + error.underlying = underlying + return error + } - static var authMechanismRequiresPassword: PSQLError { PSQLError(code: .authMechanismRequiresPassword) } + static let authMechanismRequiresPassword = PSQLError(code: .authMechanismRequiresPassword) - static var sslUnsupported: PSQLError { PSQLError(code: .sslUnsupported) } + static let sslUnsupported = PSQLError(code: .sslUnsupported) - static var queryCancelled: PSQLError { PSQLError(code: .queryCancelled) } + static let queryCancelled = PSQLError(code: .queryCancelled) - static var uncleanShutdown: PSQLError { PSQLError(code: .uncleanShutdown) } + static let uncleanShutdown = PSQLError(code: .uncleanShutdown) - static var receivedUnencryptedDataAfterSSLRequest: PSQLError { PSQLError(code: .receivedUnencryptedDataAfterSSLRequest) } + static let receivedUnencryptedDataAfterSSLRequest = PSQLError(code: .receivedUnencryptedDataAfterSSLRequest) static func server(_ response: PostgresBackendMessage.ErrorResponse) -> PSQLError { var error = PSQLError(code: .server) @@ -385,6 +448,12 @@ public struct PSQLError: Error { return error } + static func unlistenError(underlying: Error) -> PSQLError { + var error = PSQLError(code: .unlistenFailed) + error.underlying = underlying + return error + } + enum UnsupportedAuthScheme { case none case kerberosV5 @@ -395,6 +464,69 @@ public struct PSQLError: Error { case sspi case sasl(mechanisms: [String]) } + + static var poolClosed: PSQLError { + Self.init(code: .poolClosed) + } +} + +extension PSQLError: CustomStringConvertible { + public var description: String { + // This may seem very odd... But we are afraid that users might accidentally send the + // unfiltered errors out to end-users. This may leak security relevant information. For this + // reason we overwrite the error description by default to this generic "Database error" + """ + PSQLError – Generic description to prevent accidental leakage of sensitive data. For debugging details, use `String(reflecting: error)`. + """ + } +} + +extension PSQLError: CustomDebugStringConvertible { + public var debugDescription: String { + var result = #"PSQLError(code: \#(self.code)"# + + if let serverInfo = self.serverInfo?.underlying { + result.append(", serverInfo: [") + result.append( + serverInfo.fields + .sorted(by: { $0.key.rawValue < $1.key.rawValue }) + .map { "\(PSQLError.ServerInfo.Field($0.0)): \($0.1)" } + .joined(separator: ", ") + ) + result.append("]") + } + + if let backendMessage = self.backendMessage { + result.append(", backendMessage: \(String(reflecting: backendMessage))") + } + + if let unsupportedAuthScheme = self.unsupportedAuthScheme { + result.append(", unsupportedAuthScheme: \(unsupportedAuthScheme)") + } + + if let invalidCommandTag = self.invalidCommandTag { + result.append(", invalidCommandTag: \(invalidCommandTag)") + } + + if let underlying = self.underlying { + result.append(", underlying: \(String(reflecting: underlying))") + } + + if let file = self.file { + result.append(", triggeredFromRequestInFile: \(file)") + if let line = self.line { + result.append(", line: \(line)") + } + } + + if let query = self.query { + result.append(", query: \(String(reflecting: query))") + } + + result.append(")") + + return result + } } /// An error that may happen when a ``PostgresRow`` or ``PostgresCell`` is decoded to native Swift types. @@ -490,7 +622,9 @@ extension PostgresDecodingError: CustomStringConvertible { // This may seem very odd... But we are afraid that users might accidentally send the // unfiltered errors out to end-users. This may leak security relevant information. For this // reason we overwrite the error description by default to this generic "Database error" - "Database error" + """ + PostgresDecodingError – Generic description to prevent accidental leakage of sensitive data. For debugging details, use `String(reflecting: error)`. + """ } } @@ -504,7 +638,7 @@ extension PostgresDecodingError: CustomDebugStringConvertible { result.append(#", postgresType: \#(self.postgresType)"#) result.append(#", postgresFormat: \#(self.postgresFormat)"#) if let postgresData = self.postgresData { - result.append(#", postgresData: \#(postgresData.debugDescription)"#) // https://github.com/apple/swift-nio/pull/2418 + result.append(#", postgresData: \#(String(reflecting: postgresData))"#) } result.append(#", file: \#(self.file)"#) result.append(#", line: \#(self.line)"#) diff --git a/Sources/PostgresNIO/New/PSQLEventsHandler.swift b/Sources/PostgresNIO/New/PSQLEventsHandler.swift index 3233fb77..0f426f20 100644 --- a/Sources/PostgresNIO/New/PSQLEventsHandler.swift +++ b/Sources/PostgresNIO/New/PSQLEventsHandler.swift @@ -7,6 +7,8 @@ enum PSQLOutgoingEvent { /// /// this shall be removed with the next breaking change and always supplied with `PSQLConnection.Configuration` case authenticate(AuthContext) + + case gracefulShutdown } enum PSQLEvent { @@ -66,10 +68,8 @@ final class PSQLEventsHandler: ChannelInboundHandler { case .authenticated: break } - case TLSUserEvent.shutdownCompleted: - break default: - preconditionFailure() + context.fireUserInboundEventTriggered(event) } } diff --git a/Sources/PostgresNIO/New/PSQLFrontendMessageEncoder.swift b/Sources/PostgresNIO/New/PSQLFrontendMessageEncoder.swift deleted file mode 100644 index 24155d84..00000000 --- a/Sources/PostgresNIO/New/PSQLFrontendMessageEncoder.swift +++ /dev/null @@ -1,85 +0,0 @@ -import NIOCore - -struct PSQLFrontendMessageEncoder: MessageToByteEncoder { - typealias OutboundIn = PostgresFrontendMessage - - init() {} - - func encode(data message: PostgresFrontendMessage, out buffer: inout ByteBuffer) { - switch message { - case .bind(let bind): - buffer.writeInteger(message.id.rawValue) - let startIndex = buffer.writerIndex - buffer.writeInteger(Int32(0)) // placeholder for length - bind.encode(into: &buffer) - let length = Int32(buffer.writerIndex - startIndex) - buffer.setInteger(length, at: startIndex) - - case .cancel(let cancel): - // cancel requests don't have an identifier - self.encode(payload: cancel, into: &buffer) - - case .close(let close): - self.encode(messageID: message.id, payload: close, into: &buffer) - - case .describe(let describe): - self.encode(messageID: message.id, payload: describe, into: &buffer) - - case .execute(let execute): - self.encode(messageID: message.id, payload: execute, into: &buffer) - - case .flush: - self.encode(messageID: message.id, payload: EmptyPayload(), into: &buffer) - - case .parse(let parse): - self.encode(messageID: message.id, payload: parse, into: &buffer) - - case .password(let password): - self.encode(messageID: message.id, payload: password, into: &buffer) - - case .saslInitialResponse(let saslInitialResponse): - self.encode(messageID: message.id, payload: saslInitialResponse, into: &buffer) - - case .saslResponse(let saslResponse): - self.encode(messageID: message.id, payload: saslResponse, into: &buffer) - - case .sslRequest(let request): - // sslRequests don't have an identifier - self.encode(payload: request, into: &buffer) - - case .startup(let startup): - // startup requests don't have an identifier - self.encode(payload: startup, into: &buffer) - - case .sync: - self.encode(messageID: message.id, payload: EmptyPayload(), into: &buffer) - - case .terminate: - self.encode(messageID: message.id, payload: EmptyPayload(), into: &buffer) - } - } - - private struct EmptyPayload: PSQLMessagePayloadEncodable { - func encode(into buffer: inout ByteBuffer) {} - } - - private func encode( - messageID: PostgresFrontendMessage.ID, - payload: Payload, - into buffer: inout ByteBuffer) - { - buffer.psqlWriteFrontendMessageID(messageID) - self.encode(payload: payload, into: &buffer) - } - - private func encode( - payload: Payload, - into buffer: inout ByteBuffer) - { - let startIndex = buffer.writerIndex - buffer.writeInteger(Int32(0)) // placeholder for length - payload.encode(into: &buffer) - let length = Int32(buffer.writerIndex - startIndex) - buffer.setInteger(length, at: startIndex) - } -} diff --git a/Sources/PostgresNIO/New/PSQLRowStream.swift b/Sources/PostgresNIO/New/PSQLRowStream.swift index 4c842275..ee925d0e 100644 --- a/Sources/PostgresNIO/New/PSQLRowStream.swift +++ b/Sources/PostgresNIO/New/PSQLRowStream.swift @@ -1,57 +1,74 @@ import NIOCore import Logging +struct QueryResult { + enum Value: Equatable { + case noRows(PSQLRowStream.StatementSummary) + case rowDescription([RowDescription.Column]) + } + + var value: Value + + var logger: Logger +} + // Thread safety is guaranteed in the RowStream through dispatching onto the NIO EventLoop. final class PSQLRowStream: @unchecked Sendable { private typealias AsyncSequenceSource = NIOThrowingAsyncSequenceProducer.Source - enum RowSource { - case stream(PSQLRowsDataSource) - case noRows(Result) + enum StatementSummary: Equatable { + case tag(String) + case emptyResponse + } + + enum Source { + case stream([RowDescription.Column], PSQLRowsDataSource) + case noRows(Result) } let eventLoop: EventLoop let logger: Logger - + private enum BufferState { case streaming(buffer: CircularBuffer, dataSource: PSQLRowsDataSource) - case finished(buffer: CircularBuffer, commandTag: String) + case finished(buffer: CircularBuffer, summary: StatementSummary) case failure(Error) } - + private enum DownstreamState { case waitingForConsumer(BufferState) case iteratingRows(onRow: (PostgresRow) throws -> (), EventLoopPromise, PSQLRowsDataSource) case waitingForAll([PostgresRow], EventLoopPromise<[PostgresRow]>, PSQLRowsDataSource) - case consumed(Result) - case asyncSequence(AsyncSequenceSource, PSQLRowsDataSource) + case consumed(Result) + case asyncSequence(AsyncSequenceSource, PSQLRowsDataSource, onFinish: @Sendable () -> ()) } internal let rowDescription: [RowDescription.Column] private let lookupTable: [String: Int] private var downstreamState: DownstreamState - init(rowDescription: [RowDescription.Column], - queryContext: ExtendedQueryContext, - eventLoop: EventLoop, - rowSource: RowSource) - { + init( + source: Source, + eventLoop: EventLoop, + logger: Logger + ) { let bufferState: BufferState - switch rowSource { - case .stream(let dataSource): + switch source { + case .stream(let rowDescription, let dataSource): + self.rowDescription = rowDescription bufferState = .streaming(buffer: .init(), dataSource: dataSource) - case .noRows(.success(let commandTag)): - bufferState = .finished(buffer: .init(), commandTag: commandTag) + case .noRows(.success(let summary)): + self.rowDescription = [] + bufferState = .finished(buffer: .init(), summary: summary) case .noRows(.failure(let error)): + self.rowDescription = [] bufferState = .failure(error) } self.downstreamState = .waitingForConsumer(bufferState) self.eventLoop = eventLoop - self.logger = queryContext.logger - - self.rowDescription = rowDescription + self.logger = logger var lookup = [String: Int]() lookup.reserveCapacity(rowDescription.count) @@ -63,7 +80,7 @@ final class PSQLRowStream: @unchecked Sendable { // MARK: Async Sequence - func asyncSequence() -> PostgresRowSequence { + func asyncSequence(onFinish: @escaping @Sendable () -> () = {}) -> PostgresRowSequence { self.eventLoop.preconditionInEventLoop() guard case .waitingForConsumer(let bufferState) = self.downstreamState else { @@ -74,6 +91,7 @@ final class PSQLRowStream: @unchecked Sendable { elementType: DataRow.self, failureType: Error.self, backPressureStrategy: AdaptiveRowBuffer(), + finishOnDeinit: false, delegate: self ) @@ -82,17 +100,15 @@ final class PSQLRowStream: @unchecked Sendable { switch bufferState { case .streaming(let bufferedRows, let dataSource): let yieldResult = source.yield(contentsOf: bufferedRows) - self.downstreamState = .asyncSequence(source, dataSource) + self.downstreamState = .asyncSequence(source, dataSource, onFinish: onFinish) + self.executeActionBasedOnYieldResult(yieldResult, source: dataSource) - self.eventLoop.execute { - self.executeActionBasedOnYieldResult(yieldResult, source: dataSource) - } - - case .finished(let buffer, let commandTag): + case .finished(let buffer, let summary): _ = source.yield(contentsOf: buffer) source.finish() - self.downstreamState = .consumed(.success(commandTag)) - + onFinish() + self.downstreamState = .consumed(.success(summary)) + case .failure(let error): source.finish(error) self.downstreamState = .consumed(.failure(error)) @@ -119,7 +135,7 @@ final class PSQLRowStream: @unchecked Sendable { case .consumed: break - case .asyncSequence(_, let dataSource): + case .asyncSequence(_, let dataSource, _): dataSource.request(for: self) } } @@ -136,9 +152,10 @@ final class PSQLRowStream: @unchecked Sendable { private func cancel0() { switch self.downstreamState { - case .asyncSequence(_, let dataSource): + case .asyncSequence(_, let dataSource, let onFinish): self.downstreamState = .consumed(.failure(CancellationError())) dataSource.cancel(for: self) + onFinish() case .consumed: return @@ -178,12 +195,12 @@ final class PSQLRowStream: @unchecked Sendable { dataSource.request(for: self) return promise.futureResult - case .finished(let buffer, let commandTag): + case .finished(let buffer, let summary): let rows = buffer.map { PostgresRow(data: $0, lookupTable: self.lookupTable, columns: self.rowDescription) } - self.downstreamState = .consumed(.success(commandTag)) + self.downstreamState = .consumed(.success(summary)) return self.eventLoop.makeSucceededFuture(rows) case .failure(let error): @@ -194,7 +211,7 @@ final class PSQLRowStream: @unchecked Sendable { // MARK: Consume on EventLoop - func onRow(_ onRow: @escaping (PostgresRow) throws -> ()) -> EventLoopFuture { + func onRow(_ onRow: @Sendable @escaping (PostgresRow) throws -> ()) -> EventLoopFuture { if self.eventLoop.inEventLoop { return self.onRow0(onRow) } else { @@ -235,8 +252,8 @@ final class PSQLRowStream: @unchecked Sendable { } return promise.futureResult - - case .finished(let buffer, let commandTag): + + case .finished(let buffer, let summary): do { for data in buffer { let row = PostgresRow( @@ -247,7 +264,7 @@ final class PSQLRowStream: @unchecked Sendable { try onRow(row) } - self.downstreamState = .consumed(.success(commandTag)) + self.downstreamState = .consumed(.success(summary)) return self.eventLoop.makeSucceededVoidFuture() } catch { self.downstreamState = .consumed(.failure(error)) @@ -280,7 +297,7 @@ final class PSQLRowStream: @unchecked Sendable { case .waitingForConsumer(.finished), .waitingForConsumer(.failure): preconditionFailure("How can new rows be received, if an end was already signalled?") - + case .iteratingRows(let onRow, let promise, let dataSource): do { for data in newRows { @@ -309,7 +326,7 @@ final class PSQLRowStream: @unchecked Sendable { // immediately request more dataSource.request(for: self) - case .asyncSequence(let consumer, let source): + case .asyncSequence(let consumer, let source, _): let yieldResult = consumer.yield(contentsOf: newRows) self.executeActionBasedOnYieldResult(yieldResult, source: source) @@ -335,24 +352,25 @@ final class PSQLRowStream: @unchecked Sendable { private func receiveEnd(_ commandTag: String) { switch self.downstreamState { case .waitingForConsumer(.streaming(buffer: let buffer, _)): - self.downstreamState = .waitingForConsumer(.finished(buffer: buffer, commandTag: commandTag)) - - case .waitingForConsumer(.finished), .waitingForConsumer(.failure): + self.downstreamState = .waitingForConsumer(.finished(buffer: buffer, summary: .tag(commandTag))) + + case .waitingForConsumer(.finished), .waitingForConsumer(.failure), .consumed(.success(.emptyResponse)): preconditionFailure("How can we get another end, if an end was already signalled?") case .iteratingRows(_, let promise, _): - self.downstreamState = .consumed(.success(commandTag)) + self.downstreamState = .consumed(.success(.tag(commandTag))) promise.succeed(()) case .waitingForAll(let rows, let promise, _): - self.downstreamState = .consumed(.success(commandTag)) + self.downstreamState = .consumed(.success(.tag(commandTag))) promise.succeed(rows) - case .asyncSequence(let source, _): + case .asyncSequence(let source, _, let onFinish): + self.downstreamState = .consumed(.success(.tag(commandTag))) source.finish() - self.downstreamState = .consumed(.success(commandTag)) - - case .consumed: + onFinish() + + case .consumed(.success(.tag)), .consumed(.failure): break } } @@ -362,7 +380,7 @@ final class PSQLRowStream: @unchecked Sendable { case .waitingForConsumer(.streaming): self.downstreamState = .waitingForConsumer(.failure(error)) - case .waitingForConsumer(.finished), .waitingForConsumer(.failure): + case .waitingForConsumer(.finished), .waitingForConsumer(.failure), .consumed(.success(.emptyResponse)): preconditionFailure("How can we get another end, if an end was already signalled?") case .iteratingRows(_, let promise, _): @@ -373,11 +391,12 @@ final class PSQLRowStream: @unchecked Sendable { self.downstreamState = .consumed(.failure(error)) promise.fail(error) - case .asyncSequence(let consumer, _): - consumer.finish(error) + case .asyncSequence(let consumer, _, let onFinish): self.downstreamState = .consumed(.failure(error)) + consumer.finish(error) + onFinish() - case .consumed: + case .consumed(.success(.tag)), .consumed(.failure): break } } @@ -399,10 +418,15 @@ final class PSQLRowStream: @unchecked Sendable { } var commandTag: String { - guard case .consumed(.success(let commandTag)) = self.downstreamState else { + guard case .consumed(.success(let consumed)) = self.downstreamState else { preconditionFailure("commandTag may only be called if all rows have been consumed") } - return commandTag + switch consumed { + case .tag(let tag): + return tag + case .emptyResponse: + return "" + } } } diff --git a/Sources/PostgresNIO/New/PSQLTask.swift b/Sources/PostgresNIO/New/PSQLTask.swift index f9ca1232..6106fd21 100644 --- a/Sources/PostgresNIO/New/PSQLTask.swift +++ b/Sources/PostgresNIO/New/PSQLTask.swift @@ -1,80 +1,115 @@ import Logging import NIOCore +enum HandlerTask: Sendable { + case extendedQuery(ExtendedQueryContext) + case closeCommand(CloseCommandContext) + case startListening(NotificationListener) + case cancelListening(String, Int) + case executePreparedStatement(PreparedStatementContext) +} + enum PSQLTask { case extendedQuery(ExtendedQueryContext) - case preparedStatement(PrepareStatementContext) case closeCommand(CloseCommandContext) - + func failWithError(_ error: PSQLError) { switch self { case .extendedQuery(let extendedQueryContext): - extendedQueryContext.promise.fail(error) - case .preparedStatement(let createPreparedStatementContext): - createPreparedStatementContext.promise.fail(error) + switch extendedQueryContext.query { + case .unnamed(_, let eventLoopPromise): + eventLoopPromise.fail(error) + case .executeStatement(_, let eventLoopPromise): + eventLoopPromise.fail(error) + case .prepareStatement(_, _, _, let eventLoopPromise): + eventLoopPromise.fail(error) + } + case .closeCommand(let closeCommandContext): closeCommandContext.promise.fail(error) } } } -final class ExtendedQueryContext { +final class ExtendedQueryContext: Sendable { enum Query { - case unnamed(PostgresQuery) - case preparedStatement(PSQLExecuteStatement) + case unnamed(PostgresQuery, EventLoopPromise) + case executeStatement(PSQLExecuteStatement, EventLoopPromise) + case prepareStatement(name: String, query: String, bindingDataTypes: [PostgresDataType], EventLoopPromise) } let query: Query let logger: Logger - - let promise: EventLoopPromise - init(query: PostgresQuery, - logger: Logger, - promise: EventLoopPromise) - { - self.query = .unnamed(query) + init( + query: PostgresQuery, + logger: Logger, + promise: EventLoopPromise + ) { + self.query = .unnamed(query, promise) self.logger = logger - self.promise = promise } - init(executeStatement: PSQLExecuteStatement, - logger: Logger, - promise: EventLoopPromise) - { - self.query = .preparedStatement(executeStatement) + init( + executeStatement: PSQLExecuteStatement, + logger: Logger, + promise: EventLoopPromise + ) { + self.query = .executeStatement(executeStatement, promise) + self.logger = logger + } + + init( + name: String, + query: String, + bindingDataTypes: [PostgresDataType], + logger: Logger, + promise: EventLoopPromise + ) { + self.query = .prepareStatement(name: name, query: query, bindingDataTypes: bindingDataTypes, promise) self.logger = logger - self.promise = promise } } -final class PrepareStatementContext { +final class PreparedStatementContext: Sendable { let name: String - let query: String + let sql: String + let bindingDataTypes: [PostgresDataType] + let bindings: PostgresBindings let logger: Logger - let promise: EventLoopPromise - - init(name: String, - query: String, - logger: Logger, - promise: EventLoopPromise) - { + let promise: EventLoopPromise + + init( + name: String, + sql: String, + bindings: PostgresBindings, + bindingDataTypes: [PostgresDataType], + logger: Logger, + promise: EventLoopPromise + ) { self.name = name - self.query = query + self.sql = sql + self.bindings = bindings + if bindingDataTypes.isEmpty { + self.bindingDataTypes = bindings.metadata.map(\.dataType) + } else { + self.bindingDataTypes = bindingDataTypes + } self.logger = logger self.promise = promise } } -final class CloseCommandContext { +final class CloseCommandContext: Sendable { let target: CloseTarget let logger: Logger let promise: EventLoopPromise - init(target: CloseTarget, - logger: Logger, - promise: EventLoopPromise) - { + init( + target: CloseTarget, + logger: Logger, + promise: EventLoopPromise + ) { self.target = target self.logger = logger self.promise = promise diff --git a/Sources/PostgresNIO/New/PostgresBackendMessage.swift b/Sources/PostgresNIO/New/PostgresBackendMessage.swift index ecccd1e9..792beec3 100644 --- a/Sources/PostgresNIO/New/PostgresBackendMessage.swift +++ b/Sources/PostgresNIO/New/PostgresBackendMessage.swift @@ -20,7 +20,7 @@ protocol PSQLMessagePayloadDecodable { /// /// All messages are defined in the official Postgres Documentation in the section /// [Frontend/Backend Protocol – Message Formats](https://www.postgresql.org/docs/13/protocol-message-formats.html) -enum PostgresBackendMessage { +enum PostgresBackendMessage: Hashable { typealias PayloadDecodable = PSQLMessagePayloadDecodable @@ -46,141 +46,31 @@ enum PostgresBackendMessage { } extension PostgresBackendMessage { - enum ID: RawRepresentable, Equatable { - typealias RawValue = UInt8 - - case authentication - case backendKeyData - case bindComplete - case closeComplete - case commandComplete - case copyData - case copyDone - case copyInResponse - case copyOutResponse - case copyBothResponse - case dataRow - case emptyQueryResponse - case error - case functionCallResponse - case negotiateProtocolVersion - case noData - case noticeResponse - case notificationResponse - case parameterDescription - case parameterStatus - case parseComplete - case portalSuspended - case readyForQuery - case rowDescription - - init?(rawValue: UInt8) { - switch rawValue { - case UInt8(ascii: "R"): - self = .authentication - case UInt8(ascii: "K"): - self = .backendKeyData - case UInt8(ascii: "2"): - self = .bindComplete - case UInt8(ascii: "3"): - self = .closeComplete - case UInt8(ascii: "C"): - self = .commandComplete - case UInt8(ascii: "d"): - self = .copyData - case UInt8(ascii: "c"): - self = .copyDone - case UInt8(ascii: "G"): - self = .copyInResponse - case UInt8(ascii: "H"): - self = .copyOutResponse - case UInt8(ascii: "W"): - self = .copyBothResponse - case UInt8(ascii: "D"): - self = .dataRow - case UInt8(ascii: "I"): - self = .emptyQueryResponse - case UInt8(ascii: "E"): - self = .error - case UInt8(ascii: "V"): - self = .functionCallResponse - case UInt8(ascii: "v"): - self = .negotiateProtocolVersion - case UInt8(ascii: "n"): - self = .noData - case UInt8(ascii: "N"): - self = .noticeResponse - case UInt8(ascii: "A"): - self = .notificationResponse - case UInt8(ascii: "t"): - self = .parameterDescription - case UInt8(ascii: "S"): - self = .parameterStatus - case UInt8(ascii: "1"): - self = .parseComplete - case UInt8(ascii: "s"): - self = .portalSuspended - case UInt8(ascii: "Z"): - self = .readyForQuery - case UInt8(ascii: "T"): - self = .rowDescription - default: - return nil - } - } - - var rawValue: UInt8 { - switch self { - case .authentication: - return UInt8(ascii: "R") - case .backendKeyData: - return UInt8(ascii: "K") - case .bindComplete: - return UInt8(ascii: "2") - case .closeComplete: - return UInt8(ascii: "3") - case .commandComplete: - return UInt8(ascii: "C") - case .copyData: - return UInt8(ascii: "d") - case .copyDone: - return UInt8(ascii: "c") - case .copyInResponse: - return UInt8(ascii: "G") - case .copyOutResponse: - return UInt8(ascii: "H") - case .copyBothResponse: - return UInt8(ascii: "W") - case .dataRow: - return UInt8(ascii: "D") - case .emptyQueryResponse: - return UInt8(ascii: "I") - case .error: - return UInt8(ascii: "E") - case .functionCallResponse: - return UInt8(ascii: "V") - case .negotiateProtocolVersion: - return UInt8(ascii: "v") - case .noData: - return UInt8(ascii: "n") - case .noticeResponse: - return UInt8(ascii: "N") - case .notificationResponse: - return UInt8(ascii: "A") - case .parameterDescription: - return UInt8(ascii: "t") - case .parameterStatus: - return UInt8(ascii: "S") - case .parseComplete: - return UInt8(ascii: "1") - case .portalSuspended: - return UInt8(ascii: "s") - case .readyForQuery: - return UInt8(ascii: "Z") - case .rowDescription: - return UInt8(ascii: "T") - } - } + enum ID: UInt8, Hashable { + case authentication = 82 // ascii: R + case backendKeyData = 75 // ascii: K + case bindComplete = 50 // ascii: 2 + case closeComplete = 51 // ascii: 3 + case commandComplete = 67 // ascii: C + case copyData = 100 // ascii: d + case copyDone = 99 // ascii: c + case copyInResponse = 71 // ascii: G + case copyOutResponse = 72 // ascii: H + case copyBothResponse = 87 // ascii: W + case dataRow = 68 // ascii: D + case emptyQueryResponse = 73 // ascii: I + case error = 69 // ascii: E + case functionCallResponse = 86 // ascii: V + case negotiateProtocolVersion = 118 // ascii: v + case noData = 110 // ascii: n + case noticeResponse = 78 // ascii: N + case notificationResponse = 65 // ascii: A + case parameterDescription = 116 // ascii: t + case parameterStatus = 83 // ascii: S + case parseComplete = 49 // ascii: 1 + case portalSuspended = 115 // ascii: s + case readyForQuery = 90 // ascii: Z + case rowDescription = 84 // ascii: T } } diff --git a/Sources/PostgresNIO/New/PostgresBackendMessageDecoder.swift b/Sources/PostgresNIO/New/PostgresBackendMessageDecoder.swift index ee7e1b84..6f6be7ec 100644 --- a/Sources/PostgresNIO/New/PostgresBackendMessageDecoder.swift +++ b/Sources/PostgresNIO/New/PostgresBackendMessageDecoder.swift @@ -107,8 +107,8 @@ struct PostgresMessageDecodingError: Error { static func withPartialError( _ partialError: PSQLPartialDecodingError, messageID: UInt8, - messageBytes: ByteBuffer) -> Self - { + messageBytes: ByteBuffer + ) -> Self { var byteBuffer = messageBytes let data = byteBuffer.readData(length: byteBuffer.readableBytes)! @@ -124,8 +124,8 @@ struct PostgresMessageDecodingError: Error { messageID: UInt8, messageBytes: ByteBuffer, file: String = #fileID, - line: Int = #line) -> Self - { + line: Int = #line + ) -> Self { var byteBuffer = messageBytes let data = byteBuffer.readData(length: byteBuffer.readableBytes)! @@ -153,8 +153,8 @@ struct PSQLPartialDecodingError: Error { value: Target.RawValue, asType: Target.Type, file: String = #fileID, - line: Int = #line) -> Self - { + line: Int = #line + ) -> Self { return PSQLPartialDecodingError( description: "Can not represent '\(value)' with type '\(asType)'.", file: file, line: line) diff --git a/Sources/PostgresNIO/New/PostgresChannelHandler.swift b/Sources/PostgresNIO/New/PostgresChannelHandler.swift index 84f07d47..0a14849a 100644 --- a/Sources/PostgresNIO/New/PostgresChannelHandler.swift +++ b/Sources/PostgresNIO/New/PostgresChannelHandler.swift @@ -3,16 +3,13 @@ import NIOTLS import Crypto import Logging -protocol PSQLChannelHandlerNotificationDelegate: AnyObject { - func notificationReceived(_: PostgresBackendMessage.NotificationResponse) -} - final class PostgresChannelHandler: ChannelDuplexHandler { - typealias OutboundIn = PSQLTask + typealias OutboundIn = HandlerTask typealias InboundIn = ByteBuffer typealias OutboundOut = ByteBuffer private let logger: Logger + private let eventLoop: EventLoop private var state: ConnectionStateMachine /// A `ChannelHandlerContext` to be used for non channel related events. (for example: More rows needed). @@ -21,18 +18,21 @@ final class PostgresChannelHandler: ChannelDuplexHandler { private var handlerContext: ChannelHandlerContext? private var rowStream: PSQLRowStream? private var decoder: NIOSingleStepByteToMessageProcessor - private var encoder: BufferedMessageEncoder! + private var encoder: PostgresFrontendMessageEncoder! private let configuration: PostgresConnection.InternalConfiguration - private let configureSSLCallback: ((Channel) throws -> Void)? - - /// this delegate should only be accessed on the connections `EventLoop` - weak var notificationDelegate: PSQLChannelHandlerNotificationDelegate? - - init(configuration: PostgresConnection.InternalConfiguration, - logger: Logger, - configureSSLCallback: ((Channel) throws -> Void)?) - { + private let configureSSLCallback: ((Channel, PostgresChannelHandler) throws -> Void)? + + private var listenState = ListenStateMachine() + private var preparedStatementState = PreparedStatementStateMachine() + + init( + configuration: PostgresConnection.InternalConfiguration, + eventLoop: EventLoop, + logger: Logger, + configureSSLCallback: ((Channel, PostgresChannelHandler) throws -> Void)? + ) { self.state = ConnectionStateMachine(requireBackendKeyData: configuration.options.requireBackendKeyData) + self.eventLoop = eventLoop self.configuration = configuration self.configureSSLCallback = configureSSLCallback self.logger = logger @@ -41,12 +41,15 @@ final class PostgresChannelHandler: ChannelDuplexHandler { #if DEBUG /// for testing purposes only - init(configuration: PostgresConnection.InternalConfiguration, - state: ConnectionStateMachine = .init(.initialized), - logger: Logger = .psqlNoOpLogger, - configureSSLCallback: ((Channel) throws -> Void)?) - { + init( + configuration: PostgresConnection.InternalConfiguration, + eventLoop: EventLoop, + state: ConnectionStateMachine = .init(.initialized), + logger: Logger = .psqlNoOpLogger, + configureSSLCallback: ((Channel, PostgresChannelHandler) throws -> Void)? + ) { self.state = state + self.eventLoop = eventLoop self.configuration = configuration self.configureSSLCallback = configureSSLCallback self.logger = logger @@ -58,10 +61,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler { func handlerAdded(context: ChannelHandlerContext) { self.handlerContext = context - self.encoder = BufferedMessageEncoder( - buffer: context.channel.allocator.buffer(capacity: 256), - encoder: PSQLFrontendMessageEncoder() - ) + self.encoder = PostgresFrontendMessageEncoder(buffer: context.channel.allocator.buffer(capacity: 256)) if context.channel.isActive { self.connected(context: context) @@ -84,6 +84,17 @@ final class PostgresChannelHandler: ChannelDuplexHandler { } func channelInactive(context: ChannelHandlerContext) { + do { + try self.decoder.finishProcessing(seenEOF: true) { message in + self.handleMessage(message, context: context) + } + } catch let error as PostgresMessageDecodingError { + let action = self.state.errorHappened(.messageDecodingFailure(error)) + self.run(action, with: context) + } catch { + preconditionFailure("Expected to only get PSQLDecodingErrors from the PSQLBackendMessageDecoder.") + } + self.logger.trace("Channel inactive.") let action = self.state.closed() self.run(action, with: context) @@ -100,51 +111,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler { do { try self.decoder.process(buffer: buffer) { message in - self.logger.trace("Backend message received", metadata: [.message: "\(message)"]) - let action: ConnectionStateMachine.ConnectionAction - - switch message { - case .authentication(let authentication): - action = self.state.authenticationMessageReceived(authentication) - case .backendKeyData(let keyData): - action = self.state.backendKeyDataReceived(keyData) - case .bindComplete: - action = self.state.bindCompleteReceived() - case .closeComplete: - action = self.state.closeCompletedReceived() - case .commandComplete(let commandTag): - action = self.state.commandCompletedReceived(commandTag) - case .dataRow(let dataRow): - action = self.state.dataRowReceived(dataRow) - case .emptyQueryResponse: - action = self.state.emptyQueryResponseReceived() - case .error(let errorResponse): - action = self.state.errorReceived(errorResponse) - case .noData: - action = self.state.noDataReceived() - case .notice(let noticeResponse): - action = self.state.noticeReceived(noticeResponse) - case .notification(let notification): - action = self.state.notificationReceived(notification) - case .parameterDescription(let parameterDescription): - action = self.state.parameterDescriptionReceived(parameterDescription) - case .parameterStatus(let parameterStatus): - action = self.state.parameterStatusReceived(parameterStatus) - case .parseComplete: - action = self.state.parseCompleteReceived() - case .portalSuspended: - action = self.state.portalSuspendedReceived() - case .readyForQuery(let transactionState): - action = self.state.readyForQueryReceived(transactionState) - case .rowDescription(let rowDescription): - action = self.state.rowDescriptionReceived(rowDescription) - case .sslSupported: - action = self.state.sslSupportedReceived(unprocessedBytes: self.decoder.unprocessedBytes) - case .sslUnsupported: - action = self.state.sslUnsupportedReceived() - } - - self.run(action, with: context) + self.handleMessage(message, context: context) } } catch let error as PostgresMessageDecodingError { let action = self.state.errorHappened(.messageDecodingFailure(error)) @@ -153,7 +120,55 @@ final class PostgresChannelHandler: ChannelDuplexHandler { preconditionFailure("Expected to only get PSQLDecodingErrors from the PSQLBackendMessageDecoder.") } } - + + private func handleMessage(_ message: PostgresBackendMessage, context: ChannelHandlerContext) { + self.logger.trace("Backend message received", metadata: [.message: "\(message)"]) + let action: ConnectionStateMachine.ConnectionAction + + switch message { + case .authentication(let authentication): + action = self.state.authenticationMessageReceived(authentication) + case .backendKeyData(let keyData): + action = self.state.backendKeyDataReceived(keyData) + case .bindComplete: + action = self.state.bindCompleteReceived() + case .closeComplete: + action = self.state.closeCompletedReceived() + case .commandComplete(let commandTag): + action = self.state.commandCompletedReceived(commandTag) + case .dataRow(let dataRow): + action = self.state.dataRowReceived(dataRow) + case .emptyQueryResponse: + action = self.state.emptyQueryResponseReceived() + case .error(let errorResponse): + action = self.state.errorReceived(errorResponse) + case .noData: + action = self.state.noDataReceived() + case .notice(let noticeResponse): + action = self.state.noticeReceived(noticeResponse) + case .notification(let notification): + action = self.state.notificationReceived(notification) + case .parameterDescription(let parameterDescription): + action = self.state.parameterDescriptionReceived(parameterDescription) + case .parameterStatus(let parameterStatus): + action = self.state.parameterStatusReceived(parameterStatus) + case .parseComplete: + action = self.state.parseCompleteReceived() + case .portalSuspended: + action = self.state.portalSuspendedReceived() + case .readyForQuery(let transactionState): + action = self.state.readyForQueryReceived(transactionState) + case .rowDescription(let rowDescription): + action = self.state.rowDescriptionReceived(rowDescription) + case .sslSupported: + action = self.state.sslSupportedReceived(unprocessedBytes: self.decoder.unprocessedBytes) + case .sslUnsupported: + action = self.state.sslUnsupportedReceived() + } + + self.run(action, with: context) + } + func channelReadComplete(context: ChannelHandlerContext) { let action = self.state.channelReadComplete() self.run(action, with: context) @@ -182,8 +197,67 @@ final class PostgresChannelHandler: ChannelDuplexHandler { } func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise?) { - let task = self.unwrapOutboundIn(data) - let action = self.state.enqueue(task: task) + let handlerTask = self.unwrapOutboundIn(data) + let psqlTask: PSQLTask + + switch handlerTask { + case .closeCommand(let command): + psqlTask = .closeCommand(command) + case .extendedQuery(let query): + psqlTask = .extendedQuery(query) + + case .startListening(let listener): + switch self.listenState.startListening(listener) { + case .startListening(let channel): + psqlTask = self.makeStartListeningQuery(channel: channel, context: context) + + case .none: + return + + case .succeedListenStart(let listener): + listener.startListeningSucceeded(handler: self) + return + } + + case .cancelListening(let channel, let id): + switch self.listenState.cancelNotificationListener(channel: channel, id: id) { + case .none: + return + + case .stopListening(let channel, let listener): + psqlTask = self.makeUnlistenQuery(channel: channel, context: context) + listener.failed(CancellationError()) + + case .cancelListener(let listener): + listener.failed(CancellationError()) + return + } + case .executePreparedStatement(let preparedStatement): + let action = self.preparedStatementState.lookup( + preparedStatement: preparedStatement + ) + switch action { + case .prepareStatement: + psqlTask = self.makePrepareStatementTask( + preparedStatement: preparedStatement, + context: context + ) + case .waitForAlreadyInFlightPreparation: + // The state machine already keeps track of this + // and will execute the statement as soon as it's prepared + return + case .executeStatement(let rowDescription): + psqlTask = self.makeExecutePreparedStatementTask( + preparedStatement: preparedStatement, + rowDescription: rowDescription + ) + case .returnError(let error): + preparedStatement.promise.fail(error) + return + } + } + + let action = self.state.enqueue(task: psqlTask) self.run(action, with: context) } @@ -195,7 +269,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler { return } - let action = self.state.close(promise) + let action = self.state.close(promise: promise) self.run(action, with: context) } @@ -206,14 +280,44 @@ final class PostgresChannelHandler: ChannelDuplexHandler { case PSQLOutgoingEvent.authenticate(let authContext): let action = self.state.provideAuthenticationContext(authContext) self.run(action, with: context) + + case PSQLOutgoingEvent.gracefulShutdown: + let action = self.state.gracefulClose(promise) + self.run(action, with: context) + default: context.triggerUserOutboundEvent(event, promise: promise) } } + // MARK: Listening + + func cancelNotificationListener(channel: String, id: Int) { + self.eventLoop.preconditionInEventLoop() + + switch self.listenState.cancelNotificationListener(channel: channel, id: id) { + case .cancelListener(let listener): + listener.cancelled() + + case .stopListening(let channel, cancelListener: let listener): + listener.cancelled() + + guard let context = self.handlerContext else { + return + } + + let query = self.makeUnlistenQuery(channel: channel, context: context) + let action = self.state.enqueue(task: query) + self.run(action, with: context) + + case .none: + break + } + } + // MARK: Channel handler actions - func run(_ action: ConnectionStateMachine.ConnectionAction, with context: ChannelHandlerContext) { + private func run(_ action: ConnectionStateMachine.ConnectionAction, with context: ChannelHandlerContext) { self.logger.trace("Run action", metadata: [.connectionAction: "\(action)"]) switch action { @@ -224,35 +328,33 @@ final class PostgresChannelHandler: ChannelDuplexHandler { case .wait: break case .sendStartupMessage(let authContext): - self.encoder.encode(.startup(.versionThree(parameters: authContext.toStartupParameters()))) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.startup(user: authContext.username, database: authContext.database, options: authContext.additionalParameters) + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) case .sendSSLRequest: - self.encoder.encode(.sslRequest(.init())) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.ssl() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) case .sendPasswordMessage(let mode, let authContext): self.sendPasswordMessage(mode: mode, authContext: authContext, context: context) case .sendSaslInitialResponse(let name, let initialResponse): - self.encoder.encode(.saslInitialResponse(.init(saslMechanism: name, initialData: initialResponse))) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.saslInitialResponse(mechanism: name, bytes: initialResponse) + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) case .sendSaslResponse(let bytes): - self.encoder.encode(.saslResponse(.init(data: bytes))) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.saslResponse(bytes) + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) case .closeConnectionAndCleanup(let cleanupContext): self.closeConnectionAndCleanup(cleanupContext, context: context) case .fireChannelInactive: context.fireChannelInactive() - case .sendParseDescribeSync(let name, let query): - self.sendParseDecribeAndSyncMessage(statementName: name, query: query, context: context) + case .sendParseDescribeSync(let name, let query, let bindingDataTypes): + self.sendParseDescribeAndSyncMessage(statementName: name, query: query, bindingDataTypes: bindingDataTypes, context: context) case .sendBindExecuteSync(let executeStatement): self.sendBindExecuteAndSyncMessage(executeStatement: executeStatement, context: context) case .sendParseDescribeBindExecuteSync(let query): self.sendParseDescribeBindExecuteAndSyncMessage(query: query, context: context) - case .succeedQuery(let queryContext, columns: let columns): - self.succeedQueryWithRowStream(queryContext, columns: columns, context: context) - case .succeedQueryNoRowsComming(let queryContext, let commandTag): - self.succeedQueryWithoutRowStream(queryContext, commandTag: commandTag, context: context) - case .failQuery(let queryContext, with: let error, let cleanupContext): - queryContext.promise.fail(error) + case .succeedQuery(let promise, with: let result): + self.succeedQuery(promise, result: result, context: context) + case .failQuery(let promise, with: let error, let cleanupContext): + promise.fail(error) if let cleanupContext = cleanupContext { self.closeConnectionAndCleanup(cleanupContext, context: context) } @@ -288,7 +390,8 @@ final class PostgresChannelHandler: ChannelDuplexHandler { let authContext = AuthContext( username: username, password: self.configuration.password, - database: self.configuration.database + database: self.configuration.database, + additionalParameters: self.configuration.options.additionalStartupParameters ) let action = self.state.provideAuthenticationContext(authContext) return self.run(action, with: context) @@ -300,14 +403,14 @@ final class PostgresChannelHandler: ChannelDuplexHandler { // The normal, graceful termination procedure is that the frontend sends a Terminate // message and immediately closes the connection. On receipt of this message, the // backend closes the connection and terminates. - self.encoder.encode(.terminate) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.terminate() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) } context.close(mode: .all, promise: promise) - case .succeedPreparedStatementCreation(let preparedContext, with: let rowDescription): - preparedContext.promise.succeed(rowDescription) - case .failPreparedStatementCreation(let preparedContext, with: let error, let cleanupContext): - preparedContext.promise.fail(error) + case .succeedPreparedStatementCreation(let promise, with: let rowDescription): + promise.succeed(rowDescription) + case .failPreparedStatementCreation(let promise, with: let error, let cleanupContext): + promise.fail(error) if let cleanupContext = cleanupContext { self.closeConnectionAndCleanup(cleanupContext, context: context) } @@ -321,16 +424,14 @@ final class PostgresChannelHandler: ChannelDuplexHandler { self.closeConnectionAndCleanup(cleanupContext, context: context) } case .forwardNotificationToListeners(let notification): - self.notificationDelegate?.notificationReceived(notification) + self.forwardNotificationToListeners(notification, context: context) } } // MARK: - Private Methods - private func connected(context: ChannelHandlerContext) { - let action = self.state.connected(tls: .init(self.configuration.tls)) - self.run(action, with: context) } @@ -338,7 +439,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler { // This method must only be called, if we signalized the StateMachine before that we are // able to setup a SSL connection. do { - try self.configureSSLCallback!(context.channel) + try self.configureSSLCallback!(context.channel, self) let action = self.state.sslHandlerAdded() self.run(action, with: context) } catch { @@ -350,8 +451,8 @@ final class PostgresChannelHandler: ChannelDuplexHandler { private func sendPasswordMessage( mode: PasswordAuthencationMode, authContext: AuthContext, - context: ChannelHandlerContext) - { + context: ChannelHandlerContext + ) { switch mode { case .md5(let salt): let hash1 = (authContext.password ?? "") + authContext.username @@ -360,134 +461,128 @@ final class PostgresChannelHandler: ChannelDuplexHandler { var hash2 = [UInt8]() hash2.reserveCapacity(pwdhash.count + 4) hash2.append(contentsOf: pwdhash) - hash2.append(salt.0) - hash2.append(salt.1) - hash2.append(salt.2) - hash2.append(salt.3) + var saltNetworkOrder = salt.bigEndian + withUnsafeBytes(of: &saltNetworkOrder) { ptr in + hash2.append(contentsOf: ptr) + } let hash = Insecure.MD5.hash(data: hash2).md5PrefixHexdigest() - self.encoder.encode(.password(.init(value: hash))) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.password(hash.utf8) + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) case .cleartext: - self.encoder.encode(.password(.init(value: authContext.password ?? ""))) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.password((authContext.password ?? "").utf8) + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) } } private func sendCloseAndSyncMessage(_ sendClose: CloseTarget, context: ChannelHandlerContext) { switch sendClose { case .preparedStatement(let name): - self.encoder.encode(.close(.preparedStatement(name))) - self.encoder.encode(.sync) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.closePreparedStatement(name) + self.encoder.sync() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) case .portal(let name): - self.encoder.encode(.close(.portal(name))) - self.encoder.encode(.sync) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.closePortal(name) + self.encoder.sync() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) } } - private func sendParseDecribeAndSyncMessage( + private func sendParseDescribeAndSyncMessage( statementName: String, query: String, - context: ChannelHandlerContext) - { + bindingDataTypes: [PostgresDataType], + context: ChannelHandlerContext + ) { precondition(self.rowStream == nil, "Expected to not have an open stream at this point") - let parse = PostgresFrontendMessage.Parse( - preparedStatementName: statementName, - query: query, - parameters: []) - - self.encoder.encode(.parse(parse)) - self.encoder.encode(.describe(.preparedStatement(statementName))) - self.encoder.encode(.sync) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + self.encoder.parse(preparedStatementName: statementName, query: query, parameters: bindingDataTypes) + self.encoder.describePreparedStatement(statementName) + self.encoder.sync() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) } private func sendBindExecuteAndSyncMessage( executeStatement: PSQLExecuteStatement, context: ChannelHandlerContext ) { - let bind = PostgresFrontendMessage.Bind( + self.encoder.bind( portalName: "", preparedStatementName: executeStatement.name, - bind: executeStatement.binds) - - self.encoder.encode(.bind(bind)) - self.encoder.encode(.execute(.init(portalName: ""))) - self.encoder.encode(.sync) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) + bind: executeStatement.binds + ) + self.encoder.execute(portalName: "") + self.encoder.sync() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) } private func sendParseDescribeBindExecuteAndSyncMessage( query: PostgresQuery, - context: ChannelHandlerContext) - { + context: ChannelHandlerContext + ) { precondition(self.rowStream == nil, "Expected to not have an open stream at this point") let unnamedStatementName = "" - let parse = PostgresFrontendMessage.Parse( + self.encoder.parse( preparedStatementName: unnamedStatementName, query: query.sql, - parameters: query.binds.metadata.map(\.dataType)) - let bind = PostgresFrontendMessage.Bind( - portalName: "", - preparedStatementName: unnamedStatementName, - bind: query.binds) - - self.encoder.encode(.parse(parse)) - self.encoder.encode(.describe(.preparedStatement(""))) - self.encoder.encode(.bind(bind)) - self.encoder.encode(.execute(.init(portalName: ""))) - self.encoder.encode(.sync) - context.writeAndFlush(self.wrapOutboundOut(self.encoder.flush()), promise: nil) - } - - private func succeedQueryWithRowStream( - _ queryContext: ExtendedQueryContext, - columns: [RowDescription.Column], - context: ChannelHandlerContext) - { - let rows = PSQLRowStream( - rowDescription: columns, - queryContext: queryContext, - eventLoop: context.channel.eventLoop, - rowSource: .stream(self)) - - self.rowStream = rows - queryContext.promise.succeed(rows) + parameters: query.binds.metadata.lazy.map(\.dataType) + ) + self.encoder.describePreparedStatement(unnamedStatementName) + self.encoder.bind(portalName: "", preparedStatementName: unnamedStatementName, bind: query.binds) + self.encoder.execute(portalName: "") + self.encoder.sync() + context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil) } - private func succeedQueryWithoutRowStream( - _ queryContext: ExtendedQueryContext, - commandTag: String, - context: ChannelHandlerContext) - { - let rows = PSQLRowStream( - rowDescription: [], - queryContext: queryContext, - eventLoop: context.channel.eventLoop, - rowSource: .noRows(.success(commandTag)) - ) - queryContext.promise.succeed(rows) + private func succeedQuery( + _ promise: EventLoopPromise, + result: QueryResult, + context: ChannelHandlerContext + ) { + let rows: PSQLRowStream + switch result.value { + case .rowDescription(let columns): + rows = PSQLRowStream( + source: .stream(columns, self), + eventLoop: context.channel.eventLoop, + logger: result.logger + ) + self.rowStream = rows + + case .noRows(let summary): + rows = PSQLRowStream( + source: .noRows(.success(summary)), + eventLoop: context.channel.eventLoop, + logger: result.logger + ) + } + + promise.succeed(rows) } private func closeConnectionAndCleanup( _ cleanup: ConnectionStateMachine.ConnectionAction.CleanUpContext, - context: ChannelHandlerContext) - { + context: ChannelHandlerContext + ) { self.logger.debug("Cleaning up and closing connection.", metadata: [.error: "\(cleanup.error)"]) // 1. fail all tasks cleanup.tasks.forEach { task in task.failWithError(cleanup.error) } - - // 2. fire an error - context.fireErrorCaught(cleanup.error) - - // 3. close the connection or fire channel inactive + + // 2. stop all listeners + for listener in self.listenState.fail(cleanup.error) { + listener.failed(cleanup.error) + } + + // 3. fire an error + if cleanup.error.code != .clientClosedConnection { + context.fireErrorCaught(cleanup.error) + } + + // 4. close the connection or fire channel inactive switch cleanup.action { case .close: context.close(mode: .all, promise: cleanup.closePromise) @@ -496,6 +591,199 @@ final class PostgresChannelHandler: ChannelDuplexHandler { context.fireChannelInactive() } } + + private func makeStartListeningQuery(channel: String, context: ChannelHandlerContext) -> PSQLTask { + let promise = context.eventLoop.makePromise(of: PSQLRowStream.self) + let query = ExtendedQueryContext( + query: PostgresQuery(unsafeSQL: #"LISTEN "\#(channel)";"#), + logger: self.logger, + promise: promise + ) + let loopBound = NIOLoopBound((self, context), eventLoop: self.eventLoop) + promise.futureResult.whenComplete { result in + let (selfTransferred, context) = loopBound.value + selfTransferred.startListenCompleted(result, for: channel, context: context) + } + + return .extendedQuery(query) + } + + private func startListenCompleted(_ result: Result, for channel: String, context: ChannelHandlerContext) { + switch result { + case .success: + switch self.listenState.startListeningSucceeded(channel: channel) { + case .activateListeners(let listeners): + for list in listeners { + list.startListeningSucceeded(handler: self) + } + + case .stopListening: + let task = self.makeUnlistenQuery(channel: channel, context: context) + let action = self.state.enqueue(task: task) + self.run(action, with: context) + } + + case .failure(let error): + let finalError: PSQLError + if var psqlError = error as? PSQLError { + psqlError.code = .listenFailed + finalError = psqlError + } else { + var psqlError = PSQLError(code: .listenFailed) + psqlError.underlying = error + finalError = psqlError + } + let listeners = self.listenState.startListeningFailed(channel: channel, error: finalError) + for list in listeners { + list.failed(finalError) + } + } + } + + private func makeUnlistenQuery(channel: String, context: ChannelHandlerContext) -> PSQLTask { + let promise = context.eventLoop.makePromise(of: PSQLRowStream.self) + let query = ExtendedQueryContext( + query: PostgresQuery(unsafeSQL: #"UNLISTEN "\#(channel)";"#), + logger: self.logger, + promise: promise + ) + let loopBound = NIOLoopBound((self, context), eventLoop: self.eventLoop) + promise.futureResult.whenComplete { result in + let (selfTransferred, context) = loopBound.value + selfTransferred.stopListenCompleted(result, for: channel, context: context) + } + + return .extendedQuery(query) + } + + private func stopListenCompleted( + _ result: Result, + for channel: String, + context: ChannelHandlerContext + ) { + switch result { + case .success: + switch self.listenState.stopListeningSucceeded(channel: channel) { + case .none: + break + + case .startListening: + let task = self.makeStartListeningQuery(channel: channel, context: context) + let action = self.state.enqueue(task: task) + self.run(action, with: context) + } + + case .failure(let error): + let action = self.state.errorHappened(.unlistenError(underlying: error)) + self.run(action, with: context) + } + } + + private func forwardNotificationToListeners( + _ notification: PostgresBackendMessage.NotificationResponse, + context: ChannelHandlerContext + ) { + switch self.listenState.notificationReceived(channel: notification.channel) { + case .none: + break + + case .notify(let listeners): + for listener in listeners { + listener.notificationReceived(notification) + } + } + } + + private func makePrepareStatementTask( + preparedStatement: PreparedStatementContext, + context: ChannelHandlerContext + ) -> PSQLTask { + let promise = self.eventLoop.makePromise(of: RowDescription?.self) + let loopBound = NIOLoopBound((self, context), eventLoop: self.eventLoop) + promise.futureResult.whenComplete { result in + let (selfTransferred, context) = loopBound.value + switch result { + case .success(let rowDescription): + selfTransferred.prepareStatementComplete( + name: preparedStatement.name, + rowDescription: rowDescription, + context: context + ) + case .failure(let error): + let psqlError: PSQLError + if let error = error as? PSQLError { + psqlError = error + } else { + psqlError = .connectionError(underlying: error) + } + selfTransferred.prepareStatementFailed( + name: preparedStatement.name, + error: psqlError, + context: context + ) + } + } + return .extendedQuery(.init( + name: preparedStatement.name, + query: preparedStatement.sql, + bindingDataTypes: preparedStatement.bindingDataTypes, + logger: preparedStatement.logger, + promise: promise + )) + } + + private func makeExecutePreparedStatementTask( + preparedStatement: PreparedStatementContext, + rowDescription: RowDescription? + ) -> PSQLTask { + return .extendedQuery(.init( + executeStatement: .init( + name: preparedStatement.name, + binds: preparedStatement.bindings, + rowDescription: rowDescription + ), + logger: preparedStatement.logger, + promise: preparedStatement.promise + )) + } + + private func prepareStatementComplete( + name: String, + rowDescription: RowDescription?, + context: ChannelHandlerContext + ) { + let action = self.preparedStatementState.preparationComplete( + name: name, + rowDescription: rowDescription + ) + for preparedStatement in action.statements { + let action = self.state.enqueue(task: .extendedQuery(.init( + executeStatement: .init( + name: preparedStatement.name, + binds: preparedStatement.bindings, + rowDescription: action.rowDescription + ), + logger: preparedStatement.logger, + promise: preparedStatement.promise + )) + ) + self.run(action, with: context) + } + } + + private func prepareStatementFailed( + name: String, + error: PSQLError, + context: ChannelHandlerContext + ) { + let action = self.preparedStatementState.errorHappened( + name: name, + error: error + ) + for statement in action.statements { + statement.promise.fail(action.error) + } + } } extension PostgresChannelHandler: PSQLRowsDataSource { @@ -516,17 +804,6 @@ extension PostgresChannelHandler: PSQLRowsDataSource { } } -extension AuthContext { - func toStartupParameters() -> PostgresFrontendMessage.Startup.Parameters { - PostgresFrontendMessage.Startup.Parameters( - user: self.username, - database: self.database, - options: nil, - replication: .false - ) - } -} - private extension Insecure.MD5.Digest { private static let lowercaseLookup: [UInt8] = [ @@ -576,16 +853,3 @@ extension ConnectionStateMachine.TLSConfiguration { } } } - -extension PostgresChannelHandler { - convenience init( - configuration: PostgresConnection.InternalConfiguration, - configureSSLCallback: ((Channel) throws -> Void)?) - { - self.init( - configuration: configuration, - logger: .psqlNoOpLogger, - configureSSLCallback: configureSSLCallback - ) - } -} diff --git a/Sources/PostgresNIO/New/PostgresCodable.swift b/Sources/PostgresNIO/New/PostgresCodable.swift index 3aa1a24f..fd82c8ea 100644 --- a/Sources/PostgresNIO/New/PostgresCodable.swift +++ b/Sources/PostgresNIO/New/PostgresCodable.swift @@ -2,29 +2,62 @@ import NIOCore import class Foundation.JSONEncoder import class Foundation.JSONDecoder +/// A type that can encode itself to a Postgres wire binary representation. +/// Dynamic types are types that don't have a well-known Postgres type OID at compile time. +/// For example, custom types created at runtime, such as enums, or extension types whose OID is not stable between +/// databases. +public protocol PostgresThrowingDynamicTypeEncodable { + /// The data type encoded into the `byteBuffer` in ``encode(into:context:)`` + var psqlType: PostgresDataType { get } + + /// The Postgres encoding format used to encode the value into `byteBuffer` in ``encode(into:context:)``. + var psqlFormat: PostgresFormat { get } + + /// Encode the entity into ``byteBuffer`` in the format specified by ``psqlFormat``, + /// using the provided ``context`` as needed, without setting the byte count. + /// + /// This method is called by ``PostgresBindings``. + func encode( + into byteBuffer: inout ByteBuffer, + context: PostgresEncodingContext + ) throws +} + +/// A type that can encode itself to a Postgres wire binary representation. +/// Dynamic types are types that don't have a well-known Postgres type OID at compile time. +/// For example, custom types created at runtime, such as enums, or extension types whose OID is not stable between +/// databases. +/// +/// This is the non-throwing alternative to ``PostgresThrowingDynamicTypeEncodable``. It allows users +/// to create ``PostgresQuery``s via `ExpressibleByStringInterpolation` without having to spell `try`. +public protocol PostgresDynamicTypeEncodable: PostgresThrowingDynamicTypeEncodable { + /// Encode the entity into ``byteBuffer`` in the format specified by ``psqlFormat``, + /// using the provided ``context`` as needed, without setting the byte count. + /// + /// This method is called by ``PostgresBindings``. + func encode( + into byteBuffer: inout ByteBuffer, + context: PostgresEncodingContext + ) +} + /// A type that can encode itself to a postgres wire binary representation. -public protocol PostgresEncodable { +public protocol PostgresEncodable: PostgresThrowingDynamicTypeEncodable { // TODO: Rename to `PostgresThrowingEncodable` with next major release - /// identifies the data type that we will encode into `byteBuffer` in `encode` + /// The data type encoded into the `byteBuffer` in ``encode(into:context:)``. static var psqlType: PostgresDataType { get } - /// identifies the postgres format that is used to encode the value into `byteBuffer` in `encode` + /// The Postgres encoding format used to encode the value into `byteBuffer` in ``encode(into:context:)``. static var psqlFormat: PostgresFormat { get } - - /// Encode the entity into the `byteBuffer` in Postgres binary format, without setting - /// the byte count. This method is called from the ``PostgresBindings``. - func encode(into byteBuffer: inout ByteBuffer, context: PostgresEncodingContext) throws } /// A type that can encode itself to a postgres wire binary representation. It enforces that the /// ``PostgresEncodable/encode(into:context:)-1jkcp`` does not throw. This allows users -/// to create ``PostgresQuery``s using the `ExpressibleByStringInterpolation` without +/// to create ``PostgresQuery``s via `ExpressibleByStringInterpolation` without /// having to spell `try`. -public protocol PostgresNonThrowingEncodable: PostgresEncodable { +public protocol PostgresNonThrowingEncodable: PostgresEncodable, PostgresDynamicTypeEncodable { // TODO: Rename to `PostgresEncodable` with next major release - - func encode(into byteBuffer: inout ByteBuffer, context: PostgresEncodingContext) } /// A type that can decode itself from a postgres wire binary representation. @@ -81,9 +114,17 @@ extension PostgresDecodable { } /// A type that can be encoded into and decoded from a postgres binary format -typealias PostgresCodable = PostgresEncodable & PostgresDecodable +public typealias PostgresCodable = PostgresEncodable & PostgresDecodable extension PostgresEncodable { + @inlinable + public var psqlType: PostgresDataType { Self.psqlType } + + @inlinable + public var psqlFormat: PostgresFormat { Self.psqlFormat } +} + +extension PostgresThrowingDynamicTypeEncodable { @inlinable func encodeRaw( into buffer: inout ByteBuffer, @@ -103,7 +144,7 @@ extension PostgresEncodable { } } -extension PostgresNonThrowingEncodable { +extension PostgresDynamicTypeEncodable { @inlinable func encodeRaw( into buffer: inout ByteBuffer, @@ -125,11 +166,10 @@ extension PostgresNonThrowingEncodable { /// A context that is passed to Swift objects that are encoded into the Postgres wire format. Used /// to pass further information to the encoding method. -public struct PostgresEncodingContext { +public struct PostgresEncodingContext: Sendable { /// A ``PostgresJSONEncoder`` used to encode the object to json. public var jsonEncoder: JSONEncoder - /// Creates a ``PostgresEncodingContext`` with the given ``PostgresJSONEncoder``. In case you want /// to use the a ``PostgresEncodingContext`` with an unconfigured Foundation `JSONEncoder` /// you can use the ``default`` context instead. @@ -147,7 +187,7 @@ extension PostgresEncodingContext where JSONEncoder == Foundation.JSONEncoder { /// A context that is passed to Swift objects that are decoded from the Postgres wire format. Used /// to pass further information to the decoding method. -public struct PostgresDecodingContext { +public struct PostgresDecodingContext: Sendable { /// A ``PostgresJSONDecoder`` used to decode the object from json. public var jsonDecoder: JSONDecoder diff --git a/Sources/PostgresNIO/New/PostgresFrontendMessage.swift b/Sources/PostgresNIO/New/PostgresFrontendMessage.swift deleted file mode 100644 index 2017cd1a..00000000 --- a/Sources/PostgresNIO/New/PostgresFrontendMessage.swift +++ /dev/null @@ -1,134 +0,0 @@ -import NIOCore - -/// A wire message that is created by a Postgres client to be consumed by Postgres server. -/// -/// All messages are defined in the official Postgres Documentation in the section -/// [Frontend/Backend Protocol – Message Formats](https://www.postgresql.org/docs/13/protocol-message-formats.html) -enum PostgresFrontendMessage: Equatable { - case bind(Bind) - case cancel(Cancel) - case close(Close) - case describe(Describe) - case execute(Execute) - case flush - case parse(Parse) - case password(Password) - case saslInitialResponse(SASLInitialResponse) - case saslResponse(SASLResponse) - case sslRequest(SSLRequest) - case sync - case startup(Startup) - case terminate - - enum ID: UInt8, Equatable { - - case bind - case close - case describe - case execute - case flush - case parse - case password - case saslInitialResponse - case saslResponse - case sync - case terminate - - init?(rawValue: UInt8) { - switch rawValue { - case UInt8(ascii: "B"): - self = .bind - case UInt8(ascii: "C"): - self = .close - case UInt8(ascii: "D"): - self = .describe - case UInt8(ascii: "E"): - self = .execute - case UInt8(ascii: "H"): - self = .flush - case UInt8(ascii: "P"): - self = .parse - case UInt8(ascii: "p"): - self = .password - case UInt8(ascii: "p"): - self = .saslInitialResponse - case UInt8(ascii: "p"): - self = .saslResponse - case UInt8(ascii: "S"): - self = .sync - case UInt8(ascii: "X"): - self = .terminate - default: - return nil - } - } - - var rawValue: UInt8 { - switch self { - case .bind: - return UInt8(ascii: "B") - case .close: - return UInt8(ascii: "C") - case .describe: - return UInt8(ascii: "D") - case .execute: - return UInt8(ascii: "E") - case .flush: - return UInt8(ascii: "H") - case .parse: - return UInt8(ascii: "P") - case .password: - return UInt8(ascii: "p") - case .saslInitialResponse: - return UInt8(ascii: "p") - case .saslResponse: - return UInt8(ascii: "p") - case .sync: - return UInt8(ascii: "S") - case .terminate: - return UInt8(ascii: "X") - } - } - } -} - -extension PostgresFrontendMessage { - - var id: ID { - switch self { - case .bind: - return .bind - case .cancel: - preconditionFailure("Cancel messages don't have an identifier") - case .close: - return .close - case .describe: - return .describe - case .execute: - return .execute - case .flush: - return .flush - case .parse: - return .parse - case .password: - return .password - case .saslInitialResponse: - return .saslInitialResponse - case .saslResponse: - return .saslResponse - case .sslRequest: - preconditionFailure("SSL requests don't have an identifier") - case .startup: - preconditionFailure("Startup messages don't have an identifier") - case .sync: - return .sync - case .terminate: - return .terminate - - } - } -} - -protocol PSQLMessagePayloadEncodable { - func encode(into buffer: inout ByteBuffer) -} diff --git a/Sources/PostgresNIO/New/PostgresFrontendMessageEncoder.swift b/Sources/PostgresNIO/New/PostgresFrontendMessageEncoder.swift new file mode 100644 index 00000000..97805418 --- /dev/null +++ b/Sources/PostgresNIO/New/PostgresFrontendMessageEncoder.swift @@ -0,0 +1,233 @@ +import NIOCore + +struct PostgresFrontendMessageEncoder { + + /// The SSL request code. The value is chosen to contain 1234 in the most significant 16 bits, + /// and 5679 in the least significant 16 bits. + static let sslRequestCode: Int32 = 80877103 + + /// The cancel request code. The value is chosen to contain 1234 in the most significant 16 bits, + /// and 5678 in the least significant 16 bits. (To avoid confusion, this code must not be the same + /// as any protocol version number.) + static let cancelRequestCode: Int32 = 80877102 + + static let startupVersionThree: Int32 = 0x00_03_00_00 + + private enum State { + case flushed + case writable + } + + private var buffer: ByteBuffer + private var state: State = .writable + + init(buffer: ByteBuffer) { + self.buffer = buffer + } + + mutating func startup(user: String, database: String?, options: [(String, String)]) { + self.clearIfNeeded() + self.buffer.psqlLengthPrefixed { buffer in + buffer.writeInteger(Self.startupVersionThree) + buffer.writeNullTerminatedString("user") + buffer.writeNullTerminatedString(user) + + if let database = database { + buffer.writeNullTerminatedString("database") + buffer.writeNullTerminatedString(database) + } + + // we don't send replication parameters, as the default is false and this is what we + // need for a client + for (key, value) in options { + buffer.writeNullTerminatedString(key) + buffer.writeNullTerminatedString(value) + } + + buffer.writeInteger(UInt8(0)) + } + } + + mutating func bind(portalName: String, preparedStatementName: String, bind: PostgresBindings) { + self.clearIfNeeded() + self.buffer.psqlLengthPrefixed(id: .bind) { buffer in + buffer.writeNullTerminatedString(portalName) + buffer.writeNullTerminatedString(preparedStatementName) + + // The number of parameter format codes that follow (denoted C below). This can be + // zero to indicate that there are no parameters or that the parameters all use the + // default format (text); or one, in which case the specified format code is applied + // to all parameters; or it can equal the actual number of parameters. + buffer.writeInteger(UInt16(bind.count)) + + // The parameter format codes. Each must presently be zero (text) or one (binary). + bind.metadata.forEach { + buffer.writeInteger($0.format.rawValue) + } + + buffer.writeInteger(UInt16(bind.count)) + + var parametersCopy = bind.bytes + buffer.writeBuffer(¶metersCopy) + + // The number of result-column format codes that follow (denoted R below). This can be + // zero to indicate that there are no result columns or that the result columns should + // all use the default format (text); or one, in which case the specified format code + // is applied to all result columns (if any); or it can equal the actual number of + // result columns of the query. + buffer.writeInteger(1, as: Int16.self) + // The result-column format codes. Each must presently be zero (text) or one (binary). + buffer.writeInteger(PostgresFormat.binary.rawValue, as: Int16.self) + } + } + + mutating func cancel(processID: Int32, secretKey: Int32) { + self.clearIfNeeded() + self.buffer.writeMultipleIntegers(UInt32(16), Self.cancelRequestCode, processID, secretKey) + } + + mutating func closePreparedStatement(_ preparedStatement: String) { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .close, length: UInt32(2 + preparedStatement.utf8.count), UInt8(ascii: "S")) + self.buffer.writeNullTerminatedString(preparedStatement) + } + + mutating func closePortal(_ portal: String) { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .close, length: UInt32(2 + portal.utf8.count), UInt8(ascii: "P")) + self.buffer.writeNullTerminatedString(portal) + } + + mutating func describePreparedStatement(_ preparedStatement: String) { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .describe, length: UInt32(2 + preparedStatement.utf8.count), UInt8(ascii: "S")) + self.buffer.writeNullTerminatedString(preparedStatement) + } + + mutating func describePortal(_ portal: String) { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .describe, length: UInt32(2 + portal.utf8.count), UInt8(ascii: "P")) + self.buffer.writeNullTerminatedString(portal) + } + + mutating func execute(portalName: String, maxNumberOfRows: Int32 = 0) { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .execute, length: UInt32(5 + portalName.utf8.count)) + self.buffer.writeNullTerminatedString(portalName) + self.buffer.writeInteger(maxNumberOfRows) + } + + mutating func parse(preparedStatementName: String, query: String, parameters: Parameters) where Parameters.Element == PostgresDataType { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers( + id: .parse, + length: UInt32(preparedStatementName.utf8.count + 1 + query.utf8.count + 1 + 2 + MemoryLayout.size * parameters.count) + ) + self.buffer.writeNullTerminatedString(preparedStatementName) + self.buffer.writeNullTerminatedString(query) + self.buffer.writeInteger(UInt16(parameters.count)) + + for dataType in parameters { + self.buffer.writeInteger(dataType.rawValue) + } + } + + mutating func password(_ bytes: Bytes) where Bytes.Element == UInt8 { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .password, length: UInt32(bytes.count) + 1) + self.buffer.writeBytes(bytes) + self.buffer.writeInteger(UInt8(0)) + } + + mutating func flush() { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .flush, length: 0) + } + + mutating func saslResponse(_ bytes: Bytes) where Bytes.Element == UInt8 { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .password, length: UInt32(bytes.count)) + self.buffer.writeBytes(bytes) + } + + mutating func saslInitialResponse(mechanism: String, bytes: Bytes) where Bytes.Element == UInt8 { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .password, length: UInt32(mechanism.utf8.count + 1 + 4 + bytes.count)) + self.buffer.writeNullTerminatedString(mechanism) + if bytes.count > 0 { + self.buffer.writeInteger(Int32(bytes.count)) + self.buffer.writeBytes(bytes) + } else { + self.buffer.writeInteger(Int32(-1)) + } + } + + mutating func ssl() { + self.clearIfNeeded() + self.buffer.writeMultipleIntegers(UInt32(8), Self.sslRequestCode) + } + + mutating func sync() { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .sync, length: 0) + } + + mutating func terminate() { + self.clearIfNeeded() + self.buffer.psqlWriteMultipleIntegers(id: .terminate, length: 0) + } + + mutating func flushBuffer() -> ByteBuffer { + self.state = .flushed + return self.buffer + } + + private mutating func clearIfNeeded() { + switch self.state { + case .flushed: + self.state = .writable + self.buffer.clear() + + case .writable: + break + } + } +} + +private enum FrontendMessageID: UInt8, Hashable, Sendable { + case bind = 66 // B + case close = 67 // C + case describe = 68 // D + case execute = 69 // E + case flush = 72 // H + case parse = 80 // P + case password = 112 // p - also both sasl values + case sync = 83 // S + case terminate = 88 // X +} + +extension ByteBuffer { + mutating fileprivate func psqlWriteMultipleIntegers(id: FrontendMessageID, length: UInt32) { + self.writeMultipleIntegers(id.rawValue, 4 + length) + } + + mutating fileprivate func psqlWriteMultipleIntegers(id: FrontendMessageID, length: UInt32, _ t1: T1) { + self.writeMultipleIntegers(id.rawValue, 4 + length, t1) + } + + mutating fileprivate func psqlLengthPrefixed(id: FrontendMessageID, _ encode: (inout ByteBuffer) -> ()) { + let lengthIndex = self.writerIndex + 1 + self.psqlWriteMultipleIntegers(id: id, length: 0) + encode(&self) + let length = UInt32(self.writerIndex - lengthIndex) + self.setInteger(length, at: lengthIndex) + } + + mutating fileprivate func psqlLengthPrefixed(_ encode: (inout ByteBuffer) -> ()) { + let lengthIndex = self.writerIndex + self.writeInteger(UInt32(0)) // placeholder + encode(&self) + let length = UInt32(self.writerIndex - lengthIndex) + self.setInteger(length, at: lengthIndex) + } +} diff --git a/Sources/PostgresNIO/New/PostgresNotificationSequence.swift b/Sources/PostgresNIO/New/PostgresNotificationSequence.swift new file mode 100644 index 00000000..d8f525eb --- /dev/null +++ b/Sources/PostgresNIO/New/PostgresNotificationSequence.swift @@ -0,0 +1,25 @@ + +public struct PostgresNotification: Sendable { + public let payload: String +} + +public struct PostgresNotificationSequence: AsyncSequence, Sendable { + public typealias Element = PostgresNotification + + let base: AsyncThrowingStream + + public func makeAsyncIterator() -> AsyncIterator { + AsyncIterator(base: self.base.makeAsyncIterator()) + } + + public struct AsyncIterator: AsyncIteratorProtocol { + var base: AsyncThrowingStream.AsyncIterator + + public mutating func next() async throws -> Element? { + try await self.base.next() + } + } +} + +@available(*, unavailable) +extension PostgresNotificationSequence.AsyncIterator: Sendable {} diff --git a/Sources/PostgresNIO/New/PostgresQuery.swift b/Sources/PostgresNIO/New/PostgresQuery.swift index 1ba75050..6449ab29 100644 --- a/Sources/PostgresNIO/New/PostgresQuery.swift +++ b/Sources/PostgresNIO/New/PostgresQuery.swift @@ -26,7 +26,7 @@ extension PostgresQuery: ExpressibleByStringInterpolation { } extension PostgresQuery { - public struct StringInterpolation: StringInterpolationProtocol { + public struct StringInterpolation: StringInterpolationProtocol, Sendable { public typealias StringLiteralType = String @usableFromInline @@ -44,13 +44,13 @@ extension PostgresQuery { } @inlinable - public mutating func appendInterpolation(_ value: Value) throws { + public mutating func appendInterpolation(_ value: Value) throws { try self.binds.append(value, context: .default) self.sql.append(contentsOf: "$\(self.binds.count)") } @inlinable - public mutating func appendInterpolation(_ value: Optional) throws { + public mutating func appendInterpolation(_ value: Optional) throws { switch value { case .none: self.binds.appendNull() @@ -62,13 +62,13 @@ extension PostgresQuery { } @inlinable - public mutating func appendInterpolation(_ value: Value) { + public mutating func appendInterpolation(_ value: Value) { self.binds.append(value, context: .default) self.sql.append(contentsOf: "$\(self.binds.count)") } @inlinable - public mutating func appendInterpolation(_ value: Optional) { + public mutating func appendInterpolation(_ value: Optional) { switch value { case .none: self.binds.appendNull() @@ -80,7 +80,7 @@ extension PostgresQuery { } @inlinable - public mutating func appendInterpolation( + public mutating func appendInterpolation( _ value: Value, context: PostgresEncodingContext ) throws { @@ -95,6 +95,20 @@ extension PostgresQuery { } } +extension PostgresQuery: CustomStringConvertible { + // See `CustomStringConvertible.description`. + public var description: String { + "\(self.sql) \(self.binds)" + } +} + +extension PostgresQuery: CustomDebugStringConvertible { + // See `CustomDebugStringConvertible.debugDescription`. + public var debugDescription: String { + "PostgresQuery(sql: \(String(describing: self.sql)), binds: \(String(reflecting: self.binds)))" + } +} + struct PSQLExecuteStatement { /// The statements name var name: String @@ -111,16 +125,19 @@ public struct PostgresBindings: Sendable, Hashable { var dataType: PostgresDataType @usableFromInline var format: PostgresFormat + @usableFromInline + var protected: Bool @inlinable - init(dataType: PostgresDataType, format: PostgresFormat) { + init(dataType: PostgresDataType, format: PostgresFormat, protected: Bool) { self.dataType = dataType self.format = format + self.protected = protected } @inlinable - init(value: Value) { - self.init(dataType: Value.psqlType, format: Value.psqlFormat) + init(value: Value, protected: Bool) { + self.init(dataType: value.psqlType, format: value.psqlFormat, protected: protected) } } @@ -147,25 +164,99 @@ public struct PostgresBindings: Sendable, Hashable { public mutating func appendNull() { self.bytes.writeInteger(-1, as: Int32.self) - self.metadata.append(.init(dataType: .null, format: .binary)) + self.metadata.append(.init(dataType: .null, format: .binary, protected: true)) + } + + @inlinable + public mutating func append(_ value: Value) throws { + try self.append(value, context: .default) + } + + @inlinable + public mutating func append(_ value: Optional) throws { + switch value { + case .none: + self.appendNull() + case let .some(value): + try self.append(value) + } + } + + @inlinable + public mutating func append( + _ value: Value, + context: PostgresEncodingContext + ) throws { + try value.encodeRaw(into: &self.bytes, context: context) + self.metadata.append(.init(value: value, protected: true)) + } + + @inlinable + public mutating func append( + _ value: Optional, + context: PostgresEncodingContext + ) throws { + switch value { + case .none: + self.appendNull() + case let .some(value): + try self.append(value, context: context) + } + } + + @inlinable + public mutating func append(_ value: Value) { + self.append(value, context: .default) + } + + @inlinable + public mutating func append(_ value: Optional) { + switch value { + case .none: + self.appendNull() + case let .some(value): + self.append(value) + } + } + + @inlinable + public mutating func append( + _ value: Value, + context: PostgresEncodingContext + ) { + value.encodeRaw(into: &self.bytes, context: context) + self.metadata.append(.init(value: value, protected: true)) } @inlinable - public mutating func append( + public mutating func append( + _ value: Optional, + context: PostgresEncodingContext + ) { + switch value { + case .none: + self.appendNull() + case let .some(value): + self.append(value, context: context) + } + } + + @inlinable + mutating func appendUnprotected( _ value: Value, context: PostgresEncodingContext ) throws { try value.encodeRaw(into: &self.bytes, context: context) - self.metadata.append(.init(value: value)) + self.metadata.append(.init(value: value, protected: false)) } @inlinable - public mutating func append( + mutating func appendUnprotected( _ value: Value, context: PostgresEncodingContext ) { value.encodeRaw(into: &self.bytes, context: context) - self.metadata.append(.init(value: value)) + self.metadata.append(.init(value: value, protected: false)) } public mutating func append(_ postgresData: PostgresData) { @@ -176,6 +267,93 @@ public struct PostgresBindings: Sendable, Hashable { self.bytes.writeInteger(Int32(input.readableBytes)) self.bytes.writeBuffer(&input) } - self.metadata.append(.init(dataType: postgresData.type, format: .binary)) + self.metadata.append(.init(dataType: postgresData.type, format: .binary, protected: true)) + } +} + +extension PostgresBindings: CustomStringConvertible, CustomDebugStringConvertible { + // See `CustomStringConvertible.description`. + public var description: String { + """ + [\(zip(self.metadata, BindingsReader(buffer: self.bytes)) + .lazy.map({ Self.makeBindingPrintable(protected: $0.protected, type: $0.dataType, format: $0.format, buffer: $1) }) + .joined(separator: ", "))] + """ + } + + // See `CustomDebugStringConvertible.description`. + public var debugDescription: String { + """ + [\(zip(self.metadata, BindingsReader(buffer: self.bytes)) + .lazy.map({ Self.makeDebugDescription(protected: $0.protected, type: $0.dataType, format: $0.format, buffer: $1) }) + .joined(separator: ", "))] + """ + } + + private static func makeDebugDescription(protected: Bool, type: PostgresDataType, format: PostgresFormat, buffer: ByteBuffer?) -> String { + "(\(Self.makeBindingPrintable(protected: protected, type: type, format: format, buffer: buffer)); \(type); format: \(format))" + } + + private static func makeBindingPrintable(protected: Bool, type: PostgresDataType, format: PostgresFormat, buffer: ByteBuffer?) -> String { + if protected { + return "****" + } + + guard var buffer = buffer else { + return "null" + } + + do { + switch (type, format) { + case (.int4, _), (.int2, _), (.int8, _): + let number = try Int64.init(from: &buffer, type: type, format: format, context: .default) + return String(describing: number) + + case (.bool, _): + let bool = try Bool.init(from: &buffer, type: type, format: format, context: .default) + return String(describing: bool) + + case (.varchar, _), (.bpchar, _), (.text, _), (.name, _): + let value = try String.init(from: &buffer, type: type, format: format, context: .default) + return String(reflecting: value) // adds quotes + + default: + return "\(buffer.readableBytes) bytes" + } + } catch { + return "\(buffer.readableBytes) bytes" + } + } +} + +/// A small helper to inspect encoded bindings +private struct BindingsReader: Sequence { + typealias Element = Optional + + var buffer: ByteBuffer + + struct Iterator: IteratorProtocol { + typealias Element = Optional + private var buffer: ByteBuffer + + init(buffer: ByteBuffer) { + self.buffer = buffer + } + + mutating func next() -> Optional> { + guard let length = self.buffer.readInteger(as: Int32.self) else { + return .none + } + + if length < 0 { + return .some(.none) + } + + return .some(self.buffer.readSlice(length: Int(length))!) + } + } + + func makeIterator() -> Iterator { + Iterator(buffer: self.buffer) } } diff --git a/Sources/PostgresNIO/New/PostgresRow-multi-decode.swift b/Sources/PostgresNIO/New/PostgresRow-multi-decode.swift deleted file mode 100644 index cb62c325..00000000 --- a/Sources/PostgresNIO/New/PostgresRow-multi-decode.swift +++ /dev/null @@ -1,1173 +0,0 @@ -/// NOTE: THIS FILE IS AUTO-GENERATED BY dev/generate-postgresrow-multi-decode.sh - -extension PostgresRow { - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0) { - precondition(self.columns.count >= 1) - let columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - let column = columnIterator.next().unsafelyUnwrapped - let swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0).Type, file: String = #fileID, line: Int = #line) throws -> (T0) { - try self.decode(T0.self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1) { - precondition(self.columns.count >= 2) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1) { - try self.decode((T0, T1).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2) { - precondition(self.columns.count >= 3) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2) { - try self.decode((T0, T1, T2).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3) { - precondition(self.columns.count >= 4) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3) { - try self.decode((T0, T1, T2, T3).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4) { - precondition(self.columns.count >= 5) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4) { - try self.decode((T0, T1, T2, T3, T4).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5) { - precondition(self.columns.count >= 6) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5) { - try self.decode((T0, T1, T2, T3, T4, T5).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6) { - precondition(self.columns.count >= 7) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6) { - try self.decode((T0, T1, T2, T3, T4, T5, T6).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7) { - precondition(self.columns.count >= 8) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8) { - precondition(self.columns.count >= 9) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9) { - precondition(self.columns.count >= 10) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 9 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T9.self - let r9 = try T9._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8, r9) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10) { - precondition(self.columns.count >= 11) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 9 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T9.self - let r9 = try T9._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 10 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T10.self - let r10 = try T10._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11) { - precondition(self.columns.count >= 12) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 9 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T9.self - let r9 = try T9._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 10 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T10.self - let r10 = try T10._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 11 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T11.self - let r11 = try T11._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12) { - precondition(self.columns.count >= 13) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 9 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T9.self - let r9 = try T9._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 10 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T10.self - let r10 = try T10._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 11 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T11.self - let r11 = try T11._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 12 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T12.self - let r12 = try T12._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13) { - precondition(self.columns.count >= 14) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 9 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T9.self - let r9 = try T9._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 10 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T10.self - let r10 = try T10._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 11 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T11.self - let r11 = try T11._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 12 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T12.self - let r12 = try T12._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 13 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T13.self - let r13 = try T13._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12, r13) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14) { - precondition(self.columns.count >= 15) - var columnIndex = 0 - var cellIterator = self.data.makeIterator() - var cellData = cellIterator.next().unsafelyUnwrapped - var columnIterator = self.columns.makeIterator() - var column = columnIterator.next().unsafelyUnwrapped - var swiftTargetType: Any.Type = T0.self - - do { - let r0 = try T0._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 1 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T1.self - let r1 = try T1._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 2 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T2.self - let r2 = try T2._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 3 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T3.self - let r3 = try T3._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 4 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T4.self - let r4 = try T4._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 5 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T5.self - let r5 = try T5._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 6 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T6.self - let r6 = try T6._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 7 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T7.self - let r7 = try T7._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 8 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T8.self - let r8 = try T8._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 9 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T9.self - let r9 = try T9._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 10 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T10.self - let r10 = try T10._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 11 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T11.self - let r11 = try T11._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 12 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T12.self - let r12 = try T12._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 13 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T13.self - let r13 = try T13._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - columnIndex = 14 - cellData = cellIterator.next().unsafelyUnwrapped - column = columnIterator.next().unsafelyUnwrapped - swiftTargetType = T14.self - let r14 = try T14._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) - - return (r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12, r13, r14) - } catch let code as PostgresDecodingError.Code { - throw PostgresDecodingError( - code: code, - columnName: column.name, - columnIndex: columnIndex, - targetType: swiftTargetType, - postgresType: column.dataType, - postgresFormat: column.format, - postgresData: cellData, - file: file, - line: line - ) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).Type, file: String = #fileID, line: Int = #line) throws -> (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14) { - try self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).self, context: .default, file: file, line: line) - } -} diff --git a/Sources/PostgresNIO/New/PostgresRowSequence-multi-decode.swift b/Sources/PostgresNIO/New/PostgresRowSequence-multi-decode.swift deleted file mode 100644 index 53d9a7ea..00000000 --- a/Sources/PostgresNIO/New/PostgresRowSequence-multi-decode.swift +++ /dev/null @@ -1,215 +0,0 @@ -/// NOTE: THIS FILE IS AUTO-GENERATED BY dev/generate-postgresrowsequence-multi-decode.sh - -#if canImport(_Concurrency) -extension AsyncSequence where Element == PostgresRow { - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode(T0.self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode(T0.self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13).self, context: .default, file: file, line: line) - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).Type, context: PostgresDecodingContext, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.map { row in - try row.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).self, context: context, file: file, line: line) - } - } - - @inlinable - @_alwaysEmitIntoClient - public func decode(_: (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).Type, file: String = #fileID, line: Int = #line) -> AsyncThrowingMapSequence { - self.decode((T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14).self, context: .default, file: file, line: line) - } -} -#endif diff --git a/Sources/PostgresNIO/New/PostgresRowSequence.swift b/Sources/PostgresNIO/New/PostgresRowSequence.swift index ccf4f69c..3936b51e 100644 --- a/Sources/PostgresNIO/New/PostgresRowSequence.swift +++ b/Sources/PostgresNIO/New/PostgresRowSequence.swift @@ -4,7 +4,7 @@ import NIOConcurrencyHelpers /// An async sequence of ``PostgresRow``s. /// /// - Note: This is a struct to allow us to move to a move only type easily once they become available. -public struct PostgresRowSequence: AsyncSequence { +public struct PostgresRowSequence: AsyncSequence, Sendable { public typealias Element = PostgresRow typealias BackingSequence = NIOThrowingAsyncSequenceProducer @@ -56,6 +56,9 @@ extension PostgresRowSequence { } } +@available(*, unavailable) +extension PostgresRowSequence.AsyncIterator: Sendable {} + extension PostgresRowSequence { public func collect() async throws -> [PostgresRow] { var result = [PostgresRow]() diff --git a/Sources/PostgresNIO/New/PostgresTransactionError.swift b/Sources/PostgresNIO/New/PostgresTransactionError.swift new file mode 100644 index 00000000..35038446 --- /dev/null +++ b/Sources/PostgresNIO/New/PostgresTransactionError.swift @@ -0,0 +1,21 @@ +/// A wrapper around the errors that can occur during a transaction. +public struct PostgresTransactionError: Error { + + /// The file in which the transaction was started + public var file: String + /// The line in which the transaction was started + public var line: Int + + /// The error thrown when running the `BEGIN` query + public var beginError: Error? + /// The error thrown in the transaction closure + public var closureError: Error? + + /// The error thrown while rolling the transaction back. If the ``closureError`` is set, + /// but the ``rollbackError`` is empty, the rollback was successful. If the ``rollbackError`` + /// is set, the rollback failed. + public var rollbackError: Error? + + /// The error thrown while commiting the transaction. + public var commitError: Error? +} diff --git a/Sources/PostgresNIO/New/PreparedStatement.swift b/Sources/PostgresNIO/New/PreparedStatement.swift new file mode 100644 index 00000000..21165388 --- /dev/null +++ b/Sources/PostgresNIO/New/PreparedStatement.swift @@ -0,0 +1,61 @@ +/// A prepared statement. +/// +/// Structs conforming to this protocol will need to provide the SQL statement to +/// send to the server and a way of creating bindings are decoding the result. +/// +/// As an example, consider this struct: +/// ```swift +/// struct Example: PostgresPreparedStatement { +/// static let sql = "SELECT pid, datname FROM pg_stat_activity WHERE state = $1" +/// typealias Row = (Int, String) +/// +/// var state: String +/// +/// func makeBindings() -> PostgresBindings { +/// var bindings = PostgresBindings() +/// bindings.append(self.state) +/// return bindings +/// } +/// +/// func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { +/// try row.decode(Row.self) +/// } +/// } +/// ``` +/// +/// Structs conforming to this protocol can then be used with `PostgresConnection.execute(_ preparedStatement:, logger:)`, +/// which will take care of preparing the statement on the server side and executing it. +public protocol PostgresPreparedStatement: Sendable { + /// The prepared statements name. + /// + /// > Note: There is a default implementation that returns the implementor's name. + static var name: String { get } + + /// The type rows returned by the statement will be decoded into + associatedtype Row + + /// The SQL statement to prepare on the database server. + static var sql: String { get } + + /// The postgres data types of the values that are bind when this statement is executed. + /// + /// If an empty array is returned the datatypes are inferred from the ``PostgresBindings`` returned + /// from ``PostgresPreparedStatement/makeBindings()``. + /// + /// > Note: There is a default implementation that returns an empty array, which will lead to + /// automatic inference. + static var bindingDataTypes: [PostgresDataType] { get } + + /// Make the bindings to provided concrete values to use when executing the prepared SQL statement. + /// The order must match ``PostgresPreparedStatement/bindingDataTypes-4b6tx``. + func makeBindings() throws -> PostgresBindings + + /// Decode a row returned by the database into an instance of `Row` + func decodeRow(_ row: PostgresRow) throws -> Row +} + +extension PostgresPreparedStatement { + public static var name: String { String(reflecting: self) } + + public static var bindingDataTypes: [PostgresDataType] { [] } +} diff --git a/Sources/PostgresNIO/New/VariadicGenerics.swift b/Sources/PostgresNIO/New/VariadicGenerics.swift new file mode 100644 index 00000000..7931c90c --- /dev/null +++ b/Sources/PostgresNIO/New/VariadicGenerics.swift @@ -0,0 +1,172 @@ + +extension PostgresRow { + // --- snip TODO: Remove once bug is fixed, that disallows tuples of one + @inlinable + public func decode( + _: Column.Type, + file: String = #fileID, + line: Int = #line + ) throws -> (Column) { + try self.decode(Column.self, context: .default, file: file, line: line) + } + + @inlinable + public func decode( + _: Column.Type, + context: PostgresDecodingContext, + file: String = #fileID, + line: Int = #line + ) throws -> (Column) { + precondition(self.columns.count >= 1) + let columnIndex = 0 + var cellIterator = self.data.makeIterator() + var cellData = cellIterator.next().unsafelyUnwrapped + var columnIterator = self.columns.makeIterator() + let column = columnIterator.next().unsafelyUnwrapped + let swiftTargetType: Any.Type = Column.self + + do { + let r0 = try Column._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) + + return (r0) + } catch let code as PostgresDecodingError.Code { + throw PostgresDecodingError( + code: code, + columnName: column.name, + columnIndex: columnIndex, + targetType: swiftTargetType, + postgresType: column.dataType, + postgresFormat: column.format, + postgresData: cellData, + file: file, + line: line + ) + } + } + // --- snap TODO: Remove once bug is fixed, that disallows tuples of one + + @inlinable + public func decode( + _ columnType: (repeat each Column).Type, + context: PostgresDecodingContext, + file: String = #fileID, + line: Int = #line + ) throws -> (repeat each Column) { + let packCount = ComputeParameterPackLength.count(ofPack: repeat (each Column).self) + precondition(self.columns.count >= packCount) + + var columnIndex = 0 + var cellIterator = self.data.makeIterator() + var columnIterator = self.columns.makeIterator() + + return ( + repeat try Self.decodeNextColumn( + (each Column).self, + cellIterator: &cellIterator, + columnIterator: &columnIterator, + columnIndex: &columnIndex, + context: context, + file: file, + line: line + ) + ) + } + + @inlinable + static func decodeNextColumn( + _ columnType: Column.Type, + cellIterator: inout IndexingIterator, + columnIterator: inout IndexingIterator<[RowDescription.Column]>, + columnIndex: inout Int, + context: PostgresDecodingContext, + file: String, + line: Int + ) throws -> Column { + defer { columnIndex += 1 } + + let column = columnIterator.next().unsafelyUnwrapped + var cellData = cellIterator.next().unsafelyUnwrapped + do { + return try Column._decodeRaw(from: &cellData, type: column.dataType, format: column.format, context: context) + } catch let code as PostgresDecodingError.Code { + throw PostgresDecodingError( + code: code, + columnName: column.name, + columnIndex: columnIndex, + targetType: Column.self, + postgresType: column.dataType, + postgresFormat: column.format, + postgresData: cellData, + file: file, + line: line + ) + } + } + + @inlinable + public func decode( + _ columnType: (repeat each Column).Type, + file: String = #fileID, + line: Int = #line + ) throws -> (repeat each Column) { + try self.decode(columnType, context: .default, file: file, line: line) + } +} + +extension AsyncSequence where Element == PostgresRow { + // --- snip TODO: Remove once bug is fixed, that disallows tuples of one + @inlinable + public func decode( + _: Column.Type, + context: PostgresDecodingContext, + file: String = #fileID, + line: Int = #line + ) -> AsyncThrowingMapSequence { + self.map { row in + try row.decode(Column.self, context: context, file: file, line: line) + } + } + + @inlinable + public func decode( + _: Column.Type, + file: String = #fileID, + line: Int = #line + ) -> AsyncThrowingMapSequence { + self.decode(Column.self, context: .default, file: file, line: line) + } + // --- snap TODO: Remove once bug is fixed, that disallows tuples of one + + public func decode( + _ columnType: (repeat each Column).Type, + context: PostgresDecodingContext, + file: String = #fileID, + line: Int = #line + ) -> AsyncThrowingMapSequence { + self.map { row in + try row.decode(columnType, context: context, file: file, line: line) + } + } + + public func decode( + _ columnType: (repeat each Column).Type, + file: String = #fileID, + line: Int = #line + ) -> AsyncThrowingMapSequence { + self.decode(columnType, context: .default, file: file, line: line) + } +} + +@usableFromInline +enum ComputeParameterPackLength { + @usableFromInline + enum BoolConverter { + @usableFromInline + typealias Bool = Swift.Bool + } + + @inlinable + static func count(ofPack t: repeat each T) -> Int { + MemoryLayout<(repeat BoolConverter.Bool)>.size / MemoryLayout.stride + } +} diff --git a/Sources/PostgresNIO/Pool/ConnectionFactory.swift b/Sources/PostgresNIO/Pool/ConnectionFactory.swift new file mode 100644 index 00000000..319b86c4 --- /dev/null +++ b/Sources/PostgresNIO/Pool/ConnectionFactory.swift @@ -0,0 +1,207 @@ +import Logging +import NIOConcurrencyHelpers +import NIOCore +import NIOSSL + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class ConnectionFactory: Sendable { + + struct ConfigCache: Sendable { + var config: PostgresClient.Configuration + } + + let configBox: NIOLockedValueBox + + struct SSLContextCache: Sendable { + enum State { + case none + case producing(TLSConfiguration, [CheckedContinuation]) + case cached(TLSConfiguration, NIOSSLContext) + case failed(TLSConfiguration, any Error) + } + + var state: State = .none + } + + let sslContextBox = NIOLockedValueBox(SSLContextCache()) + + let eventLoopGroup: any EventLoopGroup + + let logger: Logger + + init(config: PostgresClient.Configuration, eventLoopGroup: any EventLoopGroup, logger: Logger) { + self.eventLoopGroup = eventLoopGroup + self.configBox = NIOLockedValueBox(ConfigCache(config: config)) + self.logger = logger + } + + func makeConnection(_ connectionID: PostgresConnection.ID, pool: PostgresClient.Pool) async throws -> PostgresConnection { + let config = try await self.makeConnectionConfig() + + var connectionLogger = self.logger + connectionLogger[postgresMetadataKey: .connectionID] = "\(connectionID)" + + return try await PostgresConnection.connect( + on: self.eventLoopGroup.any(), + configuration: config, + id: connectionID, + logger: connectionLogger + ).get() + } + + func makeConnectionConfig() async throws -> PostgresConnection.Configuration { + let config = self.configBox.withLockedValue { $0.config } + + let tls: PostgresConnection.Configuration.TLS + switch config.tls.base { + case .prefer(let tlsConfiguration): + let sslContext = try await self.getSSLContext(for: tlsConfiguration) + tls = .prefer(sslContext) + + case .require(let tlsConfiguration): + let sslContext = try await self.getSSLContext(for: tlsConfiguration) + tls = .require(sslContext) + case .disable: + tls = .disable + } + + var connectionConfig: PostgresConnection.Configuration + switch config.endpointInfo { + case .bindUnixDomainSocket(let path): + connectionConfig = PostgresConnection.Configuration( + unixSocketPath: path, + username: config.username, + password: config.password, + database: config.database + ) + + case .connectTCP(let host, let port): + connectionConfig = PostgresConnection.Configuration( + host: host, + port: port, + username: config.username, + password: config.password, + database: config.database, + tls: tls + ) + } + + connectionConfig.options.connectTimeout = TimeAmount(config.options.connectTimeout) + connectionConfig.options.tlsServerName = config.options.tlsServerName + connectionConfig.options.requireBackendKeyData = config.options.requireBackendKeyData + connectionConfig.options.additionalStartupParameters = config.options.additionalStartupParameters + + return connectionConfig + } + + private func getSSLContext(for tlsConfiguration: TLSConfiguration) async throws -> NIOSSLContext { + enum Action { + case produce + case succeed(NIOSSLContext) + case fail(any Error) + case wait + } + + return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + let action = self.sslContextBox.withLockedValue { cache -> Action in + switch cache.state { + case .none: + cache.state = .producing(tlsConfiguration, [continuation]) + return .produce + + case .cached(let cachedTLSConfiguration, let context): + if cachedTLSConfiguration.bestEffortEquals(tlsConfiguration) { + return .succeed(context) + } else { + cache.state = .producing(tlsConfiguration, [continuation]) + return .produce + } + + case .failed(let cachedTLSConfiguration, let error): + if cachedTLSConfiguration.bestEffortEquals(tlsConfiguration) { + return .fail(error) + } else { + cache.state = .producing(tlsConfiguration, [continuation]) + return .produce + } + + case .producing(let cachedTLSConfiguration, var continuations): + continuations.append(continuation) + if cachedTLSConfiguration.bestEffortEquals(tlsConfiguration) { + cache.state = .producing(cachedTLSConfiguration, continuations) + return .wait + } else { + cache.state = .producing(tlsConfiguration, continuations) + return .produce + } + } + } + + switch action { + case .wait: + break + + case .produce: + // TBD: we might want to consider moving this off the concurrent executor + self.reportProduceSSLContextResult( + Result(catching: {try NIOSSLContext(configuration: tlsConfiguration)}), + for: tlsConfiguration + ) + + case .succeed(let context): + continuation.resume(returning: context) + + case .fail(let error): + continuation.resume(throwing: error) + } + } + } + + private func reportProduceSSLContextResult(_ result: Result, for tlsConfiguration: TLSConfiguration) { + enum Action { + case fail(any Error, [CheckedContinuation]) + case succeed(NIOSSLContext, [CheckedContinuation]) + case none + } + + let action = self.sslContextBox.withLockedValue { cache -> Action in + switch cache.state { + case .none: + preconditionFailure("Invalid state: \(cache.state)") + + case .cached, .failed: + return .none + + case .producing(let cachedTLSConfiguration, let continuations): + if cachedTLSConfiguration.bestEffortEquals(tlsConfiguration) { + switch result { + case .success(let context): + cache.state = .cached(cachedTLSConfiguration, context) + return .succeed(context, continuations) + + case .failure(let failure): + cache.state = .failed(cachedTLSConfiguration, failure) + return .fail(failure, continuations) + } + } else { + return .none + } + } + } + + switch action { + case .none: + break + + case .succeed(let context, let continuations): + for continuation in continuations { + continuation.resume(returning: context) + } + + case .fail(let error, let continuations): + for continuation in continuations { + continuation.resume(throwing: error) + } + } + } +} diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift new file mode 100644 index 00000000..d54e34eb --- /dev/null +++ b/Sources/PostgresNIO/Pool/PostgresClient.swift @@ -0,0 +1,581 @@ +import NIOCore +import NIOSSL +import Atomics +import Logging +import ServiceLifecycle +import _ConnectionPoolModule + +/// A Postgres client that is backed by an underlying connection pool. Use ``Configuration`` to change the client's +/// behavior. +/// +/// ## Creating a client +/// +/// You create a ``PostgresClient`` by first creating a ``PostgresClient/Configuration`` struct that you can +/// use to modify the client's behavior. +/// +/// @Snippet(path: "postgres-nio/Snippets/PostgresClient", slice: "configuration") +/// +/// Now you can create a client with your configuration object: +/// +/// @Snippet(path: "postgres-nio/Snippets/PostgresClient", slice: "makeClient") +/// +/// ## Running a client +/// +/// ``PostgresClient`` relies on structured concurrency. Because of this it needs a task in which it can schedule all the +/// background work that it needs to do in order to manage connections on the users behave. For this reason, developers +/// must provide a task to the client by scheduling the client's run method in a long running task: +/// +/// @Snippet(path: "postgres-nio/Snippets/PostgresClient", slice: "run") +/// +/// ``PostgresClient`` can not lease connections, if its ``run()`` method isn't active. Cancelling the ``run()`` method +/// is equivalent to closing the client. Once a client's ``run()`` method has been cancelled, executing queries or prepared +/// statements will fail. +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +public final class PostgresClient: Sendable, ServiceLifecycle.Service { + public struct Configuration: Sendable { + public struct TLS: Sendable { + enum Base { + case disable + case prefer(NIOSSL.TLSConfiguration) + case require(NIOSSL.TLSConfiguration) + } + + var base: Base + + private init(_ base: Base) { + self.base = base + } + + /// Do not try to create a TLS connection to the server. + public static let disable: Self = Self.init(.disable) + + /// Try to create a TLS connection to the server. If the server supports TLS, create a TLS connection. + /// If the server does not support TLS, create an insecure connection. + public static func prefer(_ sslContext: NIOSSL.TLSConfiguration) -> Self { + self.init(.prefer(sslContext)) + } + + /// Try to create a TLS connection to the server. If the server supports TLS, create a TLS connection. + /// If the server does not support TLS, fail the connection creation. + public static func require(_ sslContext: NIOSSL.TLSConfiguration) -> Self { + self.init(.require(sslContext)) + } + } + + // MARK: Client options + + /// Describes general client behavior options. Those settings are considered advanced options. + public struct Options: Sendable { + /// A keep-alive behavior for Postgres connections. The ``frequency`` defines after which time an idle + /// connection shall run a keep-alive ``query``. + public struct KeepAliveBehavior: Sendable { + /// The amount of time that shall pass before an idle connection runs a keep-alive ``query``. + public var frequency: Duration + + /// The ``query`` that is run on an idle connection after it has been idle for ``frequency``. + public var query: PostgresQuery + + /// Create a new `KeepAliveBehavior`. + /// - Parameters: + /// - frequency: The amount of time that shall pass before an idle connection runs a keep-alive `query`. + /// Defaults to `30` seconds. + /// - query: The `query` that is run on an idle connection after it has been idle for `frequency`. + /// Defaults to `SELECT 1;`. + public init(frequency: Duration = .seconds(30), query: PostgresQuery = "SELECT 1;") { + self.frequency = frequency + self.query = query + } + } + + /// A timeout for creating a TCP/Unix domain socket connection. Defaults to `10` seconds. + public var connectTimeout: Duration = .seconds(10) + + /// The server name to use for certificate validation and SNI (Server Name Indication) when TLS is enabled. + /// Defaults to none (but see below). + /// + /// > When set to `nil`: + /// If the connection is made to a server over TCP using + /// ``PostgresConnection/Configuration/init(host:port:username:password:database:tls:)``, the given `host` + /// is used, unless it was an IP address string. If it _was_ an IP, or the connection is made by any other + /// method, SNI is disabled. + public var tlsServerName: String? = nil + + /// Whether the connection is required to provide backend key data (internal Postgres stuff). + /// + /// This property is provided for compatibility with Amazon RDS Proxy, which requires it to be `false`. + /// If you are not using Amazon RDS Proxy, you should leave this set to `true` (the default). + public var requireBackendKeyData: Bool = true + + /// Additional parameters to send to the server on startup. The name value pairs are added to the initial + /// startup message that the client sends to the server. + public var additionalStartupParameters: [(String, String)] = [] + + /// The minimum number of connections that the client shall keep open at any time, even if there is no + /// demand. Default to `0`. + /// + /// If the open connection count becomes less than ``minimumConnections`` new connections + /// are created immidiatly. Must be greater or equal to zero and less than ``maximumConnections``. + /// + /// Idle connections are kept alive using the ``keepAliveBehavior``. + public var minimumConnections: Int = 0 + + /// The maximum number of connections that the client may open to the server at any time. Must be greater + /// than ``minimumConnections``. Defaults to `20` connections. + /// + /// Connections, that are created in response to demand are kept alive for the ``connectionIdleTimeout`` + /// before they are dropped. + public var maximumConnections: Int = 20 + + /// The maximum amount time that a connection that is not part of the ``minimumConnections`` is kept + /// open without being leased. Defaults to `60` seconds. + public var connectionIdleTimeout: Duration = .seconds(60) + + /// The ``KeepAliveBehavior-swift.struct`` to ensure that the underlying tcp-connection is still active + /// for idle connections. `Nil` means that the client shall not run keep alive queries to the server. Defaults to a + /// keep alive query of `SELECT 1;` every `30` seconds. + public var keepAliveBehavior: KeepAliveBehavior? = KeepAliveBehavior() + + /// Create an options structure with default values. + /// + /// Most users should not need to adjust the defaults. + public init() {} + } + + // MARK: - Accessors + + /// The hostname to connect to for TCP configurations. + /// + /// Always `nil` for other configurations. + public var host: String? { + if case let .connectTCP(host, _) = self.endpointInfo { return host } + else { return nil } + } + + /// The port to connect to for TCP configurations. + /// + /// Always `nil` for other configurations. + public var port: Int? { + if case let .connectTCP(_, port) = self.endpointInfo { return port } + else { return nil } + } + + /// The socket path to connect to for Unix domain socket connections. + /// + /// Always `nil` for other configurations. + public var unixSocketPath: String? { + if case let .bindUnixDomainSocket(path) = self.endpointInfo { return path } + else { return nil } + } + + /// The TLS mode to use for the connection. Valid for all configurations. + /// + /// See ``TLS-swift.struct``. + public var tls: TLS = .prefer(.makeClientConfiguration()) + + /// Options for handling the communication channel. Most users don't need to change these. + /// + /// See ``Options-swift.struct``. + public var options: Options = .init() + + /// The username to connect with. + public var username: String + + /// The password, if any, for the user specified by ``username``. + /// + /// - Warning: `nil` means "no password provided", whereas `""` (the empty string) is a password of zero + /// length; these are not the same thing. + public var password: String? + + /// The name of the database to open. + /// + /// - Note: If set to `nil` or an empty string, the provided ``username`` is used. + public var database: String? + + // MARK: - Initializers + + /// Create a configuration for connecting to a server with a hostname and optional port. + /// + /// This specifies a TCP connection. If you're unsure which kind of connection you want, you almost + /// definitely want this one. + /// + /// - Parameters: + /// - host: The hostname to connect to. + /// - port: The TCP port to connect to (defaults to 5432). + /// - tls: The TLS mode to use. + public init(host: String, port: Int = 5432, username: String, password: String?, database: String?, tls: TLS) { + self.init(endpointInfo: .connectTCP(host: host, port: port), tls: tls, username: username, password: password, database: database) + } + + /// Create a configuration for connecting to a server through a UNIX domain socket. + /// + /// - Parameters: + /// - path: The filesystem path of the socket to connect to. + /// - tls: The TLS mode to use. Defaults to ``TLS-swift.struct/disable``. + public init(unixSocketPath: String, username: String, password: String?, database: String?) { + self.init(endpointInfo: .bindUnixDomainSocket(path: unixSocketPath), tls: .disable, username: username, password: password, database: database) + } + + // MARK: - Implementation details + + enum EndpointInfo { + case bindUnixDomainSocket(path: String) + case connectTCP(host: String, port: Int) + } + + var endpointInfo: EndpointInfo + + init(endpointInfo: EndpointInfo, tls: TLS, username: String, password: String?, database: String?) { + self.endpointInfo = endpointInfo + self.tls = tls + self.username = username + self.password = password + self.database = database + } + } + + typealias Pool = ConnectionPool< + PostgresConnection, + PostgresConnection.ID, + ConnectionIDGenerator, + ConnectionRequest, + ConnectionRequest.ID, + PostgresKeepAliveBehavor, + PostgresClientMetrics, + ContinuousClock + > + + let pool: Pool + let factory: ConnectionFactory + let runningAtomic = ManagedAtomic(false) + let backgroundLogger: Logger + + /// Creates a new ``PostgresClient``, that does not log any background information. + /// + /// > Warning: + /// The client can only lease connections if the user is running the client's ``run()`` method in a long running task. + /// + /// - Parameters: + /// - configuration: The client's configuration. See ``Configuration`` for details. + /// - eventLoopGroup: The underlying NIO `EventLoopGroup`. Defaults to ``defaultEventLoopGroup``. + public convenience init( + configuration: Configuration, + eventLoopGroup: any EventLoopGroup = PostgresClient.defaultEventLoopGroup + ) { + self.init(configuration: configuration, eventLoopGroup: eventLoopGroup, backgroundLogger: Self.loggingDisabled) + } + + /// Creates a new ``PostgresClient``. Don't forget to run ``run()`` the client in a long running task. + /// + /// - Parameters: + /// - configuration: The client's configuration. See ``Configuration`` for details. + /// - eventLoopGroup: The underlying NIO `EventLoopGroup`. Defaults to ``defaultEventLoopGroup``. + /// - backgroundLogger: A `swift-log` `Logger` to log background messages to. A copy of this logger is also + /// forwarded to the created connections as a background logger. + public init( + configuration: Configuration, + eventLoopGroup: any EventLoopGroup = PostgresClient.defaultEventLoopGroup, + backgroundLogger: Logger + ) { + let factory = ConnectionFactory(config: configuration, eventLoopGroup: eventLoopGroup, logger: backgroundLogger) + self.factory = factory + self.backgroundLogger = backgroundLogger + + self.pool = ConnectionPool( + configuration: .init(configuration), + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: .init(configuration.options.keepAliveBehavior, logger: backgroundLogger), + observabilityDelegate: .init(logger: backgroundLogger), + clock: ContinuousClock() + ) { (connectionID, pool) in + let connection = try await factory.makeConnection(connectionID, pool: pool) + + return ConnectionAndMetadata(connection: connection, maximalStreamsOnConnection: 1) + } + } + + /// Lease a connection for the provided `closure`'s lifetime. + /// + /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture + /// the provided `PostgresConnection`. + /// - Returns: The closure's return value. + @_disfavoredOverload + public func withConnection(_ closure: (PostgresConnection) async throws -> Result) async throws -> Result { + let connection = try await self.leaseConnection() + + defer { self.pool.releaseConnection(connection) } + + return try await closure(connection) + } + + #if compiler(>=6.0) + /// Lease a connection for the provided `closure`'s lifetime. + /// + /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture + /// the provided `PostgresConnection`. + /// - Returns: The closure's return value. + public func withConnection( + isolation: isolated (any Actor)? = #isolation, + // DO NOT FIX THE WHITESPACE IN THE NEXT LINE UNTIL 5.10 IS UNSUPPORTED + // https://github.com/swiftlang/swift/issues/79285 + _ closure: (PostgresConnection) async throws -> sending Result) async throws -> sending Result { + let connection = try await self.leaseConnection() + + defer { self.pool.releaseConnection(connection) } + + return try await closure(connection) + } + + /// Lease a connection, which is in an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function leases a connection from the underlying connection pool and starts a transaction by running a `BEGIN` + /// query on the leased connection against the database. It then lends the connection to the user provided closure. + /// The user can then modify the database as they wish. If the user provided closure returns successfully, the function + /// will attempt to commit the changes by running a `COMMIT` query against the database. If the user provided closure + /// throws an error, the function will attempt to rollback the changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + isolation: isolated (any Actor)? = #isolation, + // DO NOT FIX THE WHITESPACE IN THE NEXT LINE UNTIL 5.10 IS UNSUPPORTED + // https://github.com/swiftlang/swift/issues/79285 + _ closure: (PostgresConnection) async throws -> sending Result) async throws -> sending Result { + try await self.withConnection { connection in + try await connection.withTransaction(logger: logger, file: file, line: line, closure) + } + } + #else + + /// Lease a connection, which is in an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function leases a connection from the underlying connection pool and starts a transaction by running a `BEGIN` + /// query on the leased connection against the database. It then lends the connection to the user provided closure. + /// The user can then modify the database as they wish. If the user provided closure returns successfully, the function + /// will attempt to commit the changes by running a `COMMIT` query against the database. If the user provided closure + /// throws an error, the function will attempt to rollback the changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + _ closure: (PostgresConnection) async throws -> Result + ) async throws -> Result { + try await self.withConnection { connection in + try await connection.withTransaction(logger: logger, file: file, line: line, closure) + } + } + #endif + + /// Run a query on the Postgres server the client is connected to. + /// + /// - Parameters: + /// - query: The ``PostgresQuery`` to run + /// - logger: The `Logger` to log into for the query + /// - file: The file, the query was started in. Used for better error reporting. + /// - line: The line, the query was started in. Used for better error reporting. + /// - Returns: A ``PostgresRowSequence`` containing the rows the server sent as the query result. + /// The sequence be discarded. + @discardableResult + public func query( + _ query: PostgresQuery, + logger: Logger? = nil, + file: String = #fileID, + line: Int = #line + ) async throws -> PostgresRowSequence { + let logger = logger ?? Self.loggingDisabled + do { + guard query.binds.count <= Int(UInt16.max) else { + throw PSQLError(code: .tooManyParameters, query: query, file: file, line: line) + } + + let connection = try await self.leaseConnection() + + var logger = logger + logger[postgresMetadataKey: .connectionID] = "\(connection.id)" + + let promise = connection.channel.eventLoop.makePromise(of: PSQLRowStream.self) + let context = ExtendedQueryContext( + query: query, + logger: logger, + promise: promise + ) + + connection.channel.write(HandlerTask.extendedQuery(context), promise: nil) + + promise.futureResult.whenFailure { _ in + self.pool.releaseConnection(connection) + } + + return try await promise.futureResult.map { + $0.asyncSequence(onFinish: { + self.pool.releaseConnection(connection) + }) + }.get() + } catch var error as PSQLError { + error.file = file + error.line = line + error.query = query + throw error // rethrow with more metadata + } + } + + /// Execute a prepared statement, taking care of the preparation when necessary + public func execute( + _ preparedStatement: Statement, + logger: Logger? = nil, + file: String = #fileID, + line: Int = #line + ) async throws -> AsyncThrowingMapSequence where Row == Statement.Row { + let bindings = try preparedStatement.makeBindings() + let logger = logger ?? Self.loggingDisabled + + do { + let connection = try await self.leaseConnection() + + let promise = connection.channel.eventLoop.makePromise(of: PSQLRowStream.self) + let task = HandlerTask.executePreparedStatement(.init( + name: String(reflecting: Statement.self), + sql: Statement.sql, + bindings: bindings, + bindingDataTypes: Statement.bindingDataTypes, + logger: logger, + promise: promise + )) + connection.channel.write(task, promise: nil) + + promise.futureResult.whenFailure { _ in + self.pool.releaseConnection(connection) + } + + return try await promise.futureResult + .map { $0.asyncSequence(onFinish: { self.pool.releaseConnection(connection) }) } + .get() + .map { try preparedStatement.decodeRow($0) } + } catch var error as PSQLError { + error.file = file + error.line = line + error.query = .init( + unsafeSQL: Statement.sql, + binds: bindings + ) + throw error // rethrow with more metadata + } + } + + /// The structured root task for the client's background work. + /// + /// > Warning: + /// Users must call this function in order to allow the client to process any background work. Executing queries, + /// prepared statements or leasing connections will hang until the developer executes the client's ``run()`` + /// method. + /// + /// Cancelling the task which executes the ``run()`` method, is equivalent to closing the client. Once the task + /// has been cancelled the client is not able to process any new queries or prepared statements. + /// + /// @Snippet(path: "postgres-nio/Snippets/PostgresClient", slice: "run") + /// + /// > Note: + /// ``PostgresClient`` implements [ServiceLifecycle](https://github.com/swift-server/swift-service-lifecycle)'s `Service` protocol. Because of this + /// ``PostgresClient`` can be passed to a `ServiceGroup` for easier lifecycle management. + public func run() async { + let atomicOp = self.runningAtomic.compareExchange(expected: false, desired: true, ordering: .relaxed) + precondition(!atomicOp.original, "PostgresClient.run() should just be called once!") + + await cancelWhenGracefulShutdown { + await self.pool.run() + } + } + + // MARK: - Private Methods - + + private func leaseConnection() async throws -> PostgresConnection { + if !self.runningAtomic.load(ordering: .relaxed) { + self.backgroundLogger.warning("Trying to lease connection from `PostgresClient`, but `PostgresClient.run()` hasn't been called yet.") + } + return try await self.pool.leaseConnection() + } + + /// Returns the default `EventLoopGroup` singleton, automatically selecting the best for the platform. + /// + /// This will select the concrete `EventLoopGroup` depending which platform this is running on. + public static var defaultEventLoopGroup: EventLoopGroup { + PostgresConnection.defaultEventLoopGroup + } + + static let loggingDisabled = Logger(label: "Postgres-do-not-log", factory: { _ in SwiftLogNoOpLogHandler() }) +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +struct PostgresKeepAliveBehavor: ConnectionKeepAliveBehavior { + let behavior: PostgresClient.Configuration.Options.KeepAliveBehavior? + let logger: Logger + + init(_ behavior: PostgresClient.Configuration.Options.KeepAliveBehavior?, logger: Logger) { + self.behavior = behavior + self.logger = logger + } + + var keepAliveFrequency: Duration? { + self.behavior?.frequency + } + + func runKeepAlive(for connection: PostgresConnection) async throws { + try await connection.query(self.behavior!.query, logger: self.logger).map { _ in }.get() + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension ConnectionPoolConfiguration { + init(_ config: PostgresClient.Configuration) { + self = ConnectionPoolConfiguration() + self.minimumConnectionCount = config.options.minimumConnections + self.maximumConnectionSoftLimit = config.options.maximumConnections + self.maximumConnectionHardLimit = config.options.maximumConnections + self.idleTimeout = config.options.connectionIdleTimeout + } +} + +extension PostgresConnection: PooledConnection { + public func close() { + self.channel.close(mode: .all, promise: nil) + } + + public func onClose(_ closure: @escaping @Sendable ((any Error)?) -> ()) { + self.closeFuture.whenComplete { _ in closure(nil) } + } +} + +extension ConnectionPoolError { + func mapToPSQLError(lastConnectError: Error?) -> Error { + var psqlError: PSQLError + switch self { + case .poolShutdown: + psqlError = PSQLError.poolClosed + psqlError.underlying = self + + case .requestCancelled: + psqlError = PSQLError.queryCancelled + psqlError.underlying = self + + default: + return self + } + return psqlError + } +} diff --git a/Sources/PostgresNIO/Pool/PostgresClientMetrics.swift b/Sources/PostgresNIO/Pool/PostgresClientMetrics.swift new file mode 100644 index 00000000..aa8215db --- /dev/null +++ b/Sources/PostgresNIO/Pool/PostgresClientMetrics.swift @@ -0,0 +1,85 @@ +import _ConnectionPoolModule +import Logging + +final class PostgresClientMetrics: ConnectionPoolObservabilityDelegate { + typealias ConnectionID = PostgresConnection.ID + + let logger: Logger + + init(logger: Logger) { + self.logger = logger + } + + func startedConnecting(id: ConnectionID) { + self.logger.debug("Creating new connection", metadata: [ + .connectionID: "\(id)", + ]) + } + + /// A connection attempt failed with the given error. After some period of + /// time ``startedConnecting(id:)`` may be called again. + func connectFailed(id: ConnectionID, error: Error) { + self.logger.debug("Connection creation failed", metadata: [ + .connectionID: "\(id)", + .error: "\(String(reflecting: error))" + ]) + } + + func connectSucceeded(id: ConnectionID) { + self.logger.debug("Connection established", metadata: [ + .connectionID: "\(id)" + ]) + } + + /// The utlization of the connection changed; a stream may have been used, returned or the + /// maximum number of concurrent streams available on the connection changed. + func connectionLeased(id: ConnectionID) { + self.logger.debug("Connection leased", metadata: [ + .connectionID: "\(id)" + ]) + } + + func connectionReleased(id: ConnectionID) { + self.logger.debug("Connection released", metadata: [ + .connectionID: "\(id)" + ]) + } + + func keepAliveTriggered(id: ConnectionID) { + self.logger.debug("run ping pong", metadata: [ + .connectionID: "\(id)", + ]) + } + + func keepAliveSucceeded(id: ConnectionID) {} + + func keepAliveFailed(id: PostgresConnection.ID, error: Error) {} + + /// The remote peer is quiescing the connection: no new streams will be created on it. The + /// connection will eventually be closed and removed from the pool. + func connectionClosing(id: ConnectionID) { + self.logger.debug("Close connection", metadata: [ + .connectionID: "\(id)" + ]) + } + + /// The connection was closed. The connection may be established again in the future (notified + /// via ``startedConnecting(id:)``). + func connectionClosed(id: ConnectionID, error: Error?) { + self.logger.debug("Connection closed", metadata: [ + .connectionID: "\(id)" + ]) + } + + func requestQueueDepthChanged(_ newDepth: Int) { + + } + + func connectSucceeded(id: PostgresConnection.ID, streamCapacity: UInt16) { + + } + + func connectionUtilizationChanged(id: PostgresConnection.ID, streamsUsed: UInt16, streamCapacity: UInt16) { + + } +} diff --git a/Sources/PostgresNIO/Postgres+PSQLCompat.swift b/Sources/PostgresNIO/Postgres+PSQLCompat.swift index ff9773f5..7d464c2b 100644 --- a/Sources/PostgresNIO/Postgres+PSQLCompat.swift +++ b/Sources/PostgresNIO/Postgres+PSQLCompat.swift @@ -5,7 +5,7 @@ extension PSQLError { switch self.code.base { case .queryCancelled: return self - case .server: + case .server, .listenFailed: guard let serverInfo = self.serverInfo else { return self } @@ -37,14 +37,17 @@ extension PSQLError { return self.underlying ?? self case .tooManyParameters, .invalidCommandTag: return self - case .connectionQuiescing: - return PostgresError.connectionClosed - case .connectionClosed: + case .clientClosedConnection, + .serverClosedConnection: return PostgresError.connectionClosed case .connectionError: return self.underlying ?? self + case .unlistenFailed: + return self.underlying ?? self case .uncleanShutdown: return PostgresError.protocol("Unexpected connection close") + case .poolClosed: + return self } } } diff --git a/Sources/PostgresNIO/PostgresDatabase+Query.swift b/Sources/PostgresNIO/PostgresDatabase+Query.swift index 95abb6fc..8de93814 100644 --- a/Sources/PostgresNIO/PostgresDatabase+Query.swift +++ b/Sources/PostgresNIO/PostgresDatabase+Query.swift @@ -1,27 +1,35 @@ import NIOCore import Logging +import NIOConcurrencyHelpers extension PostgresDatabase { public func query( _ string: String, _ binds: [PostgresData] = [] ) -> EventLoopFuture { - var rows: [PostgresRow] = [] - var metadata: PostgresQueryMetadata? - return self.query(string, binds, onMetadata: { - metadata = $0 - }) { - rows.append($0) + let box = NIOLockedValueBox((metadata: PostgresQueryMetadata?.none, rows: [PostgresRow]())) + + return self.query(string, binds, onMetadata: { metadata in + box.withLockedValue { + $0.metadata = metadata + } + }) { row in + box.withLockedValue { + $0.rows.append(row) + } }.map { - .init(metadata: metadata!, rows: rows) + box.withLockedValue { + PostgresQueryResult(metadata: $0.metadata!, rows: $0.rows) + } } } + @preconcurrency public func query( _ string: String, _ binds: [PostgresData] = [], - onMetadata: @escaping (PostgresQueryMetadata) -> () = { _ in }, - onRow: @escaping (PostgresRow) throws -> () + onMetadata: @Sendable @escaping (PostgresQueryMetadata) -> () = { _ in }, + onRow: @Sendable @escaping (PostgresRow) throws -> () ) -> EventLoopFuture { var bindings = PostgresBindings(capacity: binds.count) binds.forEach { bindings.append($0) } @@ -32,7 +40,7 @@ extension PostgresDatabase { } } -public struct PostgresQueryResult { +public struct PostgresQueryResult: Sendable { public let metadata: PostgresQueryMetadata public let rows: [PostgresRow] } @@ -58,17 +66,14 @@ extension PostgresQueryResult: Collection { } } -public struct PostgresQueryMetadata { +public struct PostgresQueryMetadata: Sendable { public let command: String public var oid: Int? public var rows: Int? init?(string: String) { let parts = string.split(separator: " ") - guard parts.count >= 1 else { - return nil - } - switch parts[0] { + switch parts.first { case "INSERT": // INSERT oid rows guard parts.count == 3 else { diff --git a/Sources/PostgresNIO/PostgresDatabase+SimpleQuery.swift b/Sources/PostgresNIO/PostgresDatabase+SimpleQuery.swift index 77f3d034..5cf2d7a4 100644 --- a/Sources/PostgresNIO/PostgresDatabase+SimpleQuery.swift +++ b/Sources/PostgresNIO/PostgresDatabase+SimpleQuery.swift @@ -1,13 +1,19 @@ import NIOCore +import NIOConcurrencyHelpers import Logging extension PostgresDatabase { public func simpleQuery(_ string: String) -> EventLoopFuture<[PostgresRow]> { - var rows: [PostgresRow] = [] - return simpleQuery(string) { rows.append($0) }.map { rows } + let rowsBoxed = NIOLockedValueBox([PostgresRow]()) + return self.simpleQuery(string) { row in + rowsBoxed.withLockedValue { + $0.append(row) + } + }.map { rowsBoxed.withLockedValue { $0 } } } - public func simpleQuery(_ string: String, _ onRow: @escaping (PostgresRow) throws -> ()) -> EventLoopFuture { + @preconcurrency + public func simpleQuery(_ string: String, _ onRow: @Sendable @escaping (PostgresRow) throws -> ()) -> EventLoopFuture { self.query(string, onRow: onRow) } } diff --git a/Sources/PostgresNIO/PostgresDatabase.swift b/Sources/PostgresNIO/PostgresDatabase.swift index 64e44abb..fcd1afc7 100644 --- a/Sources/PostgresNIO/PostgresDatabase.swift +++ b/Sources/PostgresNIO/PostgresDatabase.swift @@ -1,14 +1,15 @@ import NIOCore import Logging -public protocol PostgresDatabase { +@preconcurrency +public protocol PostgresDatabase: Sendable { var logger: Logger { get } var eventLoop: EventLoop { get } func send( _ request: PostgresRequest, logger: Logger ) -> EventLoopFuture - + func withConnection(_ closure: @escaping (PostgresConnection) -> EventLoopFuture) -> EventLoopFuture } diff --git a/Sources/PostgresNIO/Utilities/Exports.swift b/Sources/PostgresNIO/Utilities/Exports.swift index 5fc86b74..144ff3c9 100644 --- a/Sources/PostgresNIO/Utilities/Exports.swift +++ b/Sources/PostgresNIO/Utilities/Exports.swift @@ -1,14 +1,3 @@ -#if compiler(>=5.8) - @_documentation(visibility: internal) @_exported import NIO @_documentation(visibility: internal) @_exported import NIOSSL @_documentation(visibility: internal) @_exported import struct Logging.Logger - -#elseif !BUILDING_DOCC - -// TODO: Remove this with the next major release! -@_exported import NIO -@_exported import NIOSSL -@_exported import struct Logging.Logger - -#endif diff --git a/Sources/PostgresNIO/Utilities/PostgresError+Code.swift b/Sources/PostgresNIO/Utilities/PostgresError+Code.swift index 11224f4b..fae903fe 100644 --- a/Sources/PostgresNIO/Utilities/PostgresError+Code.swift +++ b/Sources/PostgresNIO/Utilities/PostgresError+Code.swift @@ -1,5 +1,5 @@ extension PostgresError { - public struct Code: ExpressibleByStringLiteral, Equatable { + public struct Code: Sendable, ExpressibleByStringLiteral, Equatable { // Class 00 — Successful Completion public static let successfulCompletion: Code = "00000" diff --git a/Sources/PostgresNIO/Utilities/PostgresJSONDecoder.swift b/Sources/PostgresNIO/Utilities/PostgresJSONDecoder.swift index fb7b4e8d..ba57ee9b 100644 --- a/Sources/PostgresNIO/Utilities/PostgresJSONDecoder.swift +++ b/Sources/PostgresNIO/Utilities/PostgresJSONDecoder.swift @@ -2,11 +2,13 @@ import class Foundation.JSONDecoder import struct Foundation.Data import NIOFoundationCompat import NIOCore +import NIOConcurrencyHelpers /// A protocol that mimicks the Foundation `JSONDecoder.decode(_:from:)` function. /// Conform a non-Foundation JSON decoder to this protocol if you want PostgresNIO to be /// able to use it when decoding JSON & JSONB values (see `PostgresNIO._defaultJSONDecoder`) -public protocol PostgresJSONDecoder { +@preconcurrency +public protocol PostgresJSONDecoder: Sendable { func decode(_ type: T.Type, from data: Data) throws -> T where T : Decodable func decode(_ type: T.Type, from buffer: ByteBuffer) throws -> T @@ -20,10 +22,20 @@ extension PostgresJSONDecoder { } } +//@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) extension JSONDecoder: PostgresJSONDecoder {} +private let jsonDecoderLocked: NIOLockedValueBox = NIOLockedValueBox(JSONDecoder()) + /// The default JSON decoder used by PostgresNIO when decoding JSON & JSONB values. /// As `_defaultJSONDecoder` will be reused for decoding all JSON & JSONB values /// from potentially multiple threads at once, you must ensure your custom JSON decoder is /// thread safe internally like `Foundation.JSONDecoder`. -public var _defaultJSONDecoder: PostgresJSONDecoder = JSONDecoder() +public var _defaultJSONDecoder: PostgresJSONDecoder { + set { + jsonDecoderLocked.withLockedValue { $0 = newValue } + } + get { + jsonDecoderLocked.withLockedValue { $0 } + } +} diff --git a/Sources/PostgresNIO/Utilities/PostgresJSONEncoder.swift b/Sources/PostgresNIO/Utilities/PostgresJSONEncoder.swift index 735e4b14..9585f20b 100644 --- a/Sources/PostgresNIO/Utilities/PostgresJSONEncoder.swift +++ b/Sources/PostgresNIO/Utilities/PostgresJSONEncoder.swift @@ -1,11 +1,13 @@ import Foundation import NIOFoundationCompat import NIOCore +import NIOConcurrencyHelpers /// A protocol that mimicks the Foundation `JSONEncoder.encode(_:)` function. /// Conform a non-Foundation JSON encoder to this protocol if you want PostgresNIO to be /// able to use it when encoding JSON & JSONB values (see `PostgresNIO._defaultJSONEncoder`) -public protocol PostgresJSONEncoder { +@preconcurrency +public protocol PostgresJSONEncoder: Sendable { func encode(_ value: T) throws -> Data where T : Encodable func encode(_ value: T, into buffer: inout ByteBuffer) throws @@ -20,8 +22,18 @@ extension PostgresJSONEncoder { extension JSONEncoder: PostgresJSONEncoder {} +private let jsonEncoderLocked: NIOLockedValueBox = NIOLockedValueBox(JSONEncoder()) + /// The default JSON encoder used by PostgresNIO when encoding JSON & JSONB values. /// As `_defaultJSONEncoder` will be reused for encoding all JSON & JSONB values /// from potentially multiple threads at once, you must ensure your custom JSON encoder is /// thread safe internally like `Foundation.JSONEncoder`. -public var _defaultJSONEncoder: PostgresJSONEncoder = JSONEncoder() +public var _defaultJSONEncoder: PostgresJSONEncoder { + set { + jsonEncoderLocked.withLockedValue { $0 = newValue } + } + get { + jsonEncoderLocked.withLockedValue { $0 } + } +} + diff --git a/Sources/PostgresNIO/Utilities/SASLAuthentication+SCRAM-SHA256.swift b/Sources/PostgresNIO/Utilities/SASLAuthentication+SCRAM-SHA256.swift index f2fd8e1a..2a717b6b 100644 --- a/Sources/PostgresNIO/Utilities/SASLAuthentication+SCRAM-SHA256.swift +++ b/Sources/PostgresNIO/Utilities/SASLAuthentication+SCRAM-SHA256.swift @@ -1,13 +1,10 @@ import Crypto import Foundation -extension UInt8: ExpressibleByUnicodeScalarLiteral { +extension UInt8 { fileprivate static var NUL: UInt8 { return 0x00 /* yeah, just U+0000 man */ } fileprivate static var comma: UInt8 { return 0x2c /* .init(ascii: ",") */ } fileprivate static var equals: UInt8 { return 0x3d /* .init(ascii: "=") */ } - public init(unicodeScalarLiteral value: Unicode.Scalar) { - self.init(ascii: value) - } } fileprivate extension String { @@ -87,7 +84,7 @@ fileprivate extension Array where Element == UInt8 { */ var isValidScramValue: Bool { // TODO: FInd a better way than doing a whole construction of String... - return self.count > 0 && !(String(bytes: self, encoding: .utf8)?.contains(",") ?? true) + return self.count > 0 && !(String(decoding: self, as: Unicode.UTF8.self).contains(",")) } } @@ -171,52 +168,51 @@ fileprivate struct SCRAMMessageParser { static func parseAttributePair(name: [UInt8], value: [UInt8], isGS2Header: Bool = false) -> SCRAMAttribute? { guard name.count == 1 || isGS2Header else { return nil } switch name.first { - case "m" where !isGS2Header: return .m(value) - case "r" where !isGS2Header: return String(printableAscii: value).map { .r($0) } - case "c" where !isGS2Header: - guard let parsedAttrs = value.decodingBase64().flatMap({ parse(raw: $0, isGS2Header: true) }) else { return nil } - guard (1...3).contains(parsedAttrs.count) else { return nil } - switch (parsedAttrs.first, parsedAttrs.dropFirst(1).first, parsedAttrs.dropFirst(2).first) { - case let (.gp(.bind(name, .none)), .a(ident), .gm(data)): return .c(binding: .bind(name, data), authIdentity: ident) - case let (.gp(.bind(name, .none)), .gm(data), .none): return .c(binding: .bind(name, data)) - case let (.gp(bind), .a(ident), .none): return .c(binding: bind, authIdentity: ident) - case let (.gp(bind), .none, .none): return .c(binding: bind) - default: return nil - } - case "n" where !isGS2Header: return String(bytes: value, encoding: .utf8)?.decodedAsSaslName.map { .n($0) } - case "s" where !isGS2Header: return value.decodingBase64().map { .s($0) } - case "i" where !isGS2Header: return String(printableAscii: value).flatMap { UInt32.init($0) }.map { .i($0) } - case "p" where !isGS2Header: return value.decodingBase64().map { .p($0) } - case "v" where !isGS2Header: return value.decodingBase64().map { .v($0) } - case "e" where !isGS2Header: // TODO: actually map the specific enum string values - guard value.isValidScramValue else { return nil } - return String(bytes: value, encoding: .utf8).flatMap { SCRAMServerError(rawValue: $0) }.map { .e($0) } - - case "y" where isGS2Header && value.count == 0: return .gp(.unused) - case "n" where isGS2Header && value.count == 0: return .gp(.unsupported) - case "p" where isGS2Header: return String(asciiAlphanumericMorse: value).map { .gp(.bind($0, nil)) } - case "a" where isGS2Header: return String(bytes: value, encoding: .utf8)?.decodedAsSaslName.map { .a($0) } - case .none where isGS2Header: return .a(nil) + case UInt8(ascii: "m") where !isGS2Header: return .m(value) + case UInt8(ascii: "r") where !isGS2Header: return String(printableAscii: value).map { .r($0) } + case UInt8(ascii: "c") where !isGS2Header: + guard let parsedAttrs = value.decodingBase64().flatMap({ parse(raw: $0, isGS2Header: true) }) else { return nil } + guard (1...3).contains(parsedAttrs.count) else { return nil } + switch (parsedAttrs.first, parsedAttrs.dropFirst(1).first, parsedAttrs.dropFirst(2).first) { + case let (.gp(.bind(name, .none)), .a(ident), .gm(data)): return .c(binding: .bind(name, data), authIdentity: ident) + case let (.gp(.bind(name, .none)), .gm(data), .none): return .c(binding: .bind(name, data)) + case let (.gp(bind), .a(ident), .none): return .c(binding: bind, authIdentity: ident) + case let (.gp(bind), .none, .none): return .c(binding: bind) + default: return nil + } + case UInt8(ascii: "n") where !isGS2Header: return String(decoding: value, as: Unicode.UTF8.self).decodedAsSaslName.map { .n($0) } + case UInt8(ascii: "s") where !isGS2Header: return value.decodingBase64().map { .s($0) } + case UInt8(ascii: "i") where !isGS2Header: return String(printableAscii: value).flatMap { UInt32.init($0) }.map { .i($0) } + case UInt8(ascii: "p") where !isGS2Header: return value.decodingBase64().map { .p($0) } + case UInt8(ascii: "v") where !isGS2Header: return value.decodingBase64().map { .v($0) } + case UInt8(ascii: "e") where !isGS2Header: // TODO: actually map the specific enum string values + guard value.isValidScramValue else { return nil } + return SCRAMServerError(rawValue: String(decoding: value, as: Unicode.UTF8.self)).flatMap { .e($0) } - default: - if isGS2Header { - return .gm(name + value) - } else { - guard value.count > 0, value.isValidScramValue else { return nil } - return .optional(name: CChar(name[0]), value: value) - } + case UInt8(ascii: "y") where isGS2Header && value.count == 0: return .gp(.unused) + case UInt8(ascii: "n") where isGS2Header && value.count == 0: return .gp(.unsupported) + case UInt8(ascii: "p") where isGS2Header: return String(asciiAlphanumericMorse: value).map { .gp(.bind($0, nil)) } + case UInt8(ascii: "a") where isGS2Header: return String(decoding: value, as: Unicode.UTF8.self).decodedAsSaslName.map { .a($0) } + case .none where isGS2Header: return .a(nil) + + default: + if isGS2Header { + return .gm(name + value) + } else { + guard value.count > 0, value.isValidScramValue else { return nil } + return .optional(name: CChar(name[0]), value: value) + } } } static func parse(raw: [UInt8], isGS2Header: Bool = false) -> [SCRAMAttribute]? { - // There are two ways to implement this parse: // 1. All-at-once: Split on comma, split each on equals, validate // each results in a valid attribute. // 2. Sequential: State machine lookahead parse. // The former is simpler. The latter provides better validation. - let likelyAttributeSets = raw.split(separator: .comma, maxSplits: isGS2Header ? 3 : Int.max, omittingEmptySubsequences: false) - let likelyAttributePairs = likelyAttributeSets.map { $0.split(separator: .equals, maxSplits: 2, omittingEmptySubsequences: false) } + let likelyAttributeSets = raw.split(separator: .comma, maxSplits: isGS2Header ? 2 : Int.max, omittingEmptySubsequences: false) + let likelyAttributePairs = likelyAttributeSets.map { $0.split(separator: .equals, maxSplits: 1, omittingEmptySubsequences: false) } let results = likelyAttributePairs.map { parseAttributePair(name: Array($0[0]), value: $0.dropFirst().first.map { Array($0) } ?? [], isGS2Header: isGS2Header) } let validResults = results.compactMap { $0 } @@ -231,45 +227,45 @@ fileprivate struct SCRAMMessageParser { for attribute in attributes { switch attribute { case .m(let value): - result.append("m"); result.append("="); result.append(contentsOf: value) + result.append(UInt8(ascii: "m")); result.append(.equals); result.append(contentsOf: value) case .r(let nonce): - result.append("r"); result.append("="); result.append(contentsOf: nonce.utf8.map { UInt8($0) }) + result.append(UInt8(ascii: "r")); result.append(.equals); result.append(contentsOf: nonce.utf8.map { UInt8($0) }) case .n(let name): - result.append("n"); result.append("="); result.append(contentsOf: name.encodedAsSaslName.utf8.map { UInt8($0) }) + result.append(UInt8(ascii: "n")); result.append(.equals); result.append(contentsOf: name.encodedAsSaslName.utf8.map { UInt8($0) }) case .s(let salt): - result.append("s"); result.append("="); result.append(contentsOf: salt.encodingBase64()) + result.append(UInt8(ascii: "s")); result.append(.equals); result.append(contentsOf: salt.encodingBase64()) case .i(let count): - result.append("i"); result.append("="); result.append(contentsOf: "\(count)".utf8.map { UInt8($0) }) + result.append(UInt8(ascii: "i")); result.append(.equals); result.append(contentsOf: "\(count)".utf8.map { UInt8($0) }) case .p(let proof): - result.append("p"); result.append("="); result.append(contentsOf: proof.encodingBase64()) + result.append(UInt8(ascii: "p")); result.append(.equals); result.append(contentsOf: proof.encodingBase64()) case .v(let signature): - result.append("v"); result.append("="); result.append(contentsOf: signature.encodingBase64()) + result.append(UInt8(ascii: "v")); result.append(.equals); result.append(contentsOf: signature.encodingBase64()) case .e(let error): - result.append("e"); result.append("="); result.append(contentsOf: error.rawValue.utf8.map { UInt8($0) }) + result.append(UInt8(ascii: "e")); result.append(.equals); result.append(contentsOf: error.rawValue.utf8.map { UInt8($0) }) case .c(let binding, let identity): if isInitialGS2Header { switch binding { - case .unsupported: result.append("n") - case .unused: result.append("y") - case .bind(let name, _): result.append("p"); result.append("="); result.append(contentsOf: name.utf8.map { UInt8($0) }) + case .unsupported: result.append(UInt8(ascii: "n")) + case .unused: result.append(UInt8(ascii: "y")) + case .bind(let name, _): result.append(UInt8(ascii: "p")); result.append(.equals); result.append(contentsOf: name.utf8.map { UInt8($0) }) } - result.append(",") + result.append(.comma) if let identity = identity { - result.append("a"); result.append("="); result.append(contentsOf: identity.encodedAsSaslName.utf8.map { UInt8($0) }) + result.append(UInt8(ascii: "a")); result.append(.equals); result.append(contentsOf: identity.encodedAsSaslName.utf8.map { UInt8($0) }) } - result.append(",") + result.append(.comma) } else { guard var partial = serialize([attribute], isInitialGS2Header: true) else { return nil } if case let .bind(_, data) = binding { guard let data = data else { return nil } partial.append(contentsOf: data) } - result.append("c"); result.append("="); result.append(contentsOf: partial.encodingBase64()) + result.append(UInt8(ascii: "c")); result.append(.equals); result.append(contentsOf: partial.encodingBase64()) } default: return nil } - result.append(",") + result.append(.comma) } return result.dropLast() } @@ -369,7 +365,7 @@ internal struct SHA256_PLUS: SASLAuthenticationMechanism { } // enum SCRAM } // enum SASLMechanism -/// Common impplementation of SCRAM-SHA-256 and SCRAM-SHA-256-PLUS +/// Common implementation of SCRAM-SHA-256 and SCRAM-SHA-256-PLUS fileprivate final class SASLMechanism_SCRAM_SHA256_Common { /// Initialized with initial client state @@ -473,7 +469,7 @@ fileprivate final class SASLMechanism_SCRAM_SHA256_Common { let saltedPassword = Hi(string: password, salt: serverSalt, iterations: serverIterations) let clientKey = HMAC.authenticationCode(for: "Client Key".data(using: .utf8)!, using: .init(data: saltedPassword)) let storedKey = SHA256.hash(data: Data(clientKey)) - var authMessage = firstMessageBare; authMessage.append(","); authMessage.append(contentsOf: message); authMessage.append(","); authMessage.append(contentsOf: clientFinalNoProof) + var authMessage = firstMessageBare; authMessage.append(.comma); authMessage.append(contentsOf: message); authMessage.append(.comma); authMessage.append(contentsOf: clientFinalNoProof) let clientSignature = HMAC.authenticationCode(for: authMessage, using: .init(data: storedKey)) var clientProof = Array(clientKey) @@ -486,7 +482,7 @@ fileprivate final class SASLMechanism_SCRAM_SHA256_Common { } // Generate a `client-final-message` - var clientFinalMessage = clientFinalNoProof; clientFinalMessage.append(",") + var clientFinalMessage = clientFinalNoProof; clientFinalMessage.append(.comma) guard let proofPart = SCRAMMessageParser.serialize([.p(Array(clientProof))]) else { throw SASLAuthenticationError.genericAuthenticationFailure } clientFinalMessage.append(contentsOf: proofPart) @@ -591,7 +587,7 @@ fileprivate final class SASLMechanism_SCRAM_SHA256_Common { // Compute client signature let clientKey = HMAC.authenticationCode(for: "Client Key".data(using: .utf8)!, using: .init(data: saltedPassword)) let storedKey = SHA256.hash(data: Data(clientKey)) - var authMessage = clientBareFirstMessage; authMessage.append(","); authMessage.append(contentsOf: serverFirstMessage); authMessage.append(","); authMessage.append(contentsOf: message.dropLast(proof.count + 3)) + var authMessage = clientBareFirstMessage; authMessage.append(.comma); authMessage.append(contentsOf: serverFirstMessage); authMessage.append(.comma); authMessage.append(contentsOf: message.dropLast(proof.count + 3)) let clientSignature = HMAC.authenticationCode(for: authMessage, using: .init(data: storedKey)) // Recompute client key from signature and proof, verify match diff --git a/Tests/ConnectionPoolModuleTests/ConnectionIDGeneratorTests.swift b/Tests/ConnectionPoolModuleTests/ConnectionIDGeneratorTests.swift new file mode 100644 index 00000000..fb0bfce1 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/ConnectionIDGeneratorTests.swift @@ -0,0 +1,22 @@ +import _ConnectionPoolModule +import XCTest + +final class ConnectionIDGeneratorTests: XCTestCase { + func testGenerateConnectionIDs() async { + let idGenerator = ConnectionIDGenerator() + + XCTAssertEqual(idGenerator.next(), 0) + XCTAssertEqual(idGenerator.next(), 1) + XCTAssertEqual(idGenerator.next(), 2) + + await withTaskGroup(of: Void.self) { taskGroup in + for _ in 0..<1000 { + taskGroup.addTask { + _ = idGenerator.next() + } + } + } + + XCTAssertEqual(idGenerator.next(), 1003) + } +} diff --git a/Tests/ConnectionPoolModuleTests/ConnectionPoolTests.swift b/Tests/ConnectionPoolModuleTests/ConnectionPoolTests.swift new file mode 100644 index 00000000..c745b4a0 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/ConnectionPoolTests.swift @@ -0,0 +1,858 @@ +@testable import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import Atomics +import NIOEmbedded +import XCTest + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class ConnectionPoolTests: XCTestCase { + + func test1000ConsecutiveRequestsOnSingleConnection() async { + let factory = MockConnectionFactory() + + var config = ConnectionPoolConfiguration() + config.minimumConnectionCount = 1 + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self), + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: ContinuousClock() + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + // the same connection is reused 1000 times + + await withTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask_ { + await pool.run() + } + + let createdConnection = await factory.nextConnectAttempt { _ in + return 1 + } + XCTAssertNotNil(createdConnection) + + do { + for _ in 0..<1000 { + async let connectionFuture = try await pool.leaseConnection() + var leasedConnection: MockConnection? + XCTAssertEqual(factory.pendingConnectionAttemptsCount, 0) + leasedConnection = try await connectionFuture + XCTAssertNotNil(leasedConnection) + XCTAssert(createdConnection === leasedConnection) + + if let leasedConnection { + pool.releaseConnection(leasedConnection) + } + } + } catch { + XCTFail("Unexpected error: \(error)") + } + + taskGroup.cancelAll() + + XCTAssertEqual(factory.pendingConnectionAttemptsCount, 0) + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + + XCTAssertEqual(factory.runningConnections.count, 0) + } + + func testShutdownPoolWhileConnectionIsBeingCreated() async { + let clock = MockClock() + let factory = MockConnectionFactory() + + var config = ConnectionPoolConfiguration() + config.minimumConnectionCount = 1 + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self), + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + await withTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask_ { + await pool.run() + } + + let (blockCancelStream, blockCancelContinuation) = AsyncStream.makeStream(of: Void.self) + let (blockConnCreationStream, blockConnCreationContinuation) = AsyncStream.makeStream(of: Void.self) + + taskGroup.addTask_ { + _ = try? await factory.nextConnectAttempt { _ in + blockCancelContinuation.yield() + var iterator = blockConnCreationStream.makeAsyncIterator() + await iterator.next() + throw ConnectionCreationError() + } + } + + var iterator = blockCancelStream.makeAsyncIterator() + await iterator.next() + + taskGroup.cancelAll() + blockConnCreationContinuation.yield() + } + + struct ConnectionCreationError: Error {} + } + + func testShutdownPoolWhileConnectionIsBackingOff() async { + let clock = MockClock() + let factory = MockConnectionFactory() + + var config = ConnectionPoolConfiguration() + config.minimumConnectionCount = 1 + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self), + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + await withTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask_ { + await pool.run() + } + + _ = try? await factory.nextConnectAttempt { _ in + throw ConnectionCreationError() + } + + await clock.nextTimerScheduled() + + taskGroup.cancelAll() + } + + struct ConnectionCreationError: Error {} + } + + func testConnectionHardLimitIsRespected() async { + let factory = MockConnectionFactory() + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 4 + mutableConfig.maximumConnectionHardLimit = 8 + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self), + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: ContinuousClock() + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + let hasFinished = ManagedAtomic(false) + let createdConnections = ManagedAtomic(0) + let iterations = 10_000 + + // the same connection is reused 1000 times + + await withTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask_ { + await pool.run() + XCTAssertFalse(hasFinished.compareExchange(expected: false, desired: true, ordering: .relaxed).original) + } + + taskGroup.addTask_ { + var usedConnectionIDs = Set() + for _ in 0..() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 1 + mutableConfig.maximumConnectionHardLimit = 1 + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + async let lease1ConnectionAsync = pool.leaseConnection() + + let connection = await factory.nextConnectAttempt { connectionID in + return 1 + } + + let lease1Connection = try await lease1ConnectionAsync + XCTAssert(connection === lease1Connection) + + pool.releaseConnection(lease1Connection) + + // keep alive 1 + + // validate that a keep alive timer and an idle timeout timer is scheduled + var expectedInstants: Set = [.init(keepAliveDuration), .init(config.idleTimeout)] + let deadline1 = await clock.nextTimerScheduled() + print(deadline1) + XCTAssertNotNil(expectedInstants.remove(deadline1)) + let deadline2 = await clock.nextTimerScheduled() + print(deadline2) + XCTAssertNotNil(expectedInstants.remove(deadline2)) + XCTAssert(expectedInstants.isEmpty) + + // move clock forward to keep alive + let newTime = clock.now.advanced(by: keepAliveDuration) + clock.advance(to: newTime) + print("clock advanced to: \(newTime)") + + await keepAlive.nextKeepAlive { keepAliveConnection in + defer { print("keep alive 1 has run") } + XCTAssertTrue(keepAliveConnection === lease1Connection) + return true + } + + // keep alive 2 + + let deadline3 = await clock.nextTimerScheduled() + XCTAssertEqual(deadline3, clock.now.advanced(by: keepAliveDuration)) + print(deadline3) + + // race keep alive vs timeout + clock.advance(to: clock.now.advanced(by: keepAliveDuration)) + + taskGroup.cancelAll() + + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } + + func testKeepAliveOnClose() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(20) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 1 + mutableConfig.maximumConnectionHardLimit = 1 + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + async let lease1ConnectionAsync = pool.leaseConnection() + + let connection = await factory.nextConnectAttempt { connectionID in + return 1 + } + + let lease1Connection = try await lease1ConnectionAsync + XCTAssert(connection === lease1Connection) + + pool.releaseConnection(lease1Connection) + + // keep alive 1 + + // validate that a keep alive timer and an idle timeout timer is scheduled + var expectedInstants: Set = [.init(keepAliveDuration), .init(config.idleTimeout)] + let deadline1 = await clock.nextTimerScheduled() + print(deadline1) + XCTAssertNotNil(expectedInstants.remove(deadline1)) + let deadline2 = await clock.nextTimerScheduled() + print(deadline2) + XCTAssertNotNil(expectedInstants.remove(deadline2)) + XCTAssert(expectedInstants.isEmpty) + + // move clock forward to keep alive + let newTime = clock.now.advanced(by: keepAliveDuration) + clock.advance(to: newTime) + + await keepAlive.nextKeepAlive { keepAliveConnection in + XCTAssertTrue(keepAliveConnection === lease1Connection) + return true + } + + // keep alive 2 + let deadline3 = await clock.nextTimerScheduled() + XCTAssertEqual(deadline3, clock.now.advanced(by: keepAliveDuration)) + clock.advance(to: clock.now.advanced(by: keepAliveDuration)) + + let failingKeepAliveDidRun = ManagedAtomic(false) + // the following keep alive should not cause a crash + _ = try? await keepAlive.nextKeepAlive { keepAliveConnection in + defer { + XCTAssertFalse(failingKeepAliveDidRun + .compareExchange(expected: false, desired: true, ordering: .relaxed).original) + } + XCTAssertTrue(keepAliveConnection === lease1Connection) + keepAliveConnection.close() + throw CancellationError() // any error + } // will fail and it's expected + XCTAssertTrue(failingKeepAliveDidRun.load(ordering: .relaxed)) + + taskGroup.cancelAll() + + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } + + func testKeepAliveWorksRacesAgainstShutdown() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 1 + mutableConfig.maximumConnectionHardLimit = 1 + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + async let lease1ConnectionAsync = pool.leaseConnection() + + let connection = await factory.nextConnectAttempt { connectionID in + return 1 + } + + let lease1Connection = try await lease1ConnectionAsync + XCTAssert(connection === lease1Connection) + + pool.releaseConnection(lease1Connection) + + // keep alive 1 + + // validate that a keep alive timer and an idle timeout timer is scheduled + var expectedInstants: Set = [.init(keepAliveDuration), .init(config.idleTimeout)] + let deadline1 = await clock.nextTimerScheduled() + print(deadline1) + XCTAssertNotNil(expectedInstants.remove(deadline1)) + let deadline2 = await clock.nextTimerScheduled() + print(deadline2) + XCTAssertNotNil(expectedInstants.remove(deadline2)) + XCTAssert(expectedInstants.isEmpty) + + clock.advance(to: clock.now.advanced(by: keepAliveDuration)) + + await keepAlive.nextKeepAlive { keepAliveConnection in + defer { print("keep alive 1 has run") } + XCTAssertTrue(keepAliveConnection === lease1Connection) + return true + } + + taskGroup.cancelAll() + print("cancelled") + + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } + + func testCancelConnectionRequestWorks() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 4 + mutableConfig.maximumConnectionHardLimit = 4 + mutableConfig.idleTimeout = .seconds(10) + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + let leaseTask = Task { + _ = try await pool.leaseConnection() + } + + let connectionAttemptWaiter = Future(of: Void.self) + + taskGroup.addTask { + try await factory.nextConnectAttempt { connectionID in + connectionAttemptWaiter.yield(value: ()) + throw CancellationError() + } + } + + try await connectionAttemptWaiter.success + leaseTask.cancel() + + let taskResult = await leaseTask.result + switch taskResult { + case .success: + XCTFail("Expected task failure") + case .failure(let failure): + XCTAssertEqual(failure as? ConnectionPoolError, .requestCancelled) + } + + taskGroup.cancelAll() + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } + + func testLeasingMultipleConnectionsAtOnceWorks() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 4 + mutableConfig.maximumConnectionSoftLimit = 4 + mutableConfig.maximumConnectionHardLimit = 4 + mutableConfig.idleTimeout = .seconds(10) + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionFuture.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + // create 4 persisted connections + for _ in 0..<4 { + await factory.nextConnectAttempt { connectionID in + return 1 + } + } + + // create 4 connection requests + let requests = (0..<4).map { ConnectionFuture(id: $0) } + + // lease 4 connections at once + pool.leaseConnections(requests) + var connections = [MockConnection]() + + for request in requests { + let connection = try await request.future.success + connections.append(connection) + } + + // Ensure that we got 4 distinct connections + XCTAssertEqual(Set(connections.lazy.map(\.id)).count, 4) + + // release all 4 leased connections + for connection in connections { + pool.releaseConnection(connection) + } + + // shutdown + taskGroup.cancelAll() + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } + + func testLeasingConnectionAfterShutdownIsInvokedFails() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 4 + mutableConfig.maximumConnectionSoftLimit = 4 + mutableConfig.maximumConnectionHardLimit = 4 + mutableConfig.idleTimeout = .seconds(10) + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionRequest.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + // create 4 persisted connections + for _ in 0..<4 { + await factory.nextConnectAttempt { connectionID in + return 1 + } + } + + // shutdown + taskGroup.cancelAll() + + do { + _ = try await pool.leaseConnection() + XCTFail("Expected a failure") + } catch { + print("failed") + XCTAssertEqual(error as? ConnectionPoolError, .poolShutdown) + } + + print("will close connections: \(factory.runningConnections)") + for connection in factory.runningConnections { + try await connection.signalToClose + connection.closeIfClosing() + } + } + } + + func testLeasingConnectionsAfterShutdownIsInvokedFails() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 4 + mutableConfig.maximumConnectionSoftLimit = 4 + mutableConfig.maximumConnectionHardLimit = 4 + mutableConfig.idleTimeout = .seconds(10) + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionFuture.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + // create 4 persisted connections + for _ in 0..<4 { + await factory.nextConnectAttempt { connectionID in + return 1 + } + } + + // shutdown + taskGroup.cancelAll() + + // create 4 connection requests + let requests = (0..<4).map { ConnectionFuture(id: $0) } + + // lease 4 connections at once + pool.leaseConnections(requests) + + for request in requests { + do { + _ = try await request.future.success + XCTFail("Expected a failure") + } catch { + XCTAssertEqual(error as? ConnectionPoolError, .poolShutdown) + } + } + + for connection in factory.runningConnections { + try await connection.signalToClose + connection.closeIfClosing() + } + } + } + + func testLeasingMultipleStreamsFromOneConnectionWorks() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 1 + mutableConfig.maximumConnectionHardLimit = 10 + mutableConfig.idleTimeout = .seconds(10) + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionFuture.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + // create 4 connection requests + let requests = (0..<10).map { ConnectionFuture(id: $0) } + pool.leaseConnections(requests) + var connections = [MockConnection]() + + await factory.nextConnectAttempt { connectionID in + return 10 + } + + for request in requests { + let connection = try await request.future.success + connections.append(connection) + } + + // Ensure that all requests got the same connection + XCTAssertEqual(Set(connections.lazy.map(\.id)).count, 1) + + // release all 10 leased streams + for connection in connections { + pool.releaseConnection(connection) + } + + for _ in 0..<9 { + _ = try? await factory.nextConnectAttempt { connectionID in + throw CancellationError() + } + } + + // shutdown + taskGroup.cancelAll() + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } + + func testIncreasingAvailableStreamsWorks() async throws { + let clock = MockClock() + let factory = MockConnectionFactory() + let keepAliveDuration = Duration.seconds(30) + let keepAlive = MockPingPongBehavior(keepAliveFrequency: keepAliveDuration, connectionType: MockConnection.self) + + var mutableConfig = ConnectionPoolConfiguration() + mutableConfig.minimumConnectionCount = 0 + mutableConfig.maximumConnectionSoftLimit = 1 + mutableConfig.maximumConnectionHardLimit = 1 + mutableConfig.idleTimeout = .seconds(10) + let config = mutableConfig + + let pool = ConnectionPool( + configuration: config, + idGenerator: ConnectionIDGenerator(), + requestType: ConnectionFuture.self, + keepAliveBehavior: keepAlive, + observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self), + clock: clock + ) { + try await factory.makeConnection(id: $0, for: $1) + } + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await pool.run() + } + + // create 4 connection requests + var requests = (0..<21).map { ConnectionFuture(id: $0) } + pool.leaseConnections(requests) + var connections = [MockConnection]() + + await factory.nextConnectAttempt { connectionID in + return 1 + } + + let connection = try await requests.first!.future.success + connections.append(connection) + requests.removeFirst() + + pool.connectionReceivedNewMaxStreamSetting(connection, newMaxStreamSetting: 21) + + for (_, request) in requests.enumerated() { + let connection = try await request.future.success + connections.append(connection) + } + + // Ensure that all requests got the same connection + XCTAssertEqual(Set(connections.lazy.map(\.id)).count, 1) + + requests = (22..<42).map { ConnectionFuture(id: $0) } + pool.leaseConnections(requests) + + // release all 21 leased streams in a single call + pool.releaseConnection(connection, streams: 21) + + // ensure all 20 new requests got fulfilled + for request in requests { + let connection = try await request.future.success + connections.append(connection) + } + + // release all 20 leased streams one by one + for _ in requests { + pool.releaseConnection(connection, streams: 1) + } + + // shutdown + taskGroup.cancelAll() + for connection in factory.runningConnections { + connection.closeIfClosing() + } + } + } +} + +struct ConnectionFuture: ConnectionRequestProtocol { + let id: Int + let future: Future + + init(id: Int) { + self.id = id + self.future = Future(of: MockConnection.self) + } + + func complete(with result: Result) { + switch result { + case .success(let success): + self.future.yield(value: success) + case .failure(let failure): + self.future.yield(error: failure) + } + } +} diff --git a/Tests/ConnectionPoolModuleTests/ConnectionRequestTests.swift b/Tests/ConnectionPoolModuleTests/ConnectionRequestTests.swift new file mode 100644 index 00000000..537efbd9 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/ConnectionRequestTests.swift @@ -0,0 +1,28 @@ +@testable import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import XCTest + +final class ConnectionRequestTests: XCTestCase { + + func testHappyPath() async throws { + let mockConnection = MockConnection(id: 1) + let connection = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + let request = ConnectionRequest(id: 42, continuation: continuation) + XCTAssertEqual(request.id, 42) + continuation.resume(with: .success(mockConnection)) + } + + XCTAssert(connection === mockConnection) + } + + func testSadPath() async throws { + do { + _ = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + continuation.resume(with: .failure(ConnectionPoolError.requestCancelled)) + } + XCTFail("This point should not be reached") + } catch { + XCTAssertEqual(error as? ConnectionPoolError, .requestCancelled) + } + } +} diff --git a/Tests/ConnectionPoolModuleTests/Max2SequenceTests.swift b/Tests/ConnectionPoolModuleTests/Max2SequenceTests.swift new file mode 100644 index 00000000..081e867b --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/Max2SequenceTests.swift @@ -0,0 +1,60 @@ +@testable import _ConnectionPoolModule +import XCTest + +final class Max2SequenceTests: XCTestCase { + func testCountAndIsEmpty() async { + var sequence = Max2Sequence() + XCTAssertEqual(sequence.count, 0) + XCTAssertEqual(sequence.isEmpty, true) + sequence.append(1) + XCTAssertEqual(sequence.count, 1) + XCTAssertEqual(sequence.isEmpty, false) + sequence.append(2) + XCTAssertEqual(sequence.count, 2) + XCTAssertEqual(sequence.isEmpty, false) + } + + func testOptionalInitializer() { + let emptySequence = Max2Sequence(nil, nil) + XCTAssertEqual(emptySequence.count, 0) + XCTAssertEqual(emptySequence.isEmpty, true) + var emptySequenceIterator = emptySequence.makeIterator() + XCTAssertNil(emptySequenceIterator.next()) + XCTAssertNil(emptySequenceIterator.next()) + XCTAssertNil(emptySequenceIterator.next()) + + let oneElemSequence1 = Max2Sequence(1, nil) + XCTAssertEqual(oneElemSequence1.count, 1) + XCTAssertEqual(oneElemSequence1.isEmpty, false) + var oneElemSequence1Iterator = oneElemSequence1.makeIterator() + XCTAssertEqual(oneElemSequence1Iterator.next(), 1) + XCTAssertNil(oneElemSequence1Iterator.next()) + XCTAssertNil(oneElemSequence1Iterator.next()) + + let oneElemSequence2 = Max2Sequence(nil, 2) + XCTAssertEqual(oneElemSequence2.count, 1) + XCTAssertEqual(oneElemSequence2.isEmpty, false) + var oneElemSequence2Iterator = oneElemSequence2.makeIterator() + XCTAssertEqual(oneElemSequence2Iterator.next(), 2) + XCTAssertNil(oneElemSequence2Iterator.next()) + XCTAssertNil(oneElemSequence2Iterator.next()) + + let twoElemSequence = Max2Sequence(1, 2) + XCTAssertEqual(twoElemSequence.count, 2) + XCTAssertEqual(twoElemSequence.isEmpty, false) + var twoElemSequenceIterator = twoElemSequence.makeIterator() + XCTAssertEqual(twoElemSequenceIterator.next(), 1) + XCTAssertEqual(twoElemSequenceIterator.next(), 2) + XCTAssertNil(twoElemSequenceIterator.next()) + } + + func testMap() { + let twoElemSequence = Max2Sequence(1, 2).map({ "\($0)" }) + XCTAssertEqual(twoElemSequence.count, 2) + XCTAssertEqual(twoElemSequence.isEmpty, false) + var twoElemSequenceIterator = twoElemSequence.makeIterator() + XCTAssertEqual(twoElemSequenceIterator.next(), "1") + XCTAssertEqual(twoElemSequenceIterator.next(), "2") + XCTAssertNil(twoElemSequenceIterator.next()) + } +} diff --git a/Tests/ConnectionPoolModuleTests/Mocks/MockTimerCancellationToken.swift b/Tests/ConnectionPoolModuleTests/Mocks/MockTimerCancellationToken.swift new file mode 100644 index 00000000..27035ee9 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/Mocks/MockTimerCancellationToken.swift @@ -0,0 +1,18 @@ +@testable import _ConnectionPoolModule + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +struct MockTimerCancellationToken: Hashable, Sendable { + enum Backing: Hashable, Sendable { + case timer(TestPoolStateMachine.Timer) + case connectionTimer(TestPoolStateMachine.ConnectionTimer) + } + var backing: Backing + + init(_ timer: TestPoolStateMachine.Timer) { + self.backing = .timer(timer) + } + + init(_ timer: TestPoolStateMachine.ConnectionTimer) { + self.backing = .connectionTimer(timer) + } +} diff --git a/Tests/ConnectionPoolModuleTests/NoKeepAliveBehaviorTests.swift b/Tests/ConnectionPoolModuleTests/NoKeepAliveBehaviorTests.swift new file mode 100644 index 00000000..b1b54023 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/NoKeepAliveBehaviorTests.swift @@ -0,0 +1,11 @@ +import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import XCTest + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class NoKeepAliveBehaviorTests: XCTestCase { + func testNoKeepAlive() { + let keepAliveBehavior = NoOpKeepAliveBehavior(connectionType: MockConnection.self) + XCTAssertNil(keepAliveBehavior.keepAliveFrequency) + } +} diff --git a/Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionGroupTests.swift b/Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionGroupTests.swift new file mode 100644 index 00000000..b09bfcb4 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionGroupTests.swift @@ -0,0 +1,328 @@ +@testable import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import XCTest + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class PoolStateMachine_ConnectionGroupTests: XCTestCase { + var idGenerator: ConnectionIDGenerator! + + override func setUp() { + self.idGenerator = ConnectionIDGenerator() + super.setUp() + } + + override func tearDown() { + self.idGenerator = nil + super.tearDown() + } + + func testRefillConnections() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 4, + maximumConcurrentConnectionSoftLimit: 4, + maximumConcurrentConnectionHardLimit: 4, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + XCTAssertTrue(connections.isEmpty) + let requests = connections.refillConnections() + XCTAssertFalse(connections.isEmpty) + + XCTAssertEqual(requests.count, 4) + XCTAssertNil(connections.createNewDemandConnectionIfPossible()) + XCTAssertNil(connections.createNewOverflowConnectionIfPossible()) + XCTAssertEqual(connections.stats, .init(connecting: 4)) + XCTAssertEqual(connections.soonAvailableConnections, 4) + + let requests2 = connections.refillConnections() + XCTAssertTrue(requests2.isEmpty) + + var connected: UInt16 = 0 + for request in requests { + let newConnection = MockConnection(id: request.connectionID) + let (_, context) = connections.newConnectionEstablished(newConnection, maxStreams: 1) + XCTAssertEqual(context.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(context.use, .persisted) + connected += 1 + XCTAssertEqual(connections.stats, .init(connecting: 4 - connected, idle: connected, availableStreams: connected)) + XCTAssertEqual(connections.soonAvailableConnections, 4 - connected) + } + + let requests3 = connections.refillConnections() + XCTAssertTrue(requests3.isEmpty) + } + + func testMakeConnectionLeaseItAndDropItHappyPath() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 0, + maximumConcurrentConnectionSoftLimit: 4, + maximumConcurrentConnectionHardLimit: 4, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + let requests = connections.refillConnections() + XCTAssertTrue(connections.isEmpty) + XCTAssertTrue(requests.isEmpty) + + guard let request = connections.createNewDemandConnectionIfPossible() else { + return XCTFail("Expected to receive a connection request") + } + XCTAssertEqual(request, .init(connectionID: 0)) + XCTAssertFalse(connections.isEmpty) + XCTAssertEqual(connections.soonAvailableConnections, 1) + XCTAssertEqual(connections.stats, .init(connecting: 1)) + + let newConnection = MockConnection(id: request.connectionID) + let (_, establishedContext) = connections.newConnectionEstablished(newConnection, maxStreams: 1) + XCTAssertEqual(establishedContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(establishedContext.use, .demand) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + XCTAssertEqual(connections.soonAvailableConnections, 0) + + guard case .leasedConnection(let leaseResult) = connections.leaseConnectionOrSoonAvailableConnectionCount() else { + return XCTFail("Expected to lease a connection") + } + XCTAssert(newConnection === leaseResult.connection) + XCTAssertEqual(connections.stats, .init(leased: 1, leasedStreams: 1)) + + guard let (index, releasedContext) = connections.releaseConnection(leaseResult.connection.id, streams: 1) else { + return XCTFail("Expected that this connection is still active") + } + XCTAssertEqual(releasedContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(releasedContext.use, .demand) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + + let parkTimers = connections.parkConnection(at: index, hasBecomeIdle: true) + XCTAssertEqual(parkTimers, [ + .init(timerID: 0, connectionID: newConnection.id, usecase: .keepAlive), + .init(timerID: 1, connectionID: newConnection.id, usecase: .idleTimeout), + ]) + + guard let keepAliveAction = connections.keepAliveIfIdle(newConnection.id) else { + return XCTFail("Expected to get a connection for ping pong") + } + XCTAssert(newConnection === keepAliveAction.connection) + XCTAssertEqual(connections.stats, .init(idle: 1, runningKeepAlive: 1, availableStreams: 0)) + + guard let (_, pingPongContext) = connections.keepAliveSucceeded(newConnection.id) else { + return XCTFail("Expected to get an AvailableContext") + } + XCTAssertEqual(pingPongContext.info, .idle(availableStreams: 1, newIdle: false)) + XCTAssertEqual(releasedContext.use, .demand) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + + guard let closeAction = connections.closeConnectionIfIdle(newConnection.id) else { + return XCTFail("Expected to get a connection for ping pong") + } + XCTAssertEqual(closeAction.timersToCancel, []) + XCTAssert(closeAction.connection === newConnection) + XCTAssertEqual(connections.stats, .init(closing: 1, availableStreams: 0)) + + let closeContext = connections.connectionClosed(newConnection.id) + XCTAssertEqual(closeContext.connectionsStarting, 0) + XCTAssertTrue(connections.isEmpty) + XCTAssertEqual(connections.stats, .init()) + } + + func testBackoffDoneCreatesANewConnectionToReachMinimumConnectionsEvenThoughRetryIsSetToFalse() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 1, + maximumConcurrentConnectionSoftLimit: 4, + maximumConcurrentConnectionHardLimit: 4, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + let requests = connections.refillConnections() + XCTAssertEqual(connections.stats, .init(connecting: 1)) + XCTAssertEqual(connections.soonAvailableConnections, 1) + XCTAssertFalse(connections.isEmpty) + XCTAssertEqual(requests.count, 1) + + guard let request = requests.first else { return XCTFail("Expected to receive a connection request") } + XCTAssertEqual(request, .init(connectionID: 0)) + + let backoffTimer = connections.backoffNextConnectionAttempt(request.connectionID) + XCTAssertEqual(connections.stats, .init(backingOff: 1)) + let backoffTimerCancellationToken = MockTimerCancellationToken(backoffTimer) + XCTAssertNil(connections.timerScheduled(backoffTimer, cancelContinuation: backoffTimerCancellationToken)) + + let backoffDoneAction = connections.backoffDone(request.connectionID, retry: false) + XCTAssertEqual(backoffDoneAction, .createConnection(.init(connectionID: 0), backoffTimerCancellationToken)) + + XCTAssertEqual(connections.stats, .init(connecting: 1)) + } + + func testBackoffDoneCancelsIdleTimerIfAPersistedConnectionIsNotRetried() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 2, + maximumConcurrentConnectionSoftLimit: 4, + maximumConcurrentConnectionHardLimit: 4, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + let requests = connections.refillConnections() + XCTAssertEqual(connections.stats, .init(connecting: 2)) + XCTAssertEqual(connections.soonAvailableConnections, 2) + XCTAssertFalse(connections.isEmpty) + XCTAssertEqual(requests.count, 2) + + var requestIterator = requests.makeIterator() + guard let firstRequest = requestIterator.next(), let secondRequest = requestIterator.next() else { + return XCTFail("Expected to get two requests") + } + + guard let thirdRequest = connections.createNewDemandConnectionIfPossible() else { + return XCTFail("Expected to get another request") + } + XCTAssertEqual(connections.stats, .init(connecting: 3)) + + let newSecondConnection = MockConnection(id: secondRequest.connectionID) + let (_, establishedSecondConnectionContext) = connections.newConnectionEstablished(newSecondConnection, maxStreams: 1) + XCTAssertEqual(establishedSecondConnectionContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(establishedSecondConnectionContext.use, .persisted) + XCTAssertEqual(connections.stats, .init(connecting: 2, idle: 1, availableStreams: 1)) + XCTAssertEqual(connections.soonAvailableConnections, 2) + + let newThirdConnection = MockConnection(id: thirdRequest.connectionID) + let (thirdConnectionIndex, establishedThirdConnectionContext) = connections.newConnectionEstablished(newThirdConnection, maxStreams: 1) + XCTAssertEqual(establishedThirdConnectionContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(establishedThirdConnectionContext.use, .demand) + XCTAssertEqual(connections.stats, .init(connecting: 1, idle: 2, availableStreams: 2)) + XCTAssertEqual(connections.soonAvailableConnections, 1) + let thirdConnKeepTimer = TestPoolStateMachine.ConnectionTimer(timerID: 0, connectionID: thirdRequest.connectionID, usecase: .keepAlive) + let thirdConnIdleTimer = TestPoolStateMachine.ConnectionTimer(timerID: 1, connectionID: thirdRequest.connectionID, usecase: .idleTimeout) + let thirdConnIdleTimerCancellationToken = MockTimerCancellationToken(thirdConnIdleTimer) + XCTAssertEqual(connections.parkConnection(at: thirdConnectionIndex, hasBecomeIdle: true), [thirdConnKeepTimer, thirdConnIdleTimer]) + + XCTAssertNil(connections.timerScheduled(thirdConnKeepTimer, cancelContinuation: .init(thirdConnKeepTimer))) + XCTAssertNil(connections.timerScheduled(thirdConnIdleTimer, cancelContinuation: thirdConnIdleTimerCancellationToken)) + + let backoffTimer = connections.backoffNextConnectionAttempt(firstRequest.connectionID) + XCTAssertEqual(connections.stats, .init(backingOff: 1, idle: 2, availableStreams: 2)) + let backoffTimerCancellationToken = MockTimerCancellationToken(backoffTimer) + XCTAssertNil(connections.timerScheduled(backoffTimer, cancelContinuation: backoffTimerCancellationToken)) + XCTAssertEqual(connections.stats, .init(backingOff: 1, idle: 2, availableStreams: 2)) + + // connection three should be moved to connection one and for this reason become permanent + + XCTAssertEqual(connections.backoffDone(firstRequest.connectionID, retry: false), .cancelTimers([backoffTimerCancellationToken, thirdConnIdleTimerCancellationToken])) + XCTAssertEqual(connections.stats, .init(idle: 2, availableStreams: 2)) + + XCTAssertNil(connections.closeConnectionIfIdle(newThirdConnection.id)) + } + + func testBackoffDoneReturnsNilIfOverflowConnection() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 0, + maximumConcurrentConnectionSoftLimit: 4, + maximumConcurrentConnectionHardLimit: 4, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + guard let firstRequest = connections.createNewDemandConnectionIfPossible() else { + return XCTFail("Expected to get two requests") + } + + guard let secondRequest = connections.createNewDemandConnectionIfPossible() else { + return XCTFail("Expected to get another request") + } + XCTAssertEqual(connections.stats, .init(connecting: 2)) + + let newFirstConnection = MockConnection(id: firstRequest.connectionID) + let (_, establishedFirstConnectionContext) = connections.newConnectionEstablished(newFirstConnection, maxStreams: 1) + XCTAssertEqual(establishedFirstConnectionContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(establishedFirstConnectionContext.use, .demand) + XCTAssertEqual(connections.stats, .init(connecting: 1, idle: 1, availableStreams: 1)) + XCTAssertEqual(connections.soonAvailableConnections, 1) + + let backoffTimer = connections.backoffNextConnectionAttempt(secondRequest.connectionID) + let backoffTimerCancellationToken = MockTimerCancellationToken(backoffTimer) + XCTAssertEqual(connections.stats, .init(backingOff: 1, idle: 1, availableStreams: 1)) + XCTAssertNil(connections.timerScheduled(backoffTimer, cancelContinuation: backoffTimerCancellationToken)) + + XCTAssertEqual(connections.backoffDone(secondRequest.connectionID, retry: false), .cancelTimers([backoffTimerCancellationToken])) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + + XCTAssertNotNil(connections.closeConnectionIfIdle(newFirstConnection.id)) + } + + func testPingPong() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 1, + maximumConcurrentConnectionSoftLimit: 4, + maximumConcurrentConnectionHardLimit: 4, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + let requests = connections.refillConnections() + XCTAssertFalse(connections.isEmpty) + XCTAssertEqual(connections.stats, .init(connecting: 1)) + + XCTAssertEqual(requests.count, 1) + guard let firstRequest = requests.first else { return XCTFail("Expected to have a request here") } + + let newConnection = MockConnection(id: firstRequest.connectionID) + let (connectionIndex, establishedConnectionContext) = connections.newConnectionEstablished(newConnection, maxStreams: 1) + XCTAssertEqual(establishedConnectionContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(establishedConnectionContext.use, .persisted) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + let timers = connections.parkConnection(at: connectionIndex, hasBecomeIdle: true) + let keepAliveTimer = TestPoolStateMachine.ConnectionTimer(timerID: 0, connectionID: firstRequest.connectionID, usecase: .keepAlive) + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + XCTAssertEqual(timers, [keepAliveTimer]) + XCTAssertNil(connections.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + let keepAliveAction = connections.keepAliveIfIdle(newConnection.id) + XCTAssertEqual(keepAliveAction, .init(connection: newConnection, keepAliveTimerCancellationContinuation: keepAliveTimerCancellationToken)) + XCTAssertEqual(connections.stats, .init(idle: 1, runningKeepAlive: 1, availableStreams: 0)) + + guard let (_, afterPingIdleContext) = connections.keepAliveSucceeded(newConnection.id) else { + return XCTFail("Expected to receive an AvailableContext") + } + XCTAssertEqual(afterPingIdleContext.info, .idle(availableStreams: 1, newIdle: false)) + XCTAssertEqual(afterPingIdleContext.use, .persisted) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + } + + func testKeepAliveShouldNotIndicateCloseConnectionAfterClosed() { + var connections = TestPoolStateMachine.ConnectionGroup( + generator: self.idGenerator, + minimumConcurrentConnections: 0, + maximumConcurrentConnectionSoftLimit: 2, + maximumConcurrentConnectionHardLimit: 2, + keepAlive: true, + keepAliveReducesAvailableStreams: true + ) + + guard let firstRequest = connections.createNewDemandConnectionIfPossible() else { return XCTFail("Expected to have a request here") } + + let newConnection = MockConnection(id: firstRequest.connectionID) + let (connectionIndex, establishedConnectionContext) = connections.newConnectionEstablished(newConnection, maxStreams: 1) + XCTAssertEqual(establishedConnectionContext.info, .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(connections.stats, .init(idle: 1, availableStreams: 1)) + _ = connections.parkConnection(at: connectionIndex, hasBecomeIdle: true) + let keepAliveTimer = TestPoolStateMachine.ConnectionTimer(timerID: 0, connectionID: firstRequest.connectionID, usecase: .keepAlive) + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + XCTAssertNil(connections.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + let keepAliveAction = connections.keepAliveIfIdle(newConnection.id) + XCTAssertEqual(keepAliveAction, .init(connection: newConnection, keepAliveTimerCancellationContinuation: keepAliveTimerCancellationToken)) + XCTAssertEqual(connections.stats, .init(idle: 1, runningKeepAlive: 1, availableStreams: 0)) + + _ = connections.closeConnectionIfIdle(newConnection.id) + guard connections.keepAliveFailed(newConnection.id) == nil else { + return XCTFail("Expected keepAliveFailed not to cause close again") + } + XCTAssertEqual(connections.stats, .init(closing: 1)) + } +} diff --git a/Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionStateTests.swift b/Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionStateTests.swift new file mode 100644 index 00000000..7dd2b726 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionStateTests.swift @@ -0,0 +1,265 @@ +@testable import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import XCTest + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class PoolStateMachine_ConnectionStateTests: XCTestCase { + + typealias TestConnectionState = TestPoolStateMachine.ConnectionState + + func testStartupLeaseReleaseParkLease() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + XCTAssertEqual(state.id, connectionID) + XCTAssertEqual(state.isIdle, false) + XCTAssertEqual(state.isAvailable, false) + XCTAssertEqual(state.isConnected, false) + XCTAssertEqual(state.isLeased, false) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 1), .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual(state.isIdle, true) + XCTAssertEqual(state.isAvailable, true) + XCTAssertEqual(state.isConnected, true) + XCTAssertEqual(state.isLeased, false) + XCTAssertEqual(state.lease(streams: 1), .init(connection: connection, timersToCancel: .init(), wasIdle: true)) + + XCTAssertEqual(state.isIdle, false) + XCTAssertEqual(state.isAvailable, false) + XCTAssertEqual(state.isConnected, true) + XCTAssertEqual(state.isLeased, true) + + XCTAssertEqual(state.release(streams: 1), .idle(availableStreams: 1, newIdle: true)) + let parkResult = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: true) + XCTAssert( + parkResult.elementsEqual([ + .init(timerID: 0, connectionID: connectionID, usecase: .keepAlive), + .init(timerID: 1, connectionID: connectionID, usecase: .idleTimeout) + ]) + ) + + guard let keepAliveTimer = parkResult.first, let idleTimer = parkResult.second else { + return XCTFail("Expected to get two timers") + } + + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + let idleTimerCancellationToken = MockTimerCancellationToken(idleTimer) + + XCTAssertNil(state.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + XCTAssertNil(state.timerScheduled(idleTimer, cancelContinuation: idleTimerCancellationToken)) + + let expectLeaseAction = TestConnectionState.LeaseAction( + connection: connection, + timersToCancel: [idleTimerCancellationToken, keepAliveTimerCancellationToken], + wasIdle: true + ) + XCTAssertEqual(state.lease(streams: 1), expectLeaseAction) + } + + func testStartupParkLeaseBeforeTimersRegistered() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 1), .idle(availableStreams: 1, newIdle: true)) + let parkResult = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: true) + XCTAssertEqual( + parkResult, + [ + .init(timerID: 0, connectionID: connectionID, usecase: .keepAlive), + .init(timerID: 1, connectionID: connectionID, usecase: .idleTimeout) + ] + ) + + guard let keepAliveTimer = parkResult.first, let idleTimer = parkResult.second else { + return XCTFail("Expected to get two timers") + } + + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + let idleTimerCancellationToken = MockTimerCancellationToken(idleTimer) + XCTAssertEqual(state.lease(streams: 1), .init(connection: connection, timersToCancel: .init(), wasIdle: true)) + + XCTAssertEqual(state.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken), keepAliveTimerCancellationToken) + XCTAssertEqual(state.timerScheduled(idleTimer, cancelContinuation: idleTimerCancellationToken), idleTimerCancellationToken) + } + + func testStartupParkLeasePark() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 1), .idle(availableStreams: 1, newIdle: true)) + let parkResult = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: true) + XCTAssert( + parkResult.elementsEqual([ + .init(timerID: 0, connectionID: connectionID, usecase: .keepAlive), + .init(timerID: 1, connectionID: connectionID, usecase: .idleTimeout) + ]) + ) + + guard let keepAliveTimer = parkResult.first, let idleTimer = parkResult.second else { + return XCTFail("Expected to get two timers") + } + + let initialKeepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + let initialIdleTimerCancellationToken = MockTimerCancellationToken(idleTimer) + XCTAssertEqual(state.lease(streams: 1), .init(connection: connection, timersToCancel: .init(), wasIdle: true)) + + XCTAssertEqual(state.release(streams: 1), .idle(availableStreams: 1, newIdle: true)) + XCTAssertEqual( + state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: true), + [ + .init(timerID: 2, connectionID: connectionID, usecase: .keepAlive), + .init(timerID: 3, connectionID: connectionID, usecase: .idleTimeout) + ] + ) + + XCTAssertEqual(state.timerScheduled(keepAliveTimer, cancelContinuation: initialKeepAliveTimerCancellationToken), initialKeepAliveTimerCancellationToken) + XCTAssertEqual(state.timerScheduled(idleTimer, cancelContinuation: initialIdleTimerCancellationToken), initialIdleTimerCancellationToken) + } + + func testStartupFailed() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let firstBackoffTimer = state.failedToConnect() + let firstBackoffTimerCancellationToken = MockTimerCancellationToken(firstBackoffTimer) + XCTAssertNil(state.timerScheduled(firstBackoffTimer, cancelContinuation: firstBackoffTimerCancellationToken)) + XCTAssertEqual(state.retryConnect(), firstBackoffTimerCancellationToken) + + let secondBackoffTimer = state.failedToConnect() + let secondBackoffTimerCancellationToken = MockTimerCancellationToken(secondBackoffTimer) + XCTAssertNil(state.retryConnect()) + XCTAssertEqual( + state.timerScheduled(secondBackoffTimer, cancelContinuation: secondBackoffTimerCancellationToken), + secondBackoffTimerCancellationToken + ) + + let thirdBackoffTimer = state.failedToConnect() + let thirdBackoffTimerCancellationToken = MockTimerCancellationToken(thirdBackoffTimer) + XCTAssertNil(state.retryConnect()) + let forthBackoffTimer = state.failedToConnect() + let forthBackoffTimerCancellationToken = MockTimerCancellationToken(forthBackoffTimer) + XCTAssertEqual( + state.timerScheduled(thirdBackoffTimer, cancelContinuation: thirdBackoffTimerCancellationToken), + thirdBackoffTimerCancellationToken + ) + XCTAssertNil( + state.timerScheduled(forthBackoffTimer, cancelContinuation: forthBackoffTimerCancellationToken) + ) + XCTAssertEqual(state.retryConnect(), forthBackoffTimerCancellationToken) + + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 1), .idle(availableStreams: 1, newIdle: true)) + } + + func testLeaseMultipleStreams() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 100), .idle(availableStreams: 100, newIdle: true)) + let timers = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: false) + guard let keepAliveTimer = timers.first else { return XCTFail("Expected to get a keepAliveTimer") } + + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + XCTAssertNil(state.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + + XCTAssertEqual( + state.lease(streams: 30), + TestConnectionState.LeaseAction(connection: connection, timersToCancel: [keepAliveTimerCancellationToken], wasIdle: true) + ) + + XCTAssertEqual(state.release(streams: 10), .leased(availableStreams: 80)) + + XCTAssertEqual( + state.lease(streams: 40), + TestConnectionState.LeaseAction(connection: connection, timersToCancel: [], wasIdle: false) + ) + + XCTAssertEqual( + state.lease(streams: 40), + TestConnectionState.LeaseAction(connection: connection, timersToCancel: [], wasIdle: false) + ) + + XCTAssertEqual(state.release(streams: 1), .leased(availableStreams: 1)) + XCTAssertEqual(state.release(streams: 98), .leased(availableStreams: 99)) + XCTAssertEqual(state.release(streams: 1), .idle(availableStreams: 100, newIdle: true)) + } + + func testRunningKeepAliveReducesAvailableStreams() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 100), .idle(availableStreams: 100, newIdle: true)) + let timers = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: false) + guard let keepAliveTimer = timers.first else { return XCTFail("Expected to get a keepAliveTimer") } + + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + XCTAssertNil(state.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + + XCTAssertEqual( + state.runKeepAliveIfIdle(reducesAvailableStreams: true), + .init(connection: connection, keepAliveTimerCancellationContinuation: keepAliveTimerCancellationToken) + ) + + XCTAssertEqual( + state.lease(streams: 30), + TestConnectionState.LeaseAction(connection: connection, timersToCancel: [], wasIdle: true) + ) + + XCTAssertEqual(state.release(streams: 10), .leased(availableStreams: 79)) + XCTAssertEqual(state.isAvailable, true) + XCTAssertEqual( + state.lease(streams: 79), + TestConnectionState.LeaseAction(connection: connection, timersToCancel: [], wasIdle: false) + ) + XCTAssertEqual(state.isAvailable, false) + XCTAssertEqual(state.keepAliveSucceeded(), .leased(availableStreams: 1)) + XCTAssertEqual(state.isAvailable, true) + } + + func testRunningKeepAliveDoesNotReduceAvailableStreams() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 100), .idle(availableStreams: 100, newIdle: true)) + let timers = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: false) + guard let keepAliveTimer = timers.first else { return XCTFail("Expected to get a keepAliveTimer") } + + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + XCTAssertNil(state.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + + XCTAssertEqual( + state.runKeepAliveIfIdle(reducesAvailableStreams: false), + .init(connection: connection, keepAliveTimerCancellationContinuation: keepAliveTimerCancellationToken) + ) + + XCTAssertEqual( + state.lease(streams: 30), + TestConnectionState.LeaseAction(connection: connection, timersToCancel: [], wasIdle: true) + ) + + XCTAssertEqual(state.release(streams: 10), .leased(availableStreams: 80)) + XCTAssertEqual(state.keepAliveSucceeded(), .leased(availableStreams: 80)) + } + + func testRunKeepAliveRacesAgainstIdleClose() { + let connectionID = 1 + var state = TestConnectionState(id: connectionID) + let connection = MockConnection(id: connectionID) + XCTAssertEqual(state.connected(connection, maxStreams: 1), .idle(availableStreams: 1, newIdle: true)) + let parkResult = state.parkConnection(scheduleKeepAliveTimer: true, scheduleIdleTimeoutTimer: true) + guard let keepAliveTimer = parkResult.first, let idleTimer = parkResult.second else { + return XCTFail("Expected to get two timers") + } + + XCTAssertEqual(keepAliveTimer, .init(timerID: 0, connectionID: connectionID, usecase: .keepAlive)) + XCTAssertEqual(idleTimer, .init(timerID: 1, connectionID: connectionID, usecase: .idleTimeout)) + + let keepAliveTimerCancellationToken = MockTimerCancellationToken(keepAliveTimer) + let idleTimerCancellationToken = MockTimerCancellationToken(idleTimer) + + XCTAssertNil(state.timerScheduled(keepAliveTimer, cancelContinuation: keepAliveTimerCancellationToken)) + XCTAssertNil(state.timerScheduled(idleTimer, cancelContinuation: idleTimerCancellationToken)) + + XCTAssertEqual(state.closeIfIdle(), .init(connection: connection, previousConnectionState: .idle, cancelTimers: [keepAliveTimerCancellationToken, idleTimerCancellationToken], usedStreams: 0, maxStreams: 1, runningKeepAlive: false)) + XCTAssertEqual(state.runKeepAliveIfIdle(reducesAvailableStreams: true), .none) + + } +} diff --git a/Tests/ConnectionPoolModuleTests/PoolStateMachine+RequestQueueTests.swift b/Tests/ConnectionPoolModuleTests/PoolStateMachine+RequestQueueTests.swift new file mode 100644 index 00000000..b74b86cc --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/PoolStateMachine+RequestQueueTests.swift @@ -0,0 +1,148 @@ +@testable import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import XCTest + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class PoolStateMachine_RequestQueueTests: XCTestCase { + + typealias TestQueue = TestPoolStateMachine.RequestQueue + + func testHappyPath() { + var queue = TestQueue() + XCTAssert(queue.isEmpty) + + let request1 = MockRequest() + queue.queue(request1) + XCTAssertEqual(queue.count, 1) + XCTAssertFalse(queue.isEmpty) + let popResult = queue.pop(max: 3) + XCTAssert(popResult.elementsEqual([request1])) + XCTAssert(queue.isEmpty) + XCTAssertEqual(queue.count, 0) + } + + func testEnqueueAndPopMultipleRequests() { + var queue = TestQueue() + XCTAssert(queue.isEmpty) + + var request1 = MockRequest() + queue.queue(request1) + var request2 = MockRequest() + queue.queue(request2) + var request3 = MockRequest() + queue.queue(request3) + + do { + XCTAssertEqual(queue.count, 3) + XCTAssertFalse(queue.isEmpty) + let popResult = queue.pop(max: 3) + XCTAssert(popResult.elementsEqual([request1, request2, request3])) + XCTAssert(queue.isEmpty) + XCTAssertEqual(queue.count, 0) + } + XCTAssert(isKnownUniquelyReferenced(&request1)) + XCTAssert(isKnownUniquelyReferenced(&request2)) + XCTAssert(isKnownUniquelyReferenced(&request3)) + } + + func testEnqueueAndPopOnlyOne() { + var queue = TestQueue() + XCTAssert(queue.isEmpty) + + var request1 = MockRequest() + queue.queue(request1) + var request2 = MockRequest() + queue.queue(request2) + var request3 = MockRequest() + queue.queue(request3) + + do { + XCTAssertEqual(queue.count, 3) + XCTAssertFalse(queue.isEmpty) + let popResult = queue.pop(max: 1) + XCTAssert(popResult.elementsEqual([request1])) + XCTAssertFalse(queue.isEmpty) + XCTAssertEqual(queue.count, 2) + + let removeAllResult = queue.removeAll() + XCTAssert(Set(removeAllResult) == [request2, request3]) + } + XCTAssert(isKnownUniquelyReferenced(&request1)) + XCTAssert(isKnownUniquelyReferenced(&request2)) + XCTAssert(isKnownUniquelyReferenced(&request3)) + } + + func testCancellation() { + var queue = TestQueue() + XCTAssert(queue.isEmpty) + + var request1 = MockRequest() + queue.queue(request1) + var request2 = MockRequest() + queue.queue(request2) + var request3 = MockRequest() + queue.queue(request3) + + do { + XCTAssertEqual(queue.count, 3) + let returnedRequest2 = queue.remove(request2.id) + XCTAssert(returnedRequest2 === request2) + XCTAssertEqual(queue.count, 2) + XCTAssertFalse(queue.isEmpty) + } + + // still retained by the deque inside the queue + XCTAssertEqual(queue.requests.count, 2) + XCTAssertEqual(queue.queue.count, 3) + + do { + XCTAssertEqual(queue.count, 2) + XCTAssertFalse(queue.isEmpty) + let popResult = queue.pop(max: 3) + XCTAssert(popResult.elementsEqual([request1, request3])) + XCTAssert(queue.isEmpty) + XCTAssertEqual(queue.count, 0) + } + + XCTAssert(isKnownUniquelyReferenced(&request1)) + XCTAssert(isKnownUniquelyReferenced(&request2)) + XCTAssert(isKnownUniquelyReferenced(&request3)) + } + + func testRemoveAllAfterCancellation() { + var queue = TestQueue() + XCTAssert(queue.isEmpty) + + var request1 = MockRequest() + queue.queue(request1) + var request2 = MockRequest() + queue.queue(request2) + var request3 = MockRequest() + queue.queue(request3) + + do { + XCTAssertEqual(queue.count, 3) + let returnedRequest2 = queue.remove(request2.id) + XCTAssert(returnedRequest2 === request2) + XCTAssertEqual(queue.count, 2) + XCTAssertFalse(queue.isEmpty) + } + + // still retained by the deque inside the queue + XCTAssertEqual(queue.requests.count, 2) + XCTAssertEqual(queue.queue.count, 3) + + do { + XCTAssertEqual(queue.count, 2) + XCTAssertFalse(queue.isEmpty) + let removeAllResult = queue.removeAll() + XCTAssert(Set(removeAllResult) == [request1, request3]) + XCTAssert(queue.isEmpty) + XCTAssertEqual(queue.count, 0) + } + + XCTAssert(isKnownUniquelyReferenced(&request1)) + XCTAssert(isKnownUniquelyReferenced(&request2)) + XCTAssert(isKnownUniquelyReferenced(&request3)) + } +} diff --git a/Tests/ConnectionPoolModuleTests/PoolStateMachineTests.swift b/Tests/ConnectionPoolModuleTests/PoolStateMachineTests.swift new file mode 100644 index 00000000..c0b6ddcd --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/PoolStateMachineTests.swift @@ -0,0 +1,385 @@ +@testable import _ConnectionPoolModule +import _ConnectionPoolTestUtils +import XCTest + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +typealias TestPoolStateMachine = PoolStateMachine< + MockConnection, + ConnectionIDGenerator, + MockConnection.ID, + MockRequest, + MockRequest.ID, + MockTimerCancellationToken +> + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class PoolStateMachineTests: XCTestCase { + + func testConnectionsAreCreatedAndParkedOnStartup() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 2 + configuration.maximumConnectionSoftLimit = 4 + configuration.maximumConnectionHardLimit = 6 + configuration.keepAliveDuration = .seconds(10) + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + let connection1 = MockConnection(id: 0) + let connection2 = MockConnection(id: 1) + + do { + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 2) + let createdAction1 = stateMachine.connectionEstablished(connection1, maxStreams: 1) + let connection1KeepAliveTimer = TestPoolStateMachine.Timer(.init(timerID: 0, connectionID: 0, usecase: .keepAlive), duration: .seconds(10)) + let connection1KeepAliveTimerCancellationToken = MockTimerCancellationToken(connection1KeepAliveTimer) + XCTAssertEqual(createdAction1.request, .none) + XCTAssertEqual(createdAction1.connection, .scheduleTimers([connection1KeepAliveTimer])) + + XCTAssertEqual(stateMachine.timerScheduled(connection1KeepAliveTimer, cancelContinuation: connection1KeepAliveTimerCancellationToken), .none) + + let createdAction2 = stateMachine.connectionEstablished(connection2, maxStreams: 1) + let connection2KeepAliveTimer = TestPoolStateMachine.Timer(.init(timerID: 0, connectionID: 1, usecase: .keepAlive), duration: .seconds(10)) + let connection2KeepAliveTimerCancellationToken = MockTimerCancellationToken(connection2KeepAliveTimer) + XCTAssertEqual(createdAction2.request, .none) + XCTAssertEqual(createdAction2.connection, .scheduleTimers([connection2KeepAliveTimer])) + XCTAssertEqual(stateMachine.timerScheduled(connection2KeepAliveTimer, cancelContinuation: connection2KeepAliveTimerCancellationToken), .none) + } + } + + func testConnectionsNoKeepAliveRun() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 1 + configuration.maximumConnectionSoftLimit = 4 + configuration.maximumConnectionHardLimit = 6 + configuration.keepAliveDuration = nil + configuration.idleTimeoutDuration = .seconds(5) + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + let connection1 = MockConnection(id: 0) + + // refill pool to at least one connection + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 1) + let createdAction1 = stateMachine.connectionEstablished(connection1, maxStreams: 1) + XCTAssertEqual(createdAction1.request, .none) + XCTAssertEqual(createdAction1.connection, .scheduleTimers([])) + + // lease connection 1 + let request1 = MockRequest() + let leaseRequest1 = stateMachine.leaseConnection(request1) + XCTAssertEqual(leaseRequest1.connection, .cancelTimers([])) + XCTAssertEqual(leaseRequest1.request, .leaseConnection(.init(element: request1), connection1)) + + // release connection 1 + XCTAssertEqual(stateMachine.releaseConnection(connection1, streams: 1), .none()) + + // lease connection 1 + let request2 = MockRequest() + let leaseRequest2 = stateMachine.leaseConnection(request2) + XCTAssertEqual(leaseRequest2.connection, .cancelTimers([])) + XCTAssertEqual(leaseRequest2.request, .leaseConnection(.init(element: request2), connection1)) + + // request connection while none is available + let request3 = MockRequest() + let leaseRequest3 = stateMachine.leaseConnection(request3) + XCTAssertEqual(leaseRequest3.connection, .makeConnection(.init(connectionID: 1), [])) + XCTAssertEqual(leaseRequest3.request, .none) + + // make connection 2 and lease immediately + let connection2 = MockConnection(id: 1) + let createdAction2 = stateMachine.connectionEstablished(connection2, maxStreams: 1) + XCTAssertEqual(createdAction2.request, .leaseConnection(.init(element: request3), connection2)) + XCTAssertEqual(createdAction2.connection, .none) + + // release connection 2 + let connection2IdleTimer = TestPoolStateMachine.Timer(.init(timerID: 0, connectionID: 1, usecase: .idleTimeout), duration: configuration.idleTimeoutDuration) + let connection2IdleTimerCancellationToken = MockTimerCancellationToken(connection2IdleTimer) + XCTAssertEqual( + stateMachine.releaseConnection(connection2, streams: 1), + .init(request: .none, connection: .scheduleTimers([connection2IdleTimer])) + ) + + XCTAssertEqual(stateMachine.timerScheduled(connection2IdleTimer, cancelContinuation: connection2IdleTimerCancellationToken), .none) + XCTAssertEqual(stateMachine.timerTriggered(connection2IdleTimer), .init(request: .none, connection: .closeConnection(connection2, [connection2IdleTimerCancellationToken]))) + } + + func testOnlyOverflowConnections() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 0 + configuration.maximumConnectionSoftLimit = 0 + configuration.maximumConnectionHardLimit = 6 + configuration.keepAliveDuration = nil + configuration.idleTimeoutDuration = .seconds(3) + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + // don't refill pool + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 0) + + // request connection while none exists + let request1 = MockRequest() + let leaseRequest1 = stateMachine.leaseConnection(request1) + XCTAssertEqual(leaseRequest1.connection, .makeConnection(.init(connectionID: 0), [])) + XCTAssertEqual(leaseRequest1.request, .none) + + // make connection 1 and lease immediately + let connection1 = MockConnection(id: 0) + let createdAction1 = stateMachine.connectionEstablished(connection1, maxStreams: 1) + XCTAssertEqual(createdAction1.request, .leaseConnection(.init(element: request1), connection1)) + XCTAssertEqual(createdAction1.connection, .none) + + // request connection while none is available + let request2 = MockRequest() + let leaseRequest2 = stateMachine.leaseConnection(request2) + XCTAssertEqual(leaseRequest2.connection, .makeConnection(.init(connectionID: 1), [])) + XCTAssertEqual(leaseRequest2.request, .none) + + // release connection 1 should be leased again immediately + let releaseRequest1 = stateMachine.releaseConnection(connection1, streams: 1) + XCTAssertEqual(releaseRequest1.request, .leaseConnection(.init(element: request2), connection1)) + XCTAssertEqual(releaseRequest1.connection, .none) + + // connection 2 comes up and should be closed right away + let connection2 = MockConnection(id: 1) + let createdAction2 = stateMachine.connectionEstablished(connection2, maxStreams: 1) + XCTAssertEqual(createdAction2.request, .none) + XCTAssertEqual(createdAction2.connection, .closeConnection(connection2, [])) + XCTAssertEqual(stateMachine.connectionClosed(connection2), .none()) + + // release connection 1 should be closed as well + let releaseRequest2 = stateMachine.releaseConnection(connection1, streams: 1) + XCTAssertEqual(releaseRequest2.request, .none) + XCTAssertEqual(releaseRequest2.connection, .closeConnection(connection1, [])) + + let shutdownAction = stateMachine.triggerForceShutdown() + XCTAssertEqual(shutdownAction.request, .failRequests(.init(), .poolShutdown)) + XCTAssertEqual(shutdownAction.connection, .shutdown(.init())) + } + + func testDemandConnectionIsMadePermanentIfPermanentIsClose() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 1 + configuration.maximumConnectionSoftLimit = 2 + configuration.maximumConnectionHardLimit = 6 + configuration.keepAliveDuration = nil + configuration.idleTimeoutDuration = .seconds(3) + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + let connection1 = MockConnection(id: 0) + + // refill pool to at least one connection + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 1) + let createdAction1 = stateMachine.connectionEstablished(connection1, maxStreams: 1) + XCTAssertEqual(createdAction1.request, .none) + XCTAssertEqual(createdAction1.connection, .scheduleTimers([])) + + // lease connection 1 + let request1 = MockRequest() + let leaseRequest1 = stateMachine.leaseConnection(request1) + XCTAssertEqual(leaseRequest1.connection, .cancelTimers([])) + XCTAssertEqual(leaseRequest1.request, .leaseConnection(.init(element: request1), connection1)) + + // request connection while none is available + let request2 = MockRequest() + let leaseRequest2 = stateMachine.leaseConnection(request2) + XCTAssertEqual(leaseRequest2.connection, .makeConnection(.init(connectionID: 1), [])) + XCTAssertEqual(leaseRequest2.request, .none) + + // make connection 2 and lease immediately + let connection2 = MockConnection(id: 1) + let createdAction2 = stateMachine.connectionEstablished(connection2, maxStreams: 1) + XCTAssertEqual(createdAction2.request, .leaseConnection(.init(element: request2), connection2)) + XCTAssertEqual(createdAction2.connection, .none) + + // release connection 2 + let connection2IdleTimer = TestPoolStateMachine.Timer(.init(timerID: 0, connectionID: 1, usecase: .idleTimeout), duration: configuration.idleTimeoutDuration) + let connection2IdleTimerCancellationToken = MockTimerCancellationToken(connection2IdleTimer) + XCTAssertEqual( + stateMachine.releaseConnection(connection2, streams: 1), + .init(request: .none, connection: .scheduleTimers([connection2IdleTimer])) + ) + + XCTAssertEqual(stateMachine.timerScheduled(connection2IdleTimer, cancelContinuation: connection2IdleTimerCancellationToken), .none) + + // connection 1 is dropped + XCTAssertEqual(stateMachine.connectionClosed(connection1), .init(request: .none, connection: .cancelTimers([connection2IdleTimerCancellationToken]))) + } + + func testReleaseLoosesRaceAgainstClosed() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 0 + configuration.maximumConnectionSoftLimit = 2 + configuration.maximumConnectionHardLimit = 2 + configuration.keepAliveDuration = nil + configuration.idleTimeoutDuration = .seconds(3) + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + // don't refill pool + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 0) + + // request connection while none exists + let request1 = MockRequest() + let leaseRequest1 = stateMachine.leaseConnection(request1) + XCTAssertEqual(leaseRequest1.connection, .makeConnection(.init(connectionID: 0), [])) + XCTAssertEqual(leaseRequest1.request, .none) + + // make connection 1 and lease immediately + let connection1 = MockConnection(id: 0) + let createdAction1 = stateMachine.connectionEstablished(connection1, maxStreams: 1) + XCTAssertEqual(createdAction1.request, .leaseConnection(.init(element: request1), connection1)) + XCTAssertEqual(createdAction1.connection, .none) + + // connection got closed + let closedAction = stateMachine.connectionClosed(connection1) + XCTAssertEqual(closedAction.connection, .none) + XCTAssertEqual(closedAction.request, .none) + + // release connection 1 should be leased again immediately + let releaseRequest1 = stateMachine.releaseConnection(connection1, streams: 1) + XCTAssertEqual(releaseRequest1.request, .none) + XCTAssertEqual(releaseRequest1.connection, .none) + } + + func testKeepAliveOnClosingConnection() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 0 + configuration.maximumConnectionSoftLimit = 2 + configuration.maximumConnectionHardLimit = 2 + configuration.keepAliveDuration = .seconds(2) + configuration.idleTimeoutDuration = .seconds(4) + + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + // don't refill pool + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 0) + + // request connection while none exists + let request1 = MockRequest() + let leaseRequest1 = stateMachine.leaseConnection(request1) + XCTAssertEqual(leaseRequest1.connection, .makeConnection(.init(connectionID: 0), [])) + XCTAssertEqual(leaseRequest1.request, .none) + + // make connection 1 + let connection1 = MockConnection(id: 0) + let createdAction1 = stateMachine.connectionEstablished(connection1, maxStreams: 1) + XCTAssertEqual(createdAction1.request, .leaseConnection(.init(element: request1), connection1)) + XCTAssertEqual(createdAction1.connection, .none) + _ = stateMachine.releaseConnection(connection1, streams: 1) + + // trigger keep alive + let keepAliveAction1 = stateMachine.connectionKeepAliveTimerTriggered(connection1.id) + XCTAssertEqual(keepAliveAction1.connection, .runKeepAlive(connection1, nil)) + + // fail keep alive and cause closed + let keepAliveFailed1 = stateMachine.connectionKeepAliveFailed(connection1.id) + XCTAssertEqual(keepAliveFailed1.connection, .closeConnection(connection1, [])) + connection1.closeIfClosing() + + // request connection while none exists anymore + let request2 = MockRequest() + let leaseRequest2 = stateMachine.leaseConnection(request2) + XCTAssertEqual(leaseRequest2.connection, .makeConnection(.init(connectionID: 1), [])) + XCTAssertEqual(leaseRequest2.request, .none) + + // make connection 2 + let connection2 = MockConnection(id: 1) + let createdAction2 = stateMachine.connectionEstablished(connection2, maxStreams: 1) + XCTAssertEqual(createdAction2.request, .leaseConnection(.init(element: request2), connection2)) + XCTAssertEqual(createdAction2.connection, .none) + _ = stateMachine.releaseConnection(connection2, streams: 1) + + // trigger keep alive while connection is still open + let keepAliveAction2 = stateMachine.connectionKeepAliveTimerTriggered(connection2.id) + XCTAssertEqual(keepAliveAction2.connection, .runKeepAlive(connection2, nil)) + + // close connection in the middle of keep alive + connection2.close() + connection2.closeIfClosing() + + // fail keep alive and cause closed + let keepAliveFailed2 = stateMachine.connectionKeepAliveFailed(connection2.id) + XCTAssertEqual(keepAliveFailed2.connection, .closeConnection(connection2, [])) + } + + func testConnectionIsEstablishedAfterFailedKeepAliveIfNotEnoughConnectionsLeft() { + var configuration = PoolConfiguration() + configuration.minimumConnectionCount = 1 + configuration.maximumConnectionSoftLimit = 2 + configuration.maximumConnectionHardLimit = 2 + configuration.keepAliveDuration = .seconds(2) + configuration.idleTimeoutDuration = .seconds(4) + + + var stateMachine = TestPoolStateMachine( + configuration: configuration, + generator: .init(), + timerCancellationTokenType: MockTimerCancellationToken.self + ) + + // refill pool + let requests = stateMachine.refillConnections() + XCTAssertEqual(requests.count, 1) + + // one connection should exist + let request = MockRequest() + let leaseRequest = stateMachine.leaseConnection(request) + XCTAssertEqual(leaseRequest.connection, .none) + XCTAssertEqual(leaseRequest.request, .none) + + // make connection 1 + let connection = MockConnection(id: 0) + let createdAction = stateMachine.connectionEstablished(connection, maxStreams: 1) + XCTAssertEqual(createdAction.request, .leaseConnection(.init(element: request), connection)) + XCTAssertEqual(createdAction.connection, .none) + _ = stateMachine.releaseConnection(connection, streams: 1) + + // trigger keep alive + let keepAliveAction = stateMachine.connectionKeepAliveTimerTriggered(connection.id) + XCTAssertEqual(keepAliveAction.connection, .runKeepAlive(connection, nil)) + + // fail keep alive, cause closed and make new connection + let keepAliveFailed = stateMachine.connectionKeepAliveFailed(connection.id) + XCTAssertEqual(keepAliveFailed.connection, .closeConnection(connection, [])) + let connectionClosed = stateMachine.connectionClosed(connection) + XCTAssertEqual(connectionClosed.connection, .makeConnection(.init(connectionID: 1), [])) + connection.closeIfClosing() + let establishAction = stateMachine.connectionEstablished(.init(id: 1), maxStreams: 1) + XCTAssertEqual(establishAction.request, .none) + guard case .scheduleTimers(let timers) = establishAction.connection else { return XCTFail("Unexpected connection action") } + XCTAssertEqual(timers, [.init(.init(timerID: 0, connectionID: 1, usecase: .keepAlive), duration: configuration.keepAliveDuration!)]) + } + +} diff --git a/Tests/ConnectionPoolModuleTests/TinyFastSequenceTests.swift b/Tests/ConnectionPoolModuleTests/TinyFastSequenceTests.swift new file mode 100644 index 00000000..1a2836b9 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/TinyFastSequenceTests.swift @@ -0,0 +1,72 @@ +@testable import _ConnectionPoolModule +import XCTest + +final class TinyFastSequenceTests: XCTestCase { + func testCountIsEmptyAndIterator() async { + var sequence = TinyFastSequence() + XCTAssertEqual(sequence.count, 0) + XCTAssertEqual(sequence.isEmpty, true) + XCTAssertEqual(sequence.first, nil) + XCTAssertEqual(Array(sequence), []) + sequence.append(1) + XCTAssertEqual(sequence.count, 1) + XCTAssertEqual(sequence.isEmpty, false) + XCTAssertEqual(sequence.first, 1) + XCTAssertEqual(Array(sequence), [1]) + sequence.append(2) + XCTAssertEqual(sequence.count, 2) + XCTAssertEqual(sequence.isEmpty, false) + XCTAssertEqual(sequence.first, 1) + XCTAssertEqual(Array(sequence), [1, 2]) + sequence.append(3) + XCTAssertEqual(sequence.count, 3) + XCTAssertEqual(sequence.isEmpty, false) + XCTAssertEqual(sequence.first, 1) + XCTAssertEqual(Array(sequence), [1, 2, 3]) + } + + func testReserveCapacityIsForwarded() { + var emptySequence = TinyFastSequence() + emptySequence.reserveCapacity(8) + emptySequence.append(1) + emptySequence.append(2) + emptySequence.append(3) + guard case .n(let array) = emptySequence.base else { + return XCTFail("Expected sequence to be backed by an array") + } + XCTAssertEqual(array.capacity, 8) + + var oneElemSequence = TinyFastSequence(element: 1) + oneElemSequence.reserveCapacity(8) + oneElemSequence.append(2) + oneElemSequence.append(3) + guard case .n(let array) = oneElemSequence.base else { + return XCTFail("Expected sequence to be backed by an array") + } + XCTAssertEqual(array.capacity, 8) + + var twoElemSequence = TinyFastSequence([1, 2]) + twoElemSequence.reserveCapacity(8) + guard case .n(let array) = twoElemSequence.base else { + return XCTFail("Expected sequence to be backed by an array") + } + XCTAssertEqual(array.capacity, 8) + } + + func testNewSequenceSlowPath() { + let sequence = TinyFastSequence("AB".utf8) + XCTAssertEqual(Array(sequence), [UInt8(ascii: "A"), UInt8(ascii: "B")]) + } + + func testSingleItem() { + let sequence = TinyFastSequence("A".utf8) + XCTAssertEqual(Array(sequence), [UInt8(ascii: "A")]) + } + + func testEmptyCollection() { + let sequence = TinyFastSequence("".utf8) + XCTAssertTrue(sequence.isEmpty) + XCTAssertEqual(sequence.count, 0) + XCTAssertEqual(Array(sequence), []) + } +} diff --git a/Tests/ConnectionPoolModuleTests/Utils/Future.swift b/Tests/ConnectionPoolModuleTests/Utils/Future.swift new file mode 100644 index 00000000..2bee3216 --- /dev/null +++ b/Tests/ConnectionPoolModuleTests/Utils/Future.swift @@ -0,0 +1,112 @@ +import Atomics +@testable import _ConnectionPoolModule + +/// This is a `Future` type that shall make writing tests a bit simpler. I'm well aware, that this is a pattern +/// that should not be embraced with structured concurrency. However writing all tests in full structured +/// concurrency is an effort, that isn't worth the endgoals in my view. +final class Future: Sendable { + struct State: Sendable { + + var result: Swift.Result? = nil + var continuations: [(Int, CheckedContinuation)] = [] + + } + + let waiterID = ManagedAtomic(0) + let stateBox: NIOLockedValueBox = NIOLockedValueBox(State()) + + init(of: Success.Type) {} + + enum GetAction { + case fail(any Error) + case succeed(Success) + case none + } + + var success: Success { + get async throws { + let waiterID = self.waiterID.loadThenWrappingIncrement(ordering: .relaxed) + + return try await withTaskCancellationHandler { + return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + let action = self.stateBox.withLockedValue { state -> GetAction in + if Task.isCancelled { + return .fail(CancellationError()) + } + + switch state.result { + case .none: + state.continuations.append((waiterID, continuation)) + return .none + + case .success(let result): + return .succeed(result) + + case .failure(let error): + return .fail(error) + } + } + + switch action { + case .fail(let error): + continuation.resume(throwing: error) + + case .succeed(let result): + continuation.resume(returning: result) + + case .none: + break + } + } + } onCancel: { + let cont = self.stateBox.withLockedValue { state -> CheckedContinuation? in + guard state.result == nil else { return nil } + + guard let contIndex = state.continuations.firstIndex(where: { $0.0 == waiterID }) else { + return nil + } + let (_, continuation) = state.continuations.remove(at: contIndex) + return continuation + } + + cont?.resume(throwing: CancellationError()) + } + } + } + + func yield(value: Success) { + let continuations = self.stateBox.withLockedValue { state in + guard state.result == nil else { + return [(Int, CheckedContinuation)]().lazy.map(\.1) + } + state.result = .success(value) + + let continuations = state.continuations + state.continuations = [] + + return continuations.lazy.map(\.1) + } + + for continuation in continuations { + continuation.resume(returning: value) + } + } + + func yield(error: any Error) { + let continuations = self.stateBox.withLockedValue { state in + guard state.result == nil else { + return [(Int, CheckedContinuation)]().lazy.map(\.1) + } + state.result = .failure(error) + + let continuations = state.continuations + state.continuations = [] + + return continuations.lazy.map(\.1) + } + + for continuation in continuations { + continuation.resume(throwing: error) + } + } +} diff --git a/Tests/IntegrationTests/AsyncTests.swift b/Tests/IntegrationTests/AsyncTests.swift index ed6910d1..b4c8e93f 100644 --- a/Tests/IntegrationTests/AsyncTests.swift +++ b/Tests/IntegrationTests/AsyncTests.swift @@ -8,7 +8,6 @@ import NIOPosix import NIOCore final class AsyncPostgresConnectionTests: XCTestCase { - func test1kRoundTrips() async throws { let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } @@ -37,7 +36,8 @@ final class AsyncPostgresConnectionTests: XCTestCase { try await withTestConnection(on: eventLoop) { connection in let rows = try await connection.query("SELECT generate_series(\(start), \(end));", logger: .psqlTest) var counter = 0 - for try await element in rows.decode(Int.self, context: .default) { + for try await row in rows { + let element = try row.decode(Int.self) XCTAssertEqual(element, counter + 1) counter += 1 } @@ -46,6 +46,74 @@ final class AsyncPostgresConnectionTests: XCTestCase { } } + func testSelectActiveConnection() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + let query: PostgresQuery = """ + SELECT + pid + ,datname + ,usename + ,application_name + ,client_hostname + ,client_port + ,backend_start + ,query_start + ,query + ,state + FROM pg_stat_activity + WHERE state = 'active'; + """ + + try await withTestConnection(on: eventLoop) { connection in + let rows = try await connection.query(query, logger: .psqlTest) + var counter = 0 + + for try await element in rows.decode((Int, String, String, String, String?, Int, Date, Date, String, String).self) { + XCTAssertEqual(element.1, env("POSTGRES_DB") ?? "test_database") + XCTAssertEqual(element.2, env("POSTGRES_USER") ?? "test_username") + + XCTAssertEqual(element.8, query.sql) + XCTAssertEqual(element.9, "active") + counter += 1 + } + + XCTAssertGreaterThanOrEqual(counter, 1) + } + } + + func testAdditionalParametersTakeEffect() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + let query: PostgresQuery = """ + SELECT + current_setting('application_name'); + """ + + let applicationName = "postgres-nio-test" + var options = PostgresConnection.Configuration.Options() + options.additionalStartupParameters = [ + ("application_name", applicationName) + ] + + try await withTestConnection(on: eventLoop, options: options) { connection in + let rows = try await connection.query(query, logger: .psqlTest) + var counter = 0 + + for try await element in rows.decode(String.self) { + XCTAssertEqual(element, applicationName) + + counter += 1 + } + + XCTAssertGreaterThanOrEqual(counter, 1) + } + } + func testSelectTimeoutWhileLongRunningQuery() async throws { let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } @@ -68,8 +136,6 @@ final class AsyncPostgresConnectionTests: XCTestCase { } catch { guard let error = error as? PSQLError else { return XCTFail("Unexpected error type") } - print(error) - XCTAssertEqual(error.code, .server) XCTAssertEqual(error.serverInfo?[.severity], "ERROR") } @@ -188,6 +254,36 @@ final class AsyncPostgresConnectionTests: XCTestCase { } } + func testListenAndNotify() async throws { + let channelNames = [ + "foo", + "default" + ] + + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + for channelName in channelNames { + try await self.withTestConnection(on: eventLoop) { connection in + let stream = try await connection.listen(channelName) + var iterator = stream.makeAsyncIterator() + + try await self.withTestConnection(on: eventLoop) { other in + try await other.query(#"NOTIFY "\#(unescaped: channelName)", 'bar';"#, logger: .psqlTest) + + try await other.query(#"NOTIFY "\#(unescaped: channelName)", 'foo';"#, logger: .psqlTest) + } + + let first = try await iterator.next() + XCTAssertEqual(first?.payload, "bar") + + let second = try await iterator.next() + XCTAssertEqual(second?.payload, "foo") + } + } + } + #if canImport(Network) func testSelect10kRowsNetworkFramework() async throws { let eventLoopGroup = NIOTSEventLoopGroup() @@ -200,7 +296,8 @@ final class AsyncPostgresConnectionTests: XCTestCase { try await withTestConnection(on: eventLoop) { connection in let rows = try await connection.query("SELECT generate_series(\(start), \(end));", logger: .psqlTest) var counter = 1 - for try await element in rows.decode(Int.self, context: .default) { + for try await row in rows { + let element = try row.decode(Int.self, context: .default) XCTAssertEqual(element, counter) counter += 1 } @@ -256,24 +353,229 @@ final class AsyncPostgresConnectionTests: XCTestCase { try await connection.query("SELECT 1;", logger: .psqlTest) } } + + func testPreparedStatement() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + struct TestPreparedStatement: PostgresPreparedStatement { + static let sql = "SELECT pid, datname FROM pg_stat_activity WHERE state = $1" + typealias Row = (Int, String) + + var state: String + + func makeBindings() -> PostgresBindings { + var bindings = PostgresBindings() + bindings.append(self.state) + return bindings + } + + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + try row.decode(Row.self) + } + } + let preparedStatement = TestPreparedStatement(state: "active") + try await withTestConnection(on: eventLoop) { connection in + var results = try await connection.execute(preparedStatement, logger: .psqlTest) + var counter = 0 + + for try await element in results { + XCTAssertEqual(element.1, env("POSTGRES_DB") ?? "test_database") + counter += 1 + } + + XCTAssertGreaterThanOrEqual(counter, 1) + + // Second execution, which reuses the existing prepared statement + results = try await connection.execute(preparedStatement, logger: .psqlTest) + for try await element in results { + XCTAssertEqual(element.1, env("POSTGRES_DB") ?? "test_database") + counter += 1 + } + } + } + + static let preparedStatementTestTable = "AsyncTestPreparedStatementTestTable" + func testPreparedStatementWithIntegerBinding() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + struct InsertPreparedStatement: PostgresPreparedStatement { + static let name = "INSERT-AsyncTestPreparedStatementTestTable" + + static let sql = #"INSERT INTO "\#(AsyncPostgresConnectionTests.preparedStatementTestTable)" (uuid) VALUES ($1);"# + typealias Row = () + + var uuid: UUID + + func makeBindings() -> PostgresBindings { + var bindings = PostgresBindings() + bindings.append(self.uuid) + return bindings + } + + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + () + } + } + + struct SelectPreparedStatement: PostgresPreparedStatement { + static let name = "SELECT-AsyncTestPreparedStatementTestTable" + + static let sql = #"SELECT id, uuid FROM "\#(AsyncPostgresConnectionTests.preparedStatementTestTable)" WHERE id <= $1;"# + typealias Row = (Int, UUID) + + var id: Int + + func makeBindings() -> PostgresBindings { + var bindings = PostgresBindings() + bindings.append(self.id) + return bindings + } + + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + try row.decode((Int, UUID).self) + } + } + + do { + try await withTestConnection(on: eventLoop) { connection in + try await connection.query(""" + CREATE TABLE IF NOT EXISTS "\(unescaped: Self.preparedStatementTestTable)" ( + id SERIAL PRIMARY KEY, + uuid UUID NOT NULL + ) + """, + logger: .psqlTest + ) + + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + + let rows = try await connection.execute(SelectPreparedStatement(id: 3), logger: .psqlTest) + var counter = 0 + for try await (id, uuid) in rows { + Logger.psqlTest.info("Received row", metadata: [ + "id": "\(id)", "uuid": "\(uuid)" + ]) + counter += 1 + } + + try await connection.query(""" + DROP TABLE "\(unescaped: Self.preparedStatementTestTable)"; + """, + logger: .psqlTest + ) + } + } catch { + XCTFail("Unexpected error: \(String(describing: error))") + } + } + + static let preparedStatementWithOptionalTestTable = "AsyncTestPreparedStatementWithOptionalTestTable" + func testPreparedStatementWithOptionalBinding() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + struct InsertPreparedStatement: PostgresPreparedStatement { + static let name = "INSERT-AsyncTestPreparedStatementWithOptionalTestTable" + + static let sql = #"INSERT INTO "\#(AsyncPostgresConnectionTests.preparedStatementWithOptionalTestTable)" (uuid) VALUES ($1);"# + typealias Row = () + + var uuid: UUID? + + func makeBindings() -> PostgresBindings { + var bindings = PostgresBindings() + bindings.append(self.uuid) + return bindings + } + + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + () + } + } + + struct SelectPreparedStatement: PostgresPreparedStatement { + static let name = "SELECT-AsyncTestPreparedStatementWithOptionalTestTable" + + static let sql = #"SELECT id, uuid FROM "\#(AsyncPostgresConnectionTests.preparedStatementWithOptionalTestTable)" WHERE id <= $1;"# + typealias Row = (Int, UUID?) + + var id: Int + + func makeBindings() -> PostgresBindings { + var bindings = PostgresBindings() + bindings.append(self.id) + return bindings + } + + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + try row.decode((Int, UUID?).self) + } + } + + do { + try await withTestConnection(on: eventLoop) { connection in + try await connection.query(""" + CREATE TABLE IF NOT EXISTS "\(unescaped: Self.preparedStatementWithOptionalTestTable)" ( + id SERIAL PRIMARY KEY, + uuid UUID + ) + """, + logger: .psqlTest + ) + + _ = try await connection.execute(InsertPreparedStatement(uuid: nil), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: nil), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest) + _ = try await connection.execute(InsertPreparedStatement(uuid: nil), logger: .psqlTest) + + let rows = try await connection.execute(SelectPreparedStatement(id: 3), logger: .psqlTest) + var counter = 0 + for try await (id, uuid) in rows { + Logger.psqlTest.info("Received row", metadata: [ + "id": "\(id)", "uuid": "\(String(describing: uuid))" + ]) + counter += 1 + } + + try await connection.query(""" + DROP TABLE "\(unescaped: Self.preparedStatementWithOptionalTestTable)"; + """, + logger: .psqlTest + ) + } + } catch { + XCTFail("Unexpected error: \(String(describing: error))") + } + } } extension XCTestCase { func withTestConnection( on eventLoop: EventLoop, + options: PostgresConnection.Configuration.Options? = nil, file: StaticString = #filePath, line: UInt = #line, _ closure: (PostgresConnection) async throws -> Result ) async throws -> Result { - let connection = try await PostgresConnection.test(on: eventLoop).get() + let connection = try await PostgresConnection.test(on: eventLoop, options: options).get() do { let result = try await closure(connection) try await connection.close() return result } catch { - XCTFail("Unexpected error: \(error)", file: file, line: line) + XCTFail("Unexpected error: \(String(reflecting: error))", file: file, line: line) try await connection.close() throw error } diff --git a/Tests/IntegrationTests/PSQLIntegrationTests.swift b/Tests/IntegrationTests/PSQLIntegrationTests.swift index 4b2b9950..d541899b 100644 --- a/Tests/IntegrationTests/PSQLIntegrationTests.swift +++ b/Tests/IntegrationTests/PSQLIntegrationTests.swift @@ -1,6 +1,7 @@ +import Atomics import XCTest import Logging -@testable import PostgresNIO +import PostgresNIO import NIOCore import NIOPosix import NIOTestUtils @@ -73,19 +74,17 @@ final class IntegrationTests: XCTestCase { defer { XCTAssertNoThrow(try conn?.close().wait()) } var metadata: PostgresQueryMetadata? - var received: Int64 = 0 + let received = ManagedAtomic(0) XCTAssertNoThrow(metadata = try conn?.query("SELECT generate_series(1, 10000);", logger: .psqlTest) { row in func workaround() { - var number: Int64? - XCTAssertNoThrow(number = try row.decode(Int64.self, context: .default)) - received += 1 - XCTAssertEqual(number, received) + let expected = received.wrappingIncrementThenLoad(ordering: .relaxed) + XCTAssertEqual(expected, try row.decode(Int64.self, context: .default)) } workaround() }.wait()) - XCTAssertEqual(received, 10000) + XCTAssertEqual(received.load(ordering: .relaxed), 10000) XCTAssertEqual(metadata?.command, "SELECT") XCTAssertEqual(metadata?.rows, 10000) } @@ -124,6 +123,25 @@ final class IntegrationTests: XCTestCase { XCTAssertEqual(foo, "hello") } + func testQueryNothing() throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + var conn: PostgresConnection? + XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) + defer { XCTAssertNoThrow(try conn?.close().wait()) } + + var _result: PostgresQueryResult? + XCTAssertNoThrow(_result = try conn?.query(""" + -- Some comments + """, logger: .psqlTest).wait()) + + let result = try XCTUnwrap(_result) + XCTAssertEqual(result.rows, []) + XCTAssertEqual(result.metadata.command, "") + } + func testDecodeIntegers() { let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } @@ -252,7 +270,7 @@ final class IntegrationTests: XCTestCase { XCTAssertNoThrow(result = try conn?.query(""" SELECT \(Decimal(string: "123456.789123")!)::numeric as numeric, - \(Decimal(string: "-123456.789123")!)::numeric as numeric_negative + \(Decimal(string: "-123456.789123")!)::numeric as numeric_negative """, logger: .psqlTest).wait()) XCTAssertEqual(result?.rows.count, 1) @@ -263,6 +281,41 @@ final class IntegrationTests: XCTestCase { XCTAssertEqual(cells?.1, Decimal(string: "-123456.789123")) } + func testDecodeRawRepresentables() { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + let eventLoop = eventLoopGroup.next() + + var conn: PostgresConnection? + XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) + defer { XCTAssertNoThrow(try conn?.close().wait()) } + + enum StringRR: String, PostgresDecodable { + case a + } + + enum IntRR: Int, PostgresDecodable { + case b + } + + let stringValue = StringRR.a + let intValue = IntRR.b + + var result: PostgresQueryResult? + XCTAssertNoThrow(result = try conn?.query(""" + SELECT + \(stringValue.rawValue)::varchar as string, + \(intValue.rawValue)::int8 as int + """, logger: .psqlTest).wait()) + XCTAssertEqual(result?.rows.count, 1) + + var cells: (StringRR, IntRR)? + XCTAssertNoThrow(cells = try result?.rows.first?.decode((StringRR, IntRR).self, context: .default)) + + XCTAssertEqual(cells?.0, stringValue) + XCTAssertEqual(cells?.1, intValue) + } + func testRoundTripUUID() { let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } diff --git a/Tests/IntegrationTests/PostgresClientTests.swift b/Tests/IntegrationTests/PostgresClientTests.swift new file mode 100644 index 00000000..34a8ad2a --- /dev/null +++ b/Tests/IntegrationTests/PostgresClientTests.swift @@ -0,0 +1,317 @@ +@_spi(ConnectionPool) import PostgresNIO +import XCTest +import NIOPosix +import NIOSSL +import Logging +import Atomics + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +final class PostgresClientTests: XCTestCase { + + func testGetConnection() async throws { + var mlogger = Logger(label: "test") + mlogger.logLevel = .debug + let logger = mlogger + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 8) + self.addTeardownBlock { + try await eventLoopGroup.shutdownGracefully() + } + + let clientConfig = PostgresClient.Configuration.makeTestConfiguration() + let client = PostgresClient(configuration: clientConfig, eventLoopGroup: eventLoopGroup, backgroundLogger: logger) + + await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + await client.run() + } + + let iterations = 1000 + + for _ in 0.. PostgresBindings { + var bindings = PostgresBindings() + bindings.append(self.id) + return bindings + } + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + try row.decode(Row.self) + } + } + + for try await (id, uuid) in try await client.execute(Example(id: 200), logger: logger) { + logger.info("id: \(id), uuid: \(uuid.uuidString)") + } + + try await client.query( + """ + DROP TABLE "\(unescaped: tableName)"; + """, + logger: logger + ) + + taskGroup.cancelAll() + } + } catch { + XCTFail("Unexpected error: \(String(reflecting: error))") + } + } +} + +@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) +extension PostgresClient.Configuration { + static func makeTestConfiguration() -> PostgresClient.Configuration { + var tlsConfiguration = TLSConfiguration.makeClientConfiguration() + tlsConfiguration.certificateVerification = .none + var clientConfig = PostgresClient.Configuration( + host: env("POSTGRES_HOSTNAME") ?? "localhost", + port: env("POSTGRES_PORT").flatMap({ Int($0) }) ?? 5432, + username: env("POSTGRES_USER") ?? "test_username", + password: env("POSTGRES_PASSWORD") ?? "test_password", + database: env("POSTGRES_DB") ?? "test_database", + tls: .prefer(tlsConfiguration) + ) + clientConfig.options.minimumConnections = 0 + clientConfig.options.maximumConnections = 12*4 + clientConfig.options.keepAliveBehavior = .init(frequency: .seconds(5)) + clientConfig.options.connectionIdleTimeout = .seconds(15) + + return clientConfig + } +} diff --git a/Tests/IntegrationTests/PostgresNIOTests.swift b/Tests/IntegrationTests/PostgresNIOTests.swift index 348e6eb6..9a58f050 100644 --- a/Tests/IntegrationTests/PostgresNIOTests.swift +++ b/Tests/IntegrationTests/PostgresNIOTests.swift @@ -1,5 +1,6 @@ import Logging @testable import PostgresNIO +import Atomics import XCTest import NIOCore import NIOPosix @@ -9,12 +10,14 @@ import NIOSSL final class PostgresNIOTests: XCTestCase { private var group: EventLoopGroup! - private var eventLoop: EventLoop { self.group.next() } + override class func setUp() { + XCTAssertTrue(isLoggingConfigured) + } + override func setUpWithError() throws { try super.setUpWithError() - XCTAssertTrue(isLoggingConfigured) self.group = MultiThreadedEventLoopGroup(numberOfThreads: 1) } @@ -110,59 +113,59 @@ final class PostgresNIOTests: XCTestCase { XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var receivedNotifications: [PostgresMessage.NotificationResponse] = [] + let receivedNotifications = ManagedAtomic(0) conn?.addListener(channel: "example") { context, notification in - receivedNotifications.append(notification) + receivedNotifications.wrappingIncrement(ordering: .relaxed) + XCTAssertEqual(notification.channel, "example") + XCTAssertEqual(notification.payload, "") } XCTAssertNoThrow(_ = try conn?.simpleQuery("LISTEN example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) // Notifications are asynchronous, so we should run at least one more query to make sure we'll have received the notification response by then XCTAssertNoThrow(_ = try conn?.simpleQuery("SELECT 1").wait()) - XCTAssertEqual(receivedNotifications.count, 1) - XCTAssertEqual(receivedNotifications.first?.channel, "example") - XCTAssertEqual(receivedNotifications.first?.payload, "") + XCTAssertEqual(receivedNotifications.load(ordering: .relaxed), 1) } func testNotificationsNonEmptyPayload() { var conn: PostgresConnection? XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var receivedNotifications: [PostgresMessage.NotificationResponse] = [] + let receivedNotifications = ManagedAtomic(0) conn?.addListener(channel: "example") { context, notification in - receivedNotifications.append(notification) + receivedNotifications.wrappingIncrement(ordering: .relaxed) + XCTAssertEqual(notification.channel, "example") + XCTAssertEqual(notification.payload, "Notification payload example") } XCTAssertNoThrow(_ = try conn?.simpleQuery("LISTEN example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example, 'Notification payload example'").wait()) // Notifications are asynchronous, so we should run at least one more query to make sure we'll have received the notification response by then XCTAssertNoThrow(_ = try conn?.simpleQuery("SELECT 1").wait()) - XCTAssertEqual(receivedNotifications.count, 1) - XCTAssertEqual(receivedNotifications.first?.channel, "example") - XCTAssertEqual(receivedNotifications.first?.payload, "Notification payload example") + XCTAssertEqual(receivedNotifications.load(ordering: .relaxed), 1) } func testNotificationsRemoveHandlerWithinHandler() { var conn: PostgresConnection? XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var receivedNotifications = 0 + let receivedNotifications = ManagedAtomic(0) conn?.addListener(channel: "example") { context, notification in - receivedNotifications += 1 + receivedNotifications.wrappingIncrement(ordering: .relaxed) context.stop() } XCTAssertNoThrow(_ = try conn?.simpleQuery("LISTEN example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("SELECT 1").wait()) - XCTAssertEqual(receivedNotifications, 1) + XCTAssertEqual(receivedNotifications.load(ordering: .relaxed), 1) } func testNotificationsRemoveHandlerOutsideHandler() { var conn: PostgresConnection? XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var receivedNotifications = 0 + let receivedNotifications = ManagedAtomic(0) let context = conn?.addListener(channel: "example") { context, notification in - receivedNotifications += 1 + receivedNotifications.wrappingIncrement(ordering: .relaxed) } XCTAssertNotNil(context) XCTAssertNoThrow(_ = try conn?.simpleQuery("LISTEN example").wait()) @@ -171,47 +174,47 @@ final class PostgresNIOTests: XCTestCase { context?.stop() XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("SELECT 1").wait()) - XCTAssertEqual(receivedNotifications, 1) + XCTAssertEqual(receivedNotifications.load(ordering: .relaxed), 1) } func testNotificationsMultipleRegisteredHandlers() { var conn: PostgresConnection? XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var receivedNotifications1 = 0 + let receivedNotifications1 = ManagedAtomic(0) conn?.addListener(channel: "example") { context, notification in - receivedNotifications1 += 1 + receivedNotifications1.wrappingIncrement(ordering: .relaxed) } - var receivedNotifications2 = 0 + let receivedNotifications2 = ManagedAtomic(0) conn?.addListener(channel: "example") { context, notification in - receivedNotifications2 += 1 + receivedNotifications2.wrappingIncrement(ordering: .relaxed) } XCTAssertNoThrow(_ = try conn?.simpleQuery("LISTEN example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("SELECT 1").wait()) - XCTAssertEqual(receivedNotifications1, 1) - XCTAssertEqual(receivedNotifications2, 1) + XCTAssertEqual(receivedNotifications1.load(ordering: .relaxed), 1) + XCTAssertEqual(receivedNotifications2.load(ordering: .relaxed), 1) } func testNotificationsMultipleRegisteredHandlersRemoval() throws { var conn: PostgresConnection? XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var receivedNotifications1 = 0 + let receivedNotifications1 = ManagedAtomic(0) XCTAssertNotNil(conn?.addListener(channel: "example") { context, notification in - receivedNotifications1 += 1 + receivedNotifications1.wrappingIncrement(ordering: .relaxed) context.stop() }) - var receivedNotifications2 = 0 + let receivedNotifications2 = ManagedAtomic(0) XCTAssertNotNil(conn?.addListener(channel: "example") { context, notification in - receivedNotifications2 += 1 + receivedNotifications2.wrappingIncrement(ordering: .relaxed) }) XCTAssertNoThrow(_ = try conn?.simpleQuery("LISTEN example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("NOTIFY example").wait()) XCTAssertNoThrow(_ = try conn?.simpleQuery("SELECT 1").wait()) - XCTAssertEqual(receivedNotifications1, 1) - XCTAssertEqual(receivedNotifications2, 2) + XCTAssertEqual(receivedNotifications1.load(ordering: .relaxed), 1) + XCTAssertEqual(receivedNotifications2.load(ordering: .relaxed), 2) } func testNotificationHandlerFiltersOnChannel() { @@ -374,6 +377,120 @@ final class PostgresNIOTests: XCTestCase { XCTAssertEqual(UUID(uuidString: row?[data: "id"].string ?? ""), UUID(uuidString: "123E4567-E89B-12D3-A456-426655440000")) } + func testInt4Range() async throws { + let conn: PostgresConnection = try await PostgresConnection.test(on: eventLoop).get() + self.addTeardownBlock { + try await conn.close() + } + struct Model: Decodable { + let range: Range + } + let results1: PostgresQueryResult = try await conn.query(""" + SELECT + '[\(Int32.min), \(Int32.max))'::int4range AS range + """).get() + XCTAssertEqual(results1.count, 1) + var row = results1.first?.makeRandomAccess() + let expectedRange: Range = Int32.min...self, context: .default) + XCTAssertEqual(decodedRange, expectedRange) + + let results2 = try await conn.query(""" + SELECT + ARRAY[ + '[0, 1)'::int4range, + '[10, 11)'::int4range + ] AS ranges + """).get() + XCTAssertEqual(results2.count, 1) + row = results2.first?.makeRandomAccess() + let decodedRangeArray = try row?.decode(column: "ranges", as: [Range].self, context: .default) + let decodedClosedRangeArray = try row?.decode(column: "ranges", as: [ClosedRange].self, context: .default) + XCTAssertEqual(decodedRangeArray, [0..<1, 10..<11]) + XCTAssertEqual(decodedClosedRangeArray, [0...0, 10...10]) + } + + func testEmptyInt4Range() async throws { + let conn: PostgresConnection = try await PostgresConnection.test(on: eventLoop).get() + self.addTeardownBlock { + try await conn.close() + } + struct Model: Decodable { + let range: Range + } + let randomValue = Int32.random(in: Int32.min...Int32.max) + let results: PostgresQueryResult = try await conn.query(""" + SELECT + '[\(randomValue),\(randomValue))'::int4range AS range + """).get() + XCTAssertEqual(results.count, 1) + let row = results.first?.makeRandomAccess() + let expectedRange: Range = Int32.valueForEmptyRange...self, context: .default) + XCTAssertEqual(decodedRange, expectedRange) + + XCTAssertThrowsError( + try row?.decode(column: "range", as: ClosedRange.self, context: .default) + ) + } + + func testInt8Range() async throws { + let conn: PostgresConnection = try await PostgresConnection.test(on: eventLoop).get() + self.addTeardownBlock { + try await conn.close() + } + struct Model: Decodable { + let range: Range + } + let results1: PostgresQueryResult = try await conn.query(""" + SELECT + '[\(Int64.min), \(Int64.max))'::int8range AS range + """).get() + XCTAssertEqual(results1.count, 1) + var row = results1.first?.makeRandomAccess() + let expectedRange: Range = Int64.min...self, context: .default) + XCTAssertEqual(decodedRange, expectedRange) + + let results2: PostgresQueryResult = try await conn.query(""" + SELECT + ARRAY[ + '[0, 1)'::int8range, + '[10, 11)'::int8range + ] AS ranges + """).get() + XCTAssertEqual(results2.count, 1) + row = results2.first?.makeRandomAccess() + let decodedRangeArray = try row?.decode(column: "ranges", as: [Range].self, context: .default) + let decodedClosedRangeArray = try row?.decode(column: "ranges", as: [ClosedRange].self, context: .default) + XCTAssertEqual(decodedRangeArray, [0..<1, 10..<11]) + XCTAssertEqual(decodedClosedRangeArray, [0...0, 10...10]) + } + + func testEmptyInt8Range() async throws { + let conn: PostgresConnection = try await PostgresConnection.test(on: eventLoop).get() + self.addTeardownBlock { + try await conn.close() + } + struct Model: Decodable { + let range: Range + } + let randomValue = Int64.random(in: Int64.min...Int64.max) + let results: PostgresQueryResult = try await conn.query(""" + SELECT + '[\(randomValue),\(randomValue))'::int8range AS range + """).get() + XCTAssertEqual(results.count, 1) + let row = results.first?.makeRandomAccess() + let expectedRange: Range = Int64.valueForEmptyRange...self, context: .default) + XCTAssertEqual(decodedRange, expectedRange) + + XCTAssertThrowsError( + try row?.decode(column: "range", as: ClosedRange.self, context: .default) + ) + } + func testDates() { var conn: PostgresConnection? XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) @@ -501,21 +618,25 @@ final class PostgresNIOTests: XCTestCase { let a = PostgresNumeric(string: "123456.789123")! let b = PostgresNumeric(string: "-123456.789123")! let c = PostgresNumeric(string: "3.14159265358979")! + let d = PostgresNumeric(string: "1234567898765")! var rows: PostgresQueryResult? XCTAssertNoThrow(rows = try conn?.query(""" select $1::numeric as a, $2::numeric as b, - $3::numeric as c + $3::numeric as c, + $4::numeric as d """, [ .init(numeric: a), .init(numeric: b), - .init(numeric: c) + .init(numeric: c), + .init(numeric: d) ]).wait()) let row = rows?.first?.makeRandomAccess() XCTAssertEqual(row?[data: "a"].decimal, Decimal(string: "123456.789123")!) XCTAssertEqual(row?[data: "b"].decimal, Decimal(string: "-123456.789123")!) XCTAssertEqual(row?[data: "c"].decimal, Decimal(string: "3.14159265358979")!) + XCTAssertEqual(row?[data: "d"].decimal, Decimal(string: "1234567898765")!) } func testDecimalStringSerialization() { @@ -669,6 +790,44 @@ final class PostgresNIOTests: XCTestCase { XCTAssertEqual(row?[data: "array"].array(of: Int64?.self), [1, nil, 3]) } + @available(*, deprecated, message: "Testing deprecated functionality") + func testDateArraySerialize() { + var conn: PostgresConnection? + XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) + defer { XCTAssertNoThrow( try conn?.close().wait() ) } + let date1 = Date(timeIntervalSince1970: 1704088800), + date2 = Date(timeIntervalSince1970: 1706767200), + date3 = Date(timeIntervalSince1970: 1709272800) + var rows: PostgresQueryResult? + XCTAssertNoThrow(rows = try conn?.query(""" + select + $1::timestamptz[] as array + """, [ + PostgresData(array: [date1, date2, date3]) + ]).wait()) + let row = rows?.first?.makeRandomAccess() + XCTAssertEqual(row?[data: "array"].array(of: Date.self), [date1, date2, date3]) + } + + @available(*, deprecated, message: "Testing deprecated functionality") + func testDateArraySerializeAsPostgresDate() { + var conn: PostgresConnection? + XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) + defer { XCTAssertNoThrow(try conn?.close().wait()) } + let date1 = Date(timeIntervalSince1970: 1704088800),//8766 + date2 = Date(timeIntervalSince1970: 1706767200),//8797 + date3 = Date(timeIntervalSince1970: 1709272800) //8826 + var data = PostgresData(array: [date1, date2, date3].map { Int32(($0.timeIntervalSince1970 - 946_684_800) / 86_400).postgresData }, elementType: .date) + data.type = .dateArray // N.B.: `.date` format is an Int32 count of days since psqlStartDate + var rows: PostgresQueryResult? + XCTAssertNoThrow(rows = try conn?.query("select $1::date[] as array", [data]).wait()) + let row = rows?.first?.makeRandomAccess() + XCTAssertEqual( + row?[data: "array"].array(of: Date.self)?.map { Int32((($0.timeIntervalSince1970 - 946_684_800) / 86_400).rounded(.toNearestOrAwayFromZero)) }, + [date1, date2, date3].map { Int32((($0.timeIntervalSince1970 - 946_684_800) / 86_400).rounded(.toNearestOrAwayFromZero)) } + ) + } + // https://github.com/vapor/postgres-nio/issues/143 func testEmptyStringFromNonNullColumn() { var conn: PostgresConnection? @@ -771,26 +930,106 @@ final class PostgresNIOTests: XCTestCase { } } - func testRemoteTLSServer() { - // postgres://uymgphwj:7_tHbREdRwkqAdu4KoIS7hQnNxr8J1LA@elmer.db.elephantsql.com:5432/uymgphwj - var conn: PostgresConnection? - let logger = Logger(label: "test") - let sslContext = try! NIOSSLContext(configuration: .makeClientConfiguration()) - let config = PostgresConnection.Configuration( - host: "elmer.db.elephantsql.com", - port: 5432, - username: "uymgphwj", - password: "7_tHbREdRwkqAdu4KoIS7hQnNxr8J1LA", - database: "uymgphwj", - tls: .require(sslContext) - ) - XCTAssertNoThrow(conn = try PostgresConnection.connect(on: eventLoop, configuration: config, id: 0, logger: logger).wait()) - defer { XCTAssertNoThrow( try conn?.close().wait() ) } - var rows: [PostgresRow]? - XCTAssertNoThrow(rows = try conn?.simpleQuery("SELECT version()").wait()) - XCTAssertEqual(rows?.count, 1) - let row = rows?.first?.makeRandomAccess() - XCTAssertEqual(row?[data: "version"].string?.contains("PostgreSQL"), true) + func testInt4RangeSerialize() async throws { + let conn: PostgresConnection = try await PostgresConnection.test(on: eventLoop).get() + self.addTeardownBlock { + try await conn.close() + } + do { + let range: Range = Int32.min..? = try row?.decode(Range.self, context: .default) + XCTAssertEqual(range, decodedRange) + } + do { + let emptyRange: Range = Int32.min..? = try row?.decode(Range.self, context: .default) + let expectedRange: Range = Int32.valueForEmptyRange.. = Int32.min...(Int32.max - 1) + var binds = PostgresBindings() + binds.append(closedRange, context: .default) + let query = PostgresQuery( + unsafeSQL: "select $1::int4range as range", + binds: binds + ) + let rowSequence: PostgresRowSequence? = try await conn.query(query, logger: .psqlTest) + var rowIterator: PostgresRowSequence.AsyncIterator? = rowSequence?.makeAsyncIterator() + let row: PostgresRow? = try await rowIterator?.next() + let decodedClosedRange: ClosedRange? = try row?.decode(ClosedRange.self, context: .default) + XCTAssertEqual(closedRange, decodedClosedRange) + } + } + + func testInt8RangeSerialize() async throws { + let conn: PostgresConnection = try await PostgresConnection.test(on: eventLoop).get() + self.addTeardownBlock { + try await conn.close() + } + do { + let range: Range = Int64.min..? = try row?.decode(Range.self, context: .default) + XCTAssertEqual(range, decodedRange) + } + do { + let emptyRange: Range = Int64.min..? = try row?.decode(Range.self, context: .default) + let expectedRange: Range = Int64.valueForEmptyRange.. = Int64.min...(Int64.max - 1) + var binds = PostgresBindings() + binds.append(closedRange, context: .default) + let query = PostgresQuery( + unsafeSQL: "select $1::int8range as range", + binds: binds + ) + let rowSequence: PostgresRowSequence? = try await conn.query(query, logger: .psqlTest) + var rowIterator: PostgresRowSequence.AsyncIterator? = rowSequence?.makeAsyncIterator() + let row: PostgresRow? = try await rowIterator?.next() + let decodedClosedRange: ClosedRange? = try row?.decode(ClosedRange.self, context: .default) + XCTAssertEqual(closedRange, decodedClosedRange) + } } @available(*, deprecated, message: "Test deprecated functionality") @@ -1023,17 +1262,17 @@ final class PostgresNIOTests: XCTestCase { XCTAssertNoThrow(conn = try PostgresConnection.test(on: eventLoop).wait()) defer { XCTAssertNoThrow( try conn?.close().wait() ) } var queries: [[PostgresRow]]? - XCTAssertNoThrow(queries = try conn?.prepare(query: "SELECT $1::text as foo;", handler: { query in + XCTAssertNoThrow(queries = try conn?.prepare(query: "SELECT $1::text as foo;", handler: { [eventLoop] query in let a = query.execute(["a"]) let b = query.execute(["b"]) let c = query.execute(["c"]) - return EventLoopFuture.whenAllSucceed([a, b, c], on: self.eventLoop) + return EventLoopFuture.whenAllSucceed([a, b, c], on: eventLoop) }).wait()) XCTAssertEqual(queries?.count, 3) - var resutIterator = queries?.makeIterator() - XCTAssertEqual(try resutIterator?.next()?.first?.decode(String.self, context: .default), "a") - XCTAssertEqual(try resutIterator?.next()?.first?.decode(String.self, context: .default), "b") - XCTAssertEqual(try resutIterator?.next()?.first?.decode(String.self, context: .default), "c") + var resultIterator = queries?.makeIterator() + XCTAssertEqual(try resultIterator?.next()?.first?.decode(String.self, context: .default), "a") + XCTAssertEqual(try resultIterator?.next()?.first?.decode(String.self, context: .default), "b") + XCTAssertEqual(try resultIterator?.next()?.first?.decode(String.self, context: .default), "c") } // https://github.com/vapor/postgres-nio/issues/122 @@ -1221,7 +1460,7 @@ final class PostgresNIOTests: XCTestCase { let isLoggingConfigured: Bool = { LoggingSystem.bootstrap { label in var handler = StreamLogHandler.standardOutput(label: label) - handler.logLevel = env("LOG_LEVEL").flatMap { Logger.Level(rawValue: $0) } ?? .debug + handler.logLevel = env("LOG_LEVEL").flatMap { .init(rawValue: $0) } ?? .info return handler } return true diff --git a/Tests/IntegrationTests/Utilities.swift b/Tests/IntegrationTests/Utilities.swift index b1788110..91dbb62e 100644 --- a/Tests/IntegrationTests/Utilities.swift +++ b/Tests/IntegrationTests/Utilities.swift @@ -24,11 +24,9 @@ extension PostgresConnection { } } - static func test(on eventLoop: EventLoop, logLevel: Logger.Level = .info) -> EventLoopFuture { - var logger = Logger(label: "postgres.connection.test") - logger.logLevel = logLevel - - let config = PostgresConnection.Configuration( + static func test(on eventLoop: EventLoop, options: Configuration.Options? = nil) -> EventLoopFuture { + let logger = Logger(label: "postgres.connection.test") + var config = PostgresConnection.Configuration( host: env("POSTGRES_HOSTNAME") ?? "localhost", port: env("POSTGRES_PORT").flatMap(Int.init(_:)) ?? 5432, username: env("POSTGRES_USER") ?? "test_username", @@ -36,14 +34,15 @@ extension PostgresConnection { database: env("POSTGRES_DB") ?? "test_database", tls: .disable ) - + if let options { + config.options = options + } + return PostgresConnection.connect(on: eventLoop, configuration: config, id: 0, logger: logger) } - static func testUDS(on eventLoop: EventLoop, logLevel: Logger.Level = .info) -> EventLoopFuture { - var logger = Logger(label: "postgres.connection.test") - logger.logLevel = logLevel - + static func testUDS(on eventLoop: EventLoop) -> EventLoopFuture { + let logger = Logger(label: "postgres.connection.test") let config = PostgresConnection.Configuration( unixSocketPath: env("POSTGRES_SOCKET") ?? "/tmp/.s.PGSQL.\(env("POSTGRES_PORT").flatMap(Int.init(_:)) ?? 5432)", username: env("POSTGRES_USER") ?? "test_username", @@ -54,10 +53,8 @@ extension PostgresConnection { return PostgresConnection.connect(on: eventLoop, configuration: config, id: 0, logger: logger) } - static func testChannel(_ channel: Channel, on eventLoop: EventLoop, logLevel: Logger.Level = .info) -> EventLoopFuture { - var logger = Logger(label: "postgres.connection.test") - logger.logLevel = logLevel - + static func testChannel(_ channel: Channel, on eventLoop: EventLoop) -> EventLoopFuture { + let logger = Logger(label: "postgres.connection.test") let config = PostgresConnection.Configuration( establishedChannel: channel, username: env("POSTGRES_USER") ?? "test_username", @@ -71,9 +68,7 @@ extension PostgresConnection { extension Logger { static var psqlTest: Logger { - var logger = Logger(label: "psql.test") - logger.logLevel = .info - return logger + .init(label: "psql.test") } } diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/AuthenticationStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/AuthenticationStateMachineTests.swift index 87478e63..df881f90 100644 --- a/Tests/PostgresNIOTests/New/Connection State Machine/AuthenticationStateMachineTests.swift +++ b/Tests/PostgresNIOTests/New/Connection State Machine/AuthenticationStateMachineTests.swift @@ -19,8 +19,8 @@ class AuthenticationStateMachineTests: XCTestCase { let authContext = AuthContext(username: "test", password: "abc123", database: "test") var state = ConnectionStateMachine(requireBackendKeyData: true) XCTAssertEqual(state.connected(tls: .disable), .provideAuthenticationContext) - let salt: (UInt8, UInt8, UInt8, UInt8) = (0, 1, 2, 3) - + let salt: UInt32 = 0x00_01_02_03 + XCTAssertEqual(state.provideAuthenticationContext(authContext), .sendStartupMessage(authContext)) XCTAssertEqual(state.authenticationMessageReceived(.md5(salt: salt)), .sendPasswordMessage(.md5(salt: salt), authContext)) XCTAssertEqual(state.authenticationMessageReceived(.ok), .wait) @@ -30,8 +30,8 @@ class AuthenticationStateMachineTests: XCTestCase { let authContext = AuthContext(username: "test", password: nil, database: "test") var state = ConnectionStateMachine(requireBackendKeyData: true) XCTAssertEqual(state.connected(tls: .disable), .provideAuthenticationContext) - let salt: (UInt8, UInt8, UInt8, UInt8) = (0, 1, 2, 3) - + let salt: UInt32 = 0x00_01_02_03 + XCTAssertEqual(state.provideAuthenticationContext(authContext), .sendStartupMessage(authContext)) XCTAssertEqual(state.authenticationMessageReceived(.md5(salt: salt)), .closeConnectionAndCleanup(.init(action: .close, tasks: [], error: .authMechanismRequiresPassword, closePromise: nil))) @@ -45,12 +45,36 @@ class AuthenticationStateMachineTests: XCTestCase { XCTAssertEqual(state.authenticationMessageReceived(.ok), .wait) } - func testAuthenticationFailure() { + func testAuthenticateSCRAMSHA256WithAtypicalEncoding() { let authContext = AuthContext(username: "test", password: "abc123", database: "test") var state = ConnectionStateMachine(requireBackendKeyData: true) XCTAssertEqual(state.connected(tls: .disable), .provideAuthenticationContext) - let salt: (UInt8, UInt8, UInt8, UInt8) = (0, 1, 2, 3) + XCTAssertEqual(state.provideAuthenticationContext(authContext), .sendStartupMessage(authContext)) + + let saslResponse = state.authenticationMessageReceived(.sasl(names: ["SCRAM-SHA-256"])) + guard case .sendSaslInitialResponse(name: let name, initialResponse: let responseData) = saslResponse else { + return XCTFail("\(saslResponse) is not .sendSaslInitialResponse") + } + let responseString = String(decoding: responseData, as: UTF8.self) + XCTAssertEqual(name, "SCRAM-SHA-256") + XCTAssert(responseString.starts(with: "n,,n=test,r=")) + let saslContinueResponse = state.authenticationMessageReceived(.saslContinue(data: .init(bytes: + "r=\(responseString.dropFirst(12))RUJSZHhkeUVFNzRLNERKMkxmU05ITU1NZWcxaQ==,s=ijgUVaWgCDLRJyF963BKNA==,i=4096".utf8 + ))) + guard case .sendSaslResponse(let responseData2) = saslContinueResponse else { + return XCTFail("\(saslContinueResponse) is not .sendSaslResponse") + } + let response2String = String(decoding: responseData2, as: UTF8.self) + XCTAssertEqual(response2String.prefix(76), "c=biws,r=\(responseString.dropFirst(12))RUJSZHhkeUVFNzRLNERKMkxmU05ITU1NZWcxaQ==,p=") + } + + func testAuthenticationFailure() { + let authContext = AuthContext(username: "test", password: "abc123", database: "test") + var state = ConnectionStateMachine(requireBackendKeyData: true) + XCTAssertEqual(state.connected(tls: .disable), .provideAuthenticationContext) + let salt: UInt32 = 0x00_01_02_03 + XCTAssertEqual(state.provideAuthenticationContext(authContext), .sendStartupMessage(authContext)) XCTAssertEqual(state.authenticationMessageReceived(.md5(salt: salt)), .sendPasswordMessage(.md5(salt: salt), authContext)) let fields: [PostgresBackendMessage.Field: String] = [ @@ -107,12 +131,12 @@ class AuthenticationStateMachineTests: XCTestCase { } func testUnexpectedMessagesAfterPasswordSent() { - let salt: (UInt8, UInt8, UInt8, UInt8) = (0, 1, 2, 3) + let salt: UInt32 = 0x00_01_02_03 var buffer = ByteBuffer() buffer.writeBytes([0, 1, 2, 3, 4, 5, 6, 7, 8]) let unexpected: [PostgresBackendMessage.Authentication] = [ .kerberosV5, - .md5(salt: (0, 1, 2, 3)), + .md5(salt: salt), .plaintext, .scmCredential, .gss, diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/ConnectionStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/ConnectionStateMachineTests.swift index 289665fb..f3d72a5e 100644 --- a/Tests/PostgresNIOTests/New/Connection State Machine/ConnectionStateMachineTests.swift +++ b/Tests/PostgresNIOTests/New/Connection State Machine/ConnectionStateMachineTests.swift @@ -23,7 +23,7 @@ class ConnectionStateMachineTests: XCTestCase { XCTAssertEqual(state.sslHandlerAdded(), .wait) XCTAssertEqual(state.sslEstablished(), .provideAuthenticationContext) XCTAssertEqual(state.provideAuthenticationContext(authContext), .sendStartupMessage(authContext)) - let salt: (UInt8, UInt8, UInt8, UInt8) = (0,1,2,3) + let salt: UInt32 = 0x00_01_02_03 XCTAssertEqual(state.authenticationMessageReceived(.md5(salt: salt)), .sendPasswordMessage(.md5(salt: salt), authContext)) } @@ -137,14 +137,14 @@ class ConnectionStateMachineTests: XCTestCase { func testErrorIsIgnoredWhenClosingConnection() { // test ignore unclean shutdown when closing connection - var stateIgnoreChannelError = ConnectionStateMachine(.closing) - + var stateIgnoreChannelError = ConnectionStateMachine(.closing(nil)) + XCTAssertEqual(stateIgnoreChannelError.errorHappened(.connectionError(underlying: NIOSSLError.uncleanShutdown)), .wait) XCTAssertEqual(stateIgnoreChannelError.closed(), .fireChannelInactive) // test ignore any other error when closing connection - var stateIgnoreErrorMessage = ConnectionStateMachine(.closing) + var stateIgnoreErrorMessage = ConnectionStateMachine(.closing(nil)) XCTAssertEqual(stateIgnoreErrorMessage.errorReceived(.init(fields: [:])), .wait) XCTAssertEqual(stateIgnoreErrorMessage.closed(), .fireChannelInactive) } @@ -154,7 +154,7 @@ class ConnectionStateMachineTests: XCTestCase { defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } let authContext = AuthContext(username: "test", password: "abc123", database: "test") - let salt: (UInt8, UInt8, UInt8, UInt8) = (0, 1, 2, 3) + let salt: UInt32 = 0x00_01_02_03 let queryPromise = eventLoopGroup.next().makePromise(of: PSQLRowStream.self) @@ -180,9 +180,9 @@ class ConnectionStateMachineTests: XCTestCase { XCTAssertEqual(state.errorReceived(.init(fields: fields)), .closeConnectionAndCleanup(.init(action: .close, tasks: [.extendedQuery(extendedQueryContext)], error: .server(.init(fields: fields)), closePromise: nil))) - XCTAssertNil(extendedQueryContext.promise.futureResult._value) - + XCTAssertNil(queryPromise.futureResult._value) + // make sure we don't crash - extendedQueryContext.promise.fail(PSQLError.server(.init(fields: fields))) + queryPromise.fail(PSQLError.server(.init(fields: fields))) } } diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift index eac46e5f..ae484acc 100644 --- a/Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift +++ b/Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift @@ -20,7 +20,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait) XCTAssertEqual(state.noDataReceived(), .wait) XCTAssertEqual(state.bindCompleteReceived(), .wait) - XCTAssertEqual(state.commandCompletedReceived("DELETE 1"), .succeedQueryNoRowsComming(queryContext, commandTag: "DELETE 1")) + XCTAssertEqual(state.commandCompletedReceived("DELETE 1"), .succeedQuery(promise, with: .init(value: .noRows(.tag("DELETE 1")), logger: logger))) XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) } @@ -49,7 +49,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { } XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait) - XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(queryContext, columns: expected)) + XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(promise, with: .init(value: .rowDescription(expected), logger: logger))) let row1: DataRow = [ByteBuffer(string: "test1")] XCTAssertEqual(state.dataRowReceived(row1), .wait) XCTAssertEqual(state.channelReadComplete(), .forwardRows([row1])) @@ -77,7 +77,25 @@ class ExtendedQueryStateMachineTests: XCTestCase { XCTAssertEqual(state.commandCompletedReceived("SELECT 2"), .forwardStreamComplete([row5, row6], commandTag: "SELECT 2")) XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) } - + + func testExtendedQueryWithNoQuery() { + var state = ConnectionStateMachine.readyForQuery() + + let logger = Logger.psqlTest + let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self) + promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all. + let query: PostgresQuery = "-- some comments" + let queryContext = ExtendedQueryContext(query: query, logger: logger, promise: promise) + + XCTAssertEqual(state.enqueue(task: .extendedQuery(queryContext)), .sendParseDescribeBindExecuteSync(query)) + XCTAssertEqual(state.parseCompleteReceived(), .wait) + XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait) + XCTAssertEqual(state.noDataReceived(), .wait) + XCTAssertEqual(state.bindCompleteReceived(), .wait) + XCTAssertEqual(state.emptyQueryResponseReceived(), .succeedQuery(promise, with: .init(value: .noRows(.emptyResponse), logger: logger))) + XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) + } + func testReceiveTotallyUnexpectedMessageInQuery() { var state = ConnectionStateMachine.readyForQuery() @@ -93,7 +111,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { let psqlError = PSQLError.unexpectedBackendMessage(.authentication(.ok)) XCTAssertEqual(state.authenticationMessageReceived(.ok), - .failQuery(queryContext, with: psqlError, cleanupContext: .init(action: .close, tasks: [], error: psqlError, closePromise: nil))) + .failQuery(promise, with: psqlError, cleanupContext: .init(action: .close, tasks: [], error: psqlError, closePromise: nil))) } func testExtendedQueryIsCancelledImmediatly() { @@ -121,7 +139,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { } XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait) - XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(queryContext, columns: expected)) + XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(promise, with: .init(value: .rowDescription(expected), logger: logger))) XCTAssertEqual(state.cancelQueryStream(), .forwardStreamError(.queryCancelled, read: false, cleanupContext: nil)) XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test1")]), .wait) XCTAssertEqual(state.channelReadComplete(), .wait) @@ -165,7 +183,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { } XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait) - XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(queryContext, columns: expected)) + XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(promise, with: .init(value: .rowDescription(expected), logger: logger))) let row1: DataRow = [ByteBuffer(string: "test1")] XCTAssertEqual(state.dataRowReceived(row1), .wait) XCTAssertEqual(state.channelReadComplete(), .forwardRows([row1])) @@ -207,7 +225,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { } XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait) - XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(queryContext, columns: expected)) + XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(promise, with: .init(value: .rowDescription(expected), logger: logger))) let dataRows1: [DataRow] = [ [ByteBuffer(string: "test1")], [ByteBuffer(string: "test2")], @@ -251,7 +269,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { let serverError = PostgresBackendMessage.ErrorResponse(fields: [.severity: "Error", .sqlState: "123"]) XCTAssertEqual( - state.errorReceived(serverError), .failQuery(queryContext, with: .server(serverError), cleanupContext: .none) + state.errorReceived(serverError), .failQuery(promise, with: .server(serverError), cleanupContext: .none) ) XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) @@ -269,7 +287,7 @@ class ExtendedQueryStateMachineTests: XCTestCase { XCTAssertEqual(state.enqueue(task: .extendedQuery(queryContext)), .sendParseDescribeBindExecuteSync(query)) XCTAssertEqual(state.parseCompleteReceived(), .wait) XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait) - XCTAssertEqual(state.cancelQueryStream(), .failQuery(queryContext, with: .queryCancelled, cleanupContext: .none)) + XCTAssertEqual(state.cancelQueryStream(), .failQuery(promise, with: .queryCancelled, cleanupContext: .none)) let serverError = PostgresBackendMessage.ErrorResponse(fields: [.severity: "Error", .sqlState: "123"]) XCTAssertEqual(state.errorReceived(serverError), .wait) diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/PrepareStatementStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/PrepareStatementStateMachineTests.swift index 6cff280e..547f5cdf 100644 --- a/Tests/PostgresNIOTests/New/Connection State Machine/PrepareStatementStateMachineTests.swift +++ b/Tests/PostgresNIOTests/New/Connection State Machine/PrepareStatementStateMachineTests.swift @@ -3,7 +3,6 @@ import NIOEmbedded @testable import PostgresNIO class PrepareStatementStateMachineTests: XCTestCase { - func testCreatePreparedStatementReturningRowDescription() { var state = ConnectionStateMachine.readyForQuery() @@ -12,11 +11,12 @@ class PrepareStatementStateMachineTests: XCTestCase { let name = "haha" let query = #"SELECT id FROM users WHERE id = $1 "# - let prepareStatementContext = PrepareStatementContext( - name: name, query: query, logger: .psqlTest, promise: promise) - - XCTAssertEqual(state.enqueue(task: .preparedStatement(prepareStatementContext)), - .sendParseDescribeSync(name: name, query: query)) + let prepareStatementContext = ExtendedQueryContext( + name: name, query: query, bindingDataTypes: [], logger: .psqlTest, promise: promise + ) + + XCTAssertEqual(state.enqueue(task: .extendedQuery(prepareStatementContext)), + .sendParseDescribeSync(name: name, query: query, bindingDataTypes: [])) XCTAssertEqual(state.parseCompleteReceived(), .wait) XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait) @@ -25,7 +25,7 @@ class PrepareStatementStateMachineTests: XCTestCase { ] XCTAssertEqual(state.rowDescriptionReceived(.init(columns: columns)), - .succeedPreparedStatementCreation(prepareStatementContext, with: .init(columns: columns))) + .succeedPreparedStatementCreation(promise, with: .init(columns: columns))) XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) } @@ -37,25 +37,42 @@ class PrepareStatementStateMachineTests: XCTestCase { let name = "haha" let query = #"DELETE FROM users WHERE id = $1 "# - let prepareStatementContext = PrepareStatementContext( - name: name, query: query, logger: .psqlTest, promise: promise) - - XCTAssertEqual(state.enqueue(task: .preparedStatement(prepareStatementContext)), - .sendParseDescribeSync(name: name, query: query)) + let prepareStatementContext = ExtendedQueryContext( + name: name, query: query, bindingDataTypes: [], logger: .psqlTest, promise: promise + ) + + XCTAssertEqual(state.enqueue(task: .extendedQuery(prepareStatementContext)), + .sendParseDescribeSync(name: name, query: query, bindingDataTypes: [])) XCTAssertEqual(state.parseCompleteReceived(), .wait) XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait) XCTAssertEqual(state.noDataReceived(), - .succeedPreparedStatementCreation(prepareStatementContext, with: nil)) + .succeedPreparedStatementCreation(promise, with: nil)) XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) } func testErrorReceivedAfter() { - let connectionContext = ConnectionStateMachine.createConnectionContext() - var state = ConnectionStateMachine(.prepareStatement(.init(.noDataMessageReceived), connectionContext)) - + var state = ConnectionStateMachine.readyForQuery() + + let promise = EmbeddedEventLoop().makePromise(of: RowDescription?.self) + promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all. + + let name = "haha" + let query = #"DELETE FROM users WHERE id = $1 "# + let prepareStatementContext = ExtendedQueryContext( + name: name, query: query, bindingDataTypes: [], logger: .psqlTest, promise: promise + ) + + XCTAssertEqual(state.enqueue(task: .extendedQuery(prepareStatementContext)), + .sendParseDescribeSync(name: name, query: query, bindingDataTypes: [])) + XCTAssertEqual(state.parseCompleteReceived(), .wait) + XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait) + + XCTAssertEqual(state.noDataReceived(), + .succeedPreparedStatementCreation(promise, with: nil)) + XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery) + XCTAssertEqual(state.authenticationMessageReceived(.ok), .closeConnectionAndCleanup(.init(action: .close, tasks: [], error: .unexpectedBackendMessage(.authentication(.ok)), closePromise: nil))) } - } diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/PreparedStatementStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/PreparedStatementStateMachineTests.swift new file mode 100644 index 00000000..e35e93f7 --- /dev/null +++ b/Tests/PostgresNIOTests/New/Connection State Machine/PreparedStatementStateMachineTests.swift @@ -0,0 +1,160 @@ +import XCTest +import NIOEmbedded +@testable import PostgresNIO + +class PreparedStatementStateMachineTests: XCTestCase { + func testPrepareAndExecuteStatement() { + let eventLoop = EmbeddedEventLoop() + var stateMachine = PreparedStatementStateMachine() + + let firstPreparedStatement = self.makePreparedStatementContext(eventLoop: eventLoop) + // Initial lookup, the statement hasn't been prepared yet + let lookupAction = stateMachine.lookup(preparedStatement: firstPreparedStatement) + guard case .preparing = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + guard case .prepareStatement = lookupAction else { + XCTFail("State machine returned the wrong action") + return + } + + // Once preparation is complete we transition to a prepared state + let preparationCompleteAction = stateMachine.preparationComplete(name: "test", rowDescription: nil) + guard case .prepared(nil) = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + XCTAssertEqual(preparationCompleteAction.statements.count, 1) + XCTAssertNil(preparationCompleteAction.rowDescription) + firstPreparedStatement.promise.succeed(PSQLRowStream( + source: .noRows(.success(.tag("tag"))), + eventLoop: eventLoop, + logger: .psqlTest + )) + + // Create a new prepared statement + let secondPreparedStatement = self.makePreparedStatementContext(eventLoop: eventLoop) + // The statement is already preparead, lookups tell us to execute it + let secondLookupAction = stateMachine.lookup(preparedStatement: secondPreparedStatement) + guard case .prepared(nil) = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + guard case .executeStatement(nil) = secondLookupAction else { + XCTFail("State machine returned the wrong action") + return + } + secondPreparedStatement.promise.succeed(PSQLRowStream( + source: .noRows(.success(.tag("tag"))), + eventLoop: eventLoop, + logger: .psqlTest + )) + } + + func testPrepareAndExecuteStatementWithError() { + let eventLoop = EmbeddedEventLoop() + var stateMachine = PreparedStatementStateMachine() + + let firstPreparedStatement = self.makePreparedStatementContext(eventLoop: eventLoop) + // Initial lookup, the statement hasn't been prepared yet + let lookupAction = stateMachine.lookup(preparedStatement: firstPreparedStatement) + guard case .preparing = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + guard case .prepareStatement = lookupAction else { + XCTFail("State machine returned the wrong action") + return + } + + // Simulate an error occurring during preparation + let error = PSQLError(code: .server) + let preparationCompleteAction = stateMachine.errorHappened( + name: "test", + error: error + ) + guard case .error = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + XCTAssertEqual(preparationCompleteAction.statements.count, 1) + firstPreparedStatement.promise.fail(error) + + // Create a new prepared statement + let secondPreparedStatement = self.makePreparedStatementContext(eventLoop: eventLoop) + // Ensure that we don't try again to prepare a statement we know will fail + let secondLookupAction = stateMachine.lookup(preparedStatement: secondPreparedStatement) + guard case .error = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + guard case .returnError = secondLookupAction else { + XCTFail("State machine returned the wrong action") + return + } + secondPreparedStatement.promise.fail(error) + } + + func testBatchStatementPreparation() { + let eventLoop = EmbeddedEventLoop() + var stateMachine = PreparedStatementStateMachine() + + let firstPreparedStatement = self.makePreparedStatementContext(eventLoop: eventLoop) + // Initial lookup, the statement hasn't been prepared yet + let lookupAction = stateMachine.lookup(preparedStatement: firstPreparedStatement) + guard case .preparing = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + guard case .prepareStatement = lookupAction else { + XCTFail("State machine returned the wrong action") + return + } + + // A new request comes in before the statement completes + let secondPreparedStatement = self.makePreparedStatementContext(eventLoop: eventLoop) + let secondLookupAction = stateMachine.lookup(preparedStatement: secondPreparedStatement) + guard case .preparing = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + guard case .waitForAlreadyInFlightPreparation = secondLookupAction else { + XCTFail("State machine returned the wrong action") + return + } + + // Once preparation is complete we transition to a prepared state. + // The action tells us to execute both the pending statements. + let preparationCompleteAction = stateMachine.preparationComplete(name: "test", rowDescription: nil) + guard case .prepared(nil) = stateMachine.preparedStatements["test"] else { + XCTFail("State machine in the wrong state") + return + } + XCTAssertEqual(preparationCompleteAction.statements.count, 2) + XCTAssertNil(preparationCompleteAction.rowDescription) + + firstPreparedStatement.promise.succeed(PSQLRowStream( + source: .noRows(.success(.tag("tag"))), + eventLoop: eventLoop, + logger: .psqlTest + )) + secondPreparedStatement.promise.succeed(PSQLRowStream( + source: .noRows(.success(.tag("tag"))), + eventLoop: eventLoop, + logger: .psqlTest + )) + } + + private func makePreparedStatementContext(eventLoop: EmbeddedEventLoop) -> PreparedStatementContext { + let promise = eventLoop.makePromise(of: PSQLRowStream.self) + return PreparedStatementContext( + name: "test", + sql: "INSERT INTO test_table (column1) VALUES (1)", + bindings: PostgresBindings(), + bindingDataTypes: [], + logger: .psqlTest, + promise: promise + ) + } +} diff --git a/Tests/PostgresNIOTests/New/Data/Array+PSQLCodableTests.swift b/Tests/PostgresNIOTests/New/Data/Array+PSQLCodableTests.swift index 0a1da7c6..bfffef52 100644 --- a/Tests/PostgresNIOTests/New/Data/Array+PSQLCodableTests.swift +++ b/Tests/PostgresNIOTests/New/Data/Array+PSQLCodableTests.swift @@ -55,6 +55,26 @@ class Array_PSQLCodableTests: XCTestCase { XCTAssertEqual(UUID.psqlArrayType, .uuidArray) XCTAssertEqual(UUID.psqlType, .uuid) XCTAssertEqual([UUID].psqlType, .uuidArray) + + XCTAssertEqual(Date.psqlArrayType, .timestamptzArray) + XCTAssertEqual(Date.psqlType, .timestamptz) + XCTAssertEqual([Date].psqlType, .timestamptzArray) + + XCTAssertEqual(Range.psqlArrayType, .int4RangeArray) + XCTAssertEqual(Range.psqlType, .int4Range) + XCTAssertEqual([Range].psqlType, .int4RangeArray) + + XCTAssertEqual(ClosedRange.psqlArrayType, .int4RangeArray) + XCTAssertEqual(ClosedRange.psqlType, .int4Range) + XCTAssertEqual([ClosedRange].psqlType, .int4RangeArray) + + XCTAssertEqual(Range.psqlArrayType, .int8RangeArray) + XCTAssertEqual(Range.psqlType, .int8Range) + XCTAssertEqual([Range].psqlType, .int8RangeArray) + + XCTAssertEqual(ClosedRange.psqlArrayType, .int8RangeArray) + XCTAssertEqual(ClosedRange.psqlType, .int8Range) + XCTAssertEqual([ClosedRange].psqlType, .int8RangeArray) } func testStringArrayRoundTrip() { diff --git a/Tests/PostgresNIOTests/New/Data/Date+PSQLCodableTests.swift b/Tests/PostgresNIOTests/New/Data/Date+PSQLCodableTests.swift index b08c2de2..3f406598 100644 --- a/Tests/PostgresNIOTests/New/Data/Date+PSQLCodableTests.swift +++ b/Tests/PostgresNIOTests/New/Data/Date+PSQLCodableTests.swift @@ -14,7 +14,7 @@ class Date_PSQLCodableTests: XCTestCase { var result: Date? XCTAssertNoThrow(result = try Date(from: &buffer, type: .timestamptz, format: .binary, context: .default)) - XCTAssertEqual(value, result) + XCTAssertEqual(value.timeIntervalSince1970, result?.timeIntervalSince1970 ?? 0, accuracy: 0.001) } func testDecodeRandomDate() { @@ -68,7 +68,7 @@ class Date_PSQLCodableTests: XCTestCase { XCTAssertNotNil(lastDate) } - func testDecodeDateFailsWithToMuchData() { + func testDecodeDateFailsWithTooMuchData() { var buffer = ByteBuffer() buffer.writeInteger(Int64(0)) diff --git a/Tests/PostgresNIOTests/New/Data/JSON+PSQLCodableTests.swift b/Tests/PostgresNIOTests/New/Data/JSON+PSQLCodableTests.swift index 858b6ede..52dead6a 100644 --- a/Tests/PostgresNIOTests/New/Data/JSON+PSQLCodableTests.swift +++ b/Tests/PostgresNIOTests/New/Data/JSON+PSQLCodableTests.swift @@ -1,4 +1,5 @@ import XCTest +import Atomics import NIOCore @testable import PostgresNIO @@ -69,11 +70,11 @@ class JSON_PSQLCodableTests: XCTestCase { } func testCustomEncoderIsUsed() { - class TestEncoder: PostgresJSONEncoder { - var encodeHits = 0 + final class TestEncoder: PostgresJSONEncoder { + let encodeHits = ManagedAtomic(0) func encode(_ value: T, into buffer: inout ByteBuffer) throws where T : Encodable { - self.encodeHits += 1 + self.encodeHits.wrappingIncrement(ordering: .relaxed) } func encode(_ value: T) throws -> Data where T : Encodable { @@ -85,6 +86,6 @@ class JSON_PSQLCodableTests: XCTestCase { let encoder = TestEncoder() var buffer = ByteBuffer() XCTAssertNoThrow(try hello.encode(into: &buffer, context: .init(jsonEncoder: encoder))) - XCTAssertEqual(encoder.encodeHits, 1) + XCTAssertEqual(encoder.encodeHits.load(ordering: .relaxed), 1) } } diff --git a/Tests/PostgresNIOTests/New/Data/Range+PSQLCodableTests.swift b/Tests/PostgresNIOTests/New/Data/Range+PSQLCodableTests.swift new file mode 100644 index 00000000..a040c3f4 --- /dev/null +++ b/Tests/PostgresNIOTests/New/Data/Range+PSQLCodableTests.swift @@ -0,0 +1,105 @@ +import XCTest +import NIOCore +@testable import PostgresNIO + +class Range_PSQLCodableTests: XCTestCase { + func testInt32RangeRoundTrip() { + let lowerBound = Int32.min + let upperBound = Int32.max + let value: Range = lowerBound...psqlType, .int4Range) + XCTAssertEqual(buffer.readableBytes, 17) + XCTAssertEqual(buffer.getInteger(at: 0, as: UInt8.self), 2) + XCTAssertEqual(buffer.getInteger(at: 1, as: UInt32.self), 4) + XCTAssertEqual(buffer.getInteger(at: 5, as: Int32.self), lowerBound) + XCTAssertEqual(buffer.getInteger(at: 9, as: UInt32.self), 4) + XCTAssertEqual(buffer.getInteger(at: 13, as: Int32.self), upperBound) + + var result: Range? + XCTAssertNoThrow(result = try Range(from: &buffer, type: .int4Range, format: .binary, context: .default)) + XCTAssertEqual(value, result) + } + + func testInt32ClosedRangeRoundTrip() { + let lowerBound = Int32.min + let upperBound = Int32.max - 1 + let value: ClosedRange = lowerBound...upperBound + + var buffer = ByteBuffer() + value.encode(into: &buffer, context: .default) + XCTAssertEqual(ClosedRange.psqlType, .int4Range) + XCTAssertEqual(buffer.readableBytes, 17) + XCTAssertEqual(buffer.getInteger(at: 0, as: UInt8.self), 6) + XCTAssertEqual(buffer.getInteger(at: 1, as: UInt32.self), 4) + XCTAssertEqual(buffer.getInteger(at: 5, as: Int32.self), lowerBound) + XCTAssertEqual(buffer.getInteger(at: 9, as: UInt32.self), 4) + XCTAssertEqual(buffer.getInteger(at: 13, as: Int32.self), upperBound) + + var result: ClosedRange? + XCTAssertNoThrow(result = try ClosedRange(from: &buffer, type: .int4Range, format: .binary, context: .default)) + XCTAssertEqual(value, result) + } + + func testInt64RangeRoundTrip() { + let lowerBound = Int64.min + let upperBound = Int64.max + let value: Range = lowerBound...psqlType, .int8Range) + XCTAssertEqual(buffer.readableBytes, 25) + XCTAssertEqual(buffer.getInteger(at: 0, as: UInt8.self), 2) + XCTAssertEqual(buffer.getInteger(at: 1, as: UInt32.self), 8) + XCTAssertEqual(buffer.getInteger(at: 5, as: Int64.self), lowerBound) + XCTAssertEqual(buffer.getInteger(at: 13, as: UInt32.self), 8) + XCTAssertEqual(buffer.getInteger(at: 17, as: Int64.self), upperBound) + + var result: Range? + XCTAssertNoThrow(result = try Range(from: &buffer, type: .int8Range, format: .binary, context: .default)) + XCTAssertEqual(value, result) + } + + func testInt64ClosedRangeRoundTrip() { + let lowerBound = Int64.min + let upperBound = Int64.max - 1 + let value: ClosedRange = lowerBound...upperBound + + var buffer = ByteBuffer() + value.encode(into: &buffer, context: .default) + XCTAssertEqual(ClosedRange.psqlType, .int8Range) + XCTAssertEqual(buffer.readableBytes, 25) + XCTAssertEqual(buffer.getInteger(at: 0, as: UInt8.self), 6) + XCTAssertEqual(buffer.getInteger(at: 1, as: UInt32.self), 8) + XCTAssertEqual(buffer.getInteger(at: 5, as: Int64.self), lowerBound) + XCTAssertEqual(buffer.getInteger(at: 13, as: UInt32.self), 8) + XCTAssertEqual(buffer.getInteger(at: 17, as: Int64.self), upperBound) + + var result: ClosedRange? + XCTAssertNoThrow(result = try ClosedRange(from: &buffer, type: .int8Range, format: .binary, context: .default)) + XCTAssertEqual(value, result) + } + + func testInt64RangeDecodeFailureInvalidLength() { + var buffer = ByteBuffer() + buffer.writeInteger(0) + buffer.writeInteger(Int64.random(in: Int64.min...Int64.max)) + buffer.writeInteger(Int64.random(in: Int64.min...Int64.max)) + + XCTAssertThrowsError(try Range(from: &buffer, type: .int8Range, format: .binary, context: .default)) { + XCTAssertEqual($0 as? PostgresDecodingError.Code, .failure) + } + } + + func testInt64RangeDecodeFailureWrongDataType() { + var buffer = ByteBuffer() + (Int64.min...Int64.max).encode(into: &buffer, context: .default) + + XCTAssertThrowsError(try Range(from: &buffer, type: .int8, format: .binary, context: .default)) { + XCTAssertEqual($0 as? PostgresDecodingError.Code, .failure) + } + } +} diff --git a/Tests/PostgresNIOTests/New/Data/String+PSQLCodableTests.swift b/Tests/PostgresNIOTests/New/Data/String+PSQLCodableTests.swift index 614749c1..6ff35130 100644 --- a/Tests/PostgresNIOTests/New/Data/String+PSQLCodableTests.swift +++ b/Tests/PostgresNIOTests/New/Data/String+PSQLCodableTests.swift @@ -20,7 +20,7 @@ class String_PSQLCodableTests: XCTestCase { buffer.writeString(expected) let dataTypes: [PostgresDataType] = [ - .text, .varchar, .name + .text, .varchar, .name, .bpchar ] for dataType in dataTypes { @@ -33,7 +33,7 @@ class String_PSQLCodableTests: XCTestCase { func testDecodeFailureFromInvalidType() { let buffer = ByteBuffer() - let dataTypes: [PostgresDataType] = [.bool, .float4Array, .float8Array, .bpchar] + let dataTypes: [PostgresDataType] = [.bool, .float4Array, .float8Array] for dataType in dataTypes { var loopBuffer = buffer diff --git a/Tests/PostgresNIOTests/New/Extensions/ByteBuffer+Utils.swift b/Tests/PostgresNIOTests/New/Extensions/ByteBuffer+Utils.swift index 71994596..7d073873 100644 --- a/Tests/PostgresNIOTests/New/Extensions/ByteBuffer+Utils.swift +++ b/Tests/PostgresNIOTests/New/Extensions/ByteBuffer+Utils.swift @@ -2,7 +2,10 @@ import NIOCore @testable import PostgresNIO extension ByteBuffer { - + mutating func psqlWriteBackendMessageID(_ messageID: PostgresBackendMessage.ID) { + self.writeInteger(messageID.rawValue) + } + static func backendMessage(id: PostgresBackendMessage.ID, _ payload: (inout ByteBuffer) throws -> ()) rethrows -> ByteBuffer { var byteBuffer = ByteBuffer() try byteBuffer.writeBackendMessage(id: id, payload) diff --git a/Tests/PostgresNIOTests/New/Extensions/ConnectionAction+TestUtils.swift b/Tests/PostgresNIOTests/New/Extensions/ConnectionAction+TestUtils.swift index 72420798..9a1224d8 100644 --- a/Tests/PostgresNIOTests/New/Extensions/ConnectionAction+TestUtils.swift +++ b/Tests/PostgresNIOTests/New/Extensions/ConnectionAction+TestUtils.swift @@ -2,7 +2,8 @@ import class Foundation.JSONEncoder import NIOCore @testable import PostgresNIO -extension ConnectionStateMachine.ConnectionAction: Equatable { +// fully-qualifying all types in the extension has the same effect as adding a `@retroactive` before the protocol +extension PostgresNIO.ConnectionStateMachine.ConnectionAction: Swift.Equatable { public static func == (lhs: Self, rhs: Self) -> Bool { switch (lhs, rhs) { case (.read, read): @@ -25,23 +26,20 @@ extension ConnectionStateMachine.ConnectionAction: Equatable { return lquery == rquery case (.fireEventReadyForQuery, .fireEventReadyForQuery): return true - - case (.succeedQueryNoRowsComming(let lhsContext, let lhsCommandTag), .succeedQueryNoRowsComming(let rhsContext, let rhsCommandTag)): - return lhsContext === rhsContext && lhsCommandTag == rhsCommandTag - case (.succeedQuery(let lhsContext, let lhsRowDescription), .succeedQuery(let rhsContext, let rhsRowDescription)): - return lhsContext === rhsContext && lhsRowDescription == rhsRowDescription - case (.failQuery(let lhsContext, let lhsError, let lhsCleanupContext), .failQuery(let rhsContext, let rhsError, let rhsCleanupContext)): - return lhsContext === rhsContext && lhsError == rhsError && lhsCleanupContext == rhsCleanupContext + case (.succeedQuery(let lhsPromise, let lhsResult), .succeedQuery(let rhsPromise, let rhsResult)): + return lhsPromise.futureResult === rhsPromise.futureResult && lhsResult.value == rhsResult.value + case (.failQuery(let lhsPromise, let lhsError, let lhsCleanupContext), .failQuery(let rhsPromise, let rhsError, let rhsCleanupContext)): + return lhsPromise.futureResult === rhsPromise.futureResult && lhsError == rhsError && lhsCleanupContext == rhsCleanupContext case (.forwardRows(let lhsRows), .forwardRows(let rhsRows)): return lhsRows == rhsRows case (.forwardStreamComplete(let lhsBuffer, let lhsCommandTag), .forwardStreamComplete(let rhsBuffer, let rhsCommandTag)): return lhsBuffer == rhsBuffer && lhsCommandTag == rhsCommandTag case (.forwardStreamError(let lhsError, let lhsRead, let lhsCleanupContext), .forwardStreamError(let rhsError , let rhsRead, let rhsCleanupContext)): return lhsError == rhsError && lhsRead == rhsRead && lhsCleanupContext == rhsCleanupContext - case (.sendParseDescribeSync(let lhsName, let lhsQuery), .sendParseDescribeSync(let rhsName, let rhsQuery)): - return lhsName == rhsName && lhsQuery == rhsQuery - case (.succeedPreparedStatementCreation(let lhsContext, let lhsRowDescription), .succeedPreparedStatementCreation(let rhsContext, let rhsRowDescription)): - return lhsContext === rhsContext && lhsRowDescription == rhsRowDescription + case (.sendParseDescribeSync(let lhsName, let lhsQuery, let lhsDataTypes), .sendParseDescribeSync(let rhsName, let rhsQuery, let rhsDataTypes)): + return lhsName == rhsName && lhsQuery == rhsQuery && lhsDataTypes == rhsDataTypes + case (.succeedPreparedStatementCreation(let lhsPromise, let lhsRowDescription), .succeedPreparedStatementCreation(let rhsPromise, let rhsRowDescription)): + return lhsPromise.futureResult === rhsPromise.futureResult && lhsRowDescription == rhsRowDescription case (.fireChannelInactive, .fireChannelInactive): return true default: @@ -50,7 +48,8 @@ extension ConnectionStateMachine.ConnectionAction: Equatable { } } -extension ConnectionStateMachine.ConnectionAction.CleanUpContext: Equatable { +// fully-qualifying all types in the extension has the same effect as adding a `@retroactive` before the protocol' +extension PostgresNIO.ConnectionStateMachine.ConnectionAction.CleanUpContext: Swift.Equatable { public static func == (lhs: Self, rhs: Self) -> Bool { guard lhs.closePromise?.futureResult === rhs.closePromise?.futureResult else { return false @@ -99,19 +98,19 @@ extension ConnectionStateMachine { } } -extension PSQLError: Equatable { +// fully-qualifying all types in the extension has the same effect as adding a `@retroactive` before the protocol +extension PostgresNIO.PSQLError: Swift.Equatable { public static func == (lhs: PSQLError, rhs: PSQLError) -> Bool { return true } } -extension PSQLTask: Equatable { +// fully-qualifying all types in the extension has the same effect as adding a `@retroactive` before the protocol +extension PostgresNIO.PSQLTask: Swift.Equatable { public static func == (lhs: PSQLTask, rhs: PSQLTask) -> Bool { switch (lhs, rhs) { case (.extendedQuery(let lhs), .extendedQuery(let rhs)): return lhs === rhs - case (.preparedStatement(let lhs), .preparedStatement(let rhs)): - return lhs === rhs case (.closeCommand(let lhs), .closeCommand(let rhs)): return lhs === rhs default: diff --git a/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessage+Equatable.swift b/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessage+Equatable.swift deleted file mode 100644 index c459ffeb..00000000 --- a/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessage+Equatable.swift +++ /dev/null @@ -1,49 +0,0 @@ -@testable import PostgresNIO - -extension PostgresBackendMessage: Equatable { - - public static func ==(lhs: Self, rhs: Self) -> Bool { - switch (lhs, rhs) { - case (.authentication(let lhs), .authentication(let rhs)): - return lhs == rhs - case (.backendKeyData(let lhs), .backendKeyData(let rhs)): - return lhs == rhs - case (.bindComplete, bindComplete): - return true - case (.closeComplete, closeComplete): - return true - case (.commandComplete(let lhs), commandComplete(let rhs)): - return lhs == rhs - case (.dataRow(let lhs), dataRow(let rhs)): - return lhs == rhs - case (.emptyQueryResponse, emptyQueryResponse): - return true - case (.error(let lhs), error(let rhs)): - return lhs == rhs - case (.noData, noData): - return true - case (.notice(let lhs), notice(let rhs)): - return lhs == rhs - case (.notification(let lhs), .notification(let rhs)): - return lhs == rhs - case (.parameterDescription(let lhs), parameterDescription(let rhs)): - return lhs == rhs - case (.parameterStatus(let lhs), parameterStatus(let rhs)): - return lhs == rhs - case (.parseComplete, parseComplete): - return true - case (.portalSuspended, portalSuspended): - return true - case (.readyForQuery(let lhs), readyForQuery(let rhs)): - return lhs == rhs - case (.rowDescription(let lhs), rowDescription(let rhs)): - return lhs == rhs - case (.sslSupported, sslSupported): - return true - case (.sslUnsupported, sslUnsupported): - return true - default: - return false - } - } -} diff --git a/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessageEncoder.swift b/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessageEncoder.swift index eea7dec3..9614bf1e 100644 --- a/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessageEncoder.swift +++ b/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessageEncoder.swift @@ -9,7 +9,7 @@ struct PSQLBackendMessageEncoder: MessageToByteEncoder { /// - parameters: /// - data: The data to encode into a `ByteBuffer`. /// - out: The `ByteBuffer` into which we want to encode. - func encode(data message: PostgresBackendMessage, out buffer: inout ByteBuffer) throws { + func encode(data message: PostgresBackendMessage, out buffer: inout ByteBuffer) { switch message { case .authentication(let authentication): self.encode(messageID: message.id, payload: authentication, into: &buffer) @@ -144,11 +144,7 @@ extension PostgresBackendMessage.Authentication: PSQLMessagePayloadEncodable { buffer.writeInteger(Int32(3)) case .md5(salt: let salt): - buffer.writeInteger(Int32(5)) - buffer.writeInteger(salt.0) - buffer.writeInteger(salt.1) - buffer.writeInteger(salt.2) - buffer.writeInteger(salt.3) + buffer.writeMultipleIntegers(Int32(5), salt) case .scmCredential: buffer.writeInteger(Int32(6)) @@ -261,3 +257,7 @@ extension RowDescription: PSQLMessagePayloadEncodable { } } } + +protocol PSQLMessagePayloadEncodable { + func encode(into buffer: inout ByteBuffer) +} diff --git a/Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift b/Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift index 311c41bd..55ccd0a9 100644 --- a/Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift +++ b/Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift @@ -34,13 +34,13 @@ struct PSQLFrontendMessageDecoder: NIOSingleStepByteToMessageDecoder { switch code { case 80877103: self.isInStartup = true - return .sslRequest(.init()) + return .sslRequest case 196608: var user: String? var database: String? - var options: String? - + var options = [(String, String)]() + while let name = messageSlice.readNullTerminatedString(), messageSlice.readerIndex < finalIndex { let value = messageSlice.readNullTerminatedString() @@ -51,11 +51,10 @@ struct PSQLFrontendMessageDecoder: NIOSingleStepByteToMessageDecoder { case "database": database = value - case "options": - options = value - default: - break + if let value = value { + options.append((name, value)) + } } } @@ -125,17 +124,90 @@ extension PostgresFrontendMessage { static func decode(from buffer: inout ByteBuffer, for messageID: ID) throws -> PostgresFrontendMessage { switch messageID { case .bind: - preconditionFailure("TODO: Unimplemented") + guard let portalName = buffer.readNullTerminatedString() else { + throw PSQLPartialDecodingError.fieldNotDecodable(type: String.self) + } + guard let preparedStatementName = buffer.readNullTerminatedString() else { + throw PSQLPartialDecodingError.fieldNotDecodable(type: String.self) + } + guard let parameterFormatCount = buffer.readInteger(as: UInt16.self) else { + preconditionFailure("TODO: Unimplemented") + } + + let parameterFormats = (0.. ByteBuffer? in + let length = buffer.readInteger(as: UInt32.self) + switch length { + case .some(..<0): + return nil + case .some(0...): + return buffer.readSlice(length: Int(length!)) + default: + preconditionFailure("TODO: Unimplemented") + } + } + + guard let resultColumnFormatCount = buffer.readInteger(as: UInt16.self) else { + preconditionFailure("TODO: Unimplemented") + } + + let resultColumnFormats = (0.. Startup { + return .init(protocolVersion: Self.versionThree, parameters: parameters) + } + + /// The protocol version number. The most significant 16 bits are the major + /// version number (3 for the protocol described here). The least significant + /// 16 bits are the minor version number (0 for the protocol described here). + var protocolVersion: Int32 + + /// The protocol version number is followed by one or more pairs of parameter + /// name and value strings. A zero byte is required as a terminator after + /// the last name/value pair. `user` is required, others are optional. + struct Parameters: Equatable { + enum Replication { + case `true` + case `false` + case database + } + + /// The database user name to connect as. Required; there is no default. + var user: String + + /// The database to connect to. Defaults to the user name. + var database: String? + + /// Command-line arguments for the backend. (This is deprecated in favor + /// of setting individual run-time parameters.) Spaces within this string are + /// considered to separate arguments, unless escaped with a + /// backslash (\); write \\ to represent a literal backslash. + var options: [(String, String)] + + /// Used to connect in streaming replication mode, where a small set of + /// replication commands can be issued instead of SQL statements. Value + /// can be true, false, or database, and the default is false. + var replication: Replication + + static func ==(lhs: Self, rhs: Self) -> Bool { + guard lhs.user == rhs.user + && lhs.database == rhs.database + && lhs.replication == rhs.replication + && lhs.options.count == rhs.options.count + else { + return false + } + + var lhsIterator = lhs.options.makeIterator() + var rhsIterator = rhs.options.makeIterator() + + while let lhsNext = lhsIterator.next(), let rhsNext = rhsIterator.next() { + guard lhsNext.0 == rhsNext.0 && lhsNext.1 == rhsNext.1 else { + return false + } + } + return true + } + + } + + var parameters: Parameters + } + + case bind(Bind) + case cancel(Cancel) + case close(Close) + case describe(Describe) + case execute(Execute) + case flush + case parse(Parse) + case password(Password) + case saslInitialResponse(SASLInitialResponse) + case saslResponse(SASLResponse) + case sslRequest + case sync + case startup(Startup) + case terminate + + enum ID: UInt8, Equatable { + + case bind + case close + case describe + case execute + case flush + case parse + case password + case saslInitialResponse + case saslResponse + case sync + case terminate + + init?(rawValue: UInt8) { + switch rawValue { + case UInt8(ascii: "B"): + self = .bind + case UInt8(ascii: "C"): + self = .close + case UInt8(ascii: "D"): + self = .describe + case UInt8(ascii: "E"): + self = .execute + case UInt8(ascii: "H"): + self = .flush + case UInt8(ascii: "P"): + self = .parse + case UInt8(ascii: "p"): + self = .password + case UInt8(ascii: "p"): + self = .saslInitialResponse + case UInt8(ascii: "p"): + self = .saslResponse + case UInt8(ascii: "S"): + self = .sync + case UInt8(ascii: "X"): + self = .terminate + default: + return nil + } + } + + var rawValue: UInt8 { + switch self { + case .bind: + return UInt8(ascii: "B") + case .close: + return UInt8(ascii: "C") + case .describe: + return UInt8(ascii: "D") + case .execute: + return UInt8(ascii: "E") + case .flush: + return UInt8(ascii: "H") + case .parse: + return UInt8(ascii: "P") + case .password: + return UInt8(ascii: "p") + case .saslInitialResponse: + return UInt8(ascii: "p") + case .saslResponse: + return UInt8(ascii: "p") + case .sync: + return UInt8(ascii: "S") + case .terminate: + return UInt8(ascii: "X") + } + } + } +} + +extension PostgresFrontendMessage { + + var id: ID { + switch self { + case .bind: + return .bind + case .cancel: + preconditionFailure("Cancel messages don't have an identifier") + case .close: + return .close + case .describe: + return .describe + case .execute: + return .execute + case .flush: + return .flush + case .parse: + return .parse + case .password: + return .password + case .saslInitialResponse: + return .saslInitialResponse + case .saslResponse: + return .saslResponse + case .sslRequest: + preconditionFailure("SSL requests don't have an identifier") + case .startup: + preconditionFailure("Startup messages don't have an identifier") + case .sync: + return .sync + case .terminate: + return .terminate + + } + } +} diff --git a/Tests/PostgresNIOTests/New/Messages/AuthenticationTests.swift b/Tests/PostgresNIOTests/New/Messages/AuthenticationTests.swift index 31a21a91..06e39aae 100644 --- a/Tests/PostgresNIOTests/New/Messages/AuthenticationTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/AuthenticationTests.swift @@ -11,35 +11,37 @@ class AuthenticationTests: XCTestCase { let encoder = PSQLBackendMessageEncoder() // add ok - XCTAssertNoThrow(try encoder.encode(data: .authentication(.ok), out: &buffer)) + encoder.encode(data: .authentication(.ok), out: &buffer) expected.append(.authentication(.ok)) // add kerberos - XCTAssertNoThrow(try encoder.encode(data: .authentication(.kerberosV5), out: &buffer)) + encoder.encode(data: .authentication(.kerberosV5), out: &buffer) expected.append(.authentication(.kerberosV5)) // add plaintext - XCTAssertNoThrow(try encoder.encode(data: .authentication(.plaintext), out: &buffer)) + encoder.encode(data: .authentication(.plaintext), out: &buffer) expected.append(.authentication(.plaintext)) // add md5 - XCTAssertNoThrow(try encoder.encode(data: .authentication(.md5(salt: (1, 2, 3, 4))), out: &buffer)) - expected.append(.authentication(.md5(salt: (1, 2, 3, 4)))) - + let salt: UInt32 = 0x01_02_03_04 + encoder.encode(data: .authentication(.md5(salt: salt)), out: &buffer) + expected.append(.authentication(.md5(salt: salt))) + // add scm credential - XCTAssertNoThrow(try encoder.encode(data: .authentication(.scmCredential), out: &buffer)) + encoder.encode(data: .authentication(.scmCredential), out: &buffer) expected.append(.authentication(.scmCredential)) // add gss - XCTAssertNoThrow(try encoder.encode(data: .authentication(.gss), out: &buffer)) + encoder.encode(data: .authentication(.gss), out: &buffer) expected.append(.authentication(.gss)) // add sspi - XCTAssertNoThrow(try encoder.encode(data: .authentication(.sspi), out: &buffer)) + encoder.encode(data: .authentication(.sspi), out: &buffer) expected.append(.authentication(.sspi)) XCTAssertNoThrow(try ByteToMessageDecoderVerifier.verifyDecoder( inputOutputPairs: [(buffer, expected)], - decoderFactory: { PostgresBackendMessageDecoder(hasAlreadyReceivedBytes: false) })) + decoderFactory: { PostgresBackendMessageDecoder(hasAlreadyReceivedBytes: false) } + )) } } diff --git a/Tests/PostgresNIOTests/New/Messages/BindTests.swift b/Tests/PostgresNIOTests/New/Messages/BindTests.swift index 85768b10..d5ec5b30 100644 --- a/Tests/PostgresNIOTests/New/Messages/BindTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/BindTests.swift @@ -5,15 +5,15 @@ import NIOCore class BindTests: XCTestCase { func testEncodeBind() { - let encoder = PSQLFrontendMessageEncoder() var bindings = PostgresBindings() bindings.append("Hello", context: .default) bindings.append("World", context: .default) - var byteBuffer = ByteBuffer() - let bind = PostgresFrontendMessage.Bind(portalName: "", preparedStatementName: "", bind: bindings) - let message = PostgresFrontendMessage.bind(bind) - encoder.encode(data: message, out: &byteBuffer) - + + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + + encoder.bind(portalName: "", preparedStatementName: "", bind: bindings) + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 37) XCTAssertEqual(PostgresFrontendMessage.ID.bind.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), 36) diff --git a/Tests/PostgresNIOTests/New/Messages/CancelTests.swift b/Tests/PostgresNIOTests/New/Messages/CancelTests.swift index c42f1999..5548aae3 100644 --- a/Tests/PostgresNIOTests/New/Messages/CancelTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/CancelTests.swift @@ -5,18 +5,17 @@ import NIOCore class CancelTests: XCTestCase { func testEncodeCancel() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let cancel = PostgresFrontendMessage.Cancel(processID: 1234, secretKey: 4567) - let message = PostgresFrontendMessage.cancel(cancel) - encoder.encode(data: message, out: &byteBuffer) + let processID: Int32 = 1234 + let secretKey: Int32 = 4567 + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.cancel(processID: processID, secretKey: secretKey) + var byteBuffer = encoder.flushBuffer() XCTAssertEqual(byteBuffer.readableBytes, 16) XCTAssertEqual(16, byteBuffer.readInteger(as: Int32.self)) // payload length XCTAssertEqual(80877102, byteBuffer.readInteger(as: Int32.self)) // cancel request code - XCTAssertEqual(cancel.processID, byteBuffer.readInteger(as: Int32.self)) - XCTAssertEqual(cancel.secretKey, byteBuffer.readInteger(as: Int32.self)) + XCTAssertEqual(processID, byteBuffer.readInteger(as: Int32.self)) + XCTAssertEqual(secretKey, byteBuffer.readInteger(as: Int32.self)) XCTAssertEqual(byteBuffer.readableBytes, 0) } - } diff --git a/Tests/PostgresNIOTests/New/Messages/CloseTests.swift b/Tests/PostgresNIOTests/New/Messages/CloseTests.swift index f6a0237b..a8e1cfeb 100644 --- a/Tests/PostgresNIOTests/New/Messages/CloseTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/CloseTests.swift @@ -3,13 +3,11 @@ import NIOCore @testable import PostgresNIO class CloseTests: XCTestCase { - func testEncodeClosePortal() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let message = PostgresFrontendMessage.close(.portal("Hello")) - encoder.encode(data: message, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.closePortal("Hello") + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 12) XCTAssertEqual(PostgresFrontendMessage.ID.close.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(11, byteBuffer.readInteger(as: Int32.self)) @@ -19,11 +17,10 @@ class CloseTests: XCTestCase { } func testEncodeCloseUnnamedStatement() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let message = PostgresFrontendMessage.close(.preparedStatement("")) - encoder.encode(data: message, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.closePreparedStatement("") + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 7) XCTAssertEqual(PostgresFrontendMessage.ID.close.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(6, byteBuffer.readInteger(as: Int32.self)) @@ -31,5 +28,4 @@ class CloseTests: XCTestCase { XCTAssertEqual("", byteBuffer.readNullTerminatedString()) XCTAssertEqual(byteBuffer.readableBytes, 0) } - } diff --git a/Tests/PostgresNIOTests/New/Messages/DataRowTests.swift b/Tests/PostgresNIOTests/New/Messages/DataRowTests.swift index db31b98a..a90d1e93 100644 --- a/Tests/PostgresNIOTests/New/Messages/DataRowTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/DataRowTests.swift @@ -113,7 +113,7 @@ class DataRowTests: XCTestCase { } } -extension DataRow: ExpressibleByArrayLiteral { +extension PostgresNIO.DataRow: Swift.ExpressibleByArrayLiteral { public typealias ArrayLiteralElement = PostgresEncodable public init(arrayLiteral elements: PostgresEncodable...) { diff --git a/Tests/PostgresNIOTests/New/Messages/DescribeTests.swift b/Tests/PostgresNIOTests/New/Messages/DescribeTests.swift index df26f3d7..cb3c745b 100644 --- a/Tests/PostgresNIOTests/New/Messages/DescribeTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/DescribeTests.swift @@ -5,11 +5,10 @@ import NIOCore class DescribeTests: XCTestCase { func testEncodeDescribePortal() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let message = PostgresFrontendMessage.describe(.portal("Hello")) - encoder.encode(data: message, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.describePortal("Hello") + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 12) XCTAssertEqual(PostgresFrontendMessage.ID.describe.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(11, byteBuffer.readInteger(as: Int32.self)) @@ -19,11 +18,10 @@ class DescribeTests: XCTestCase { } func testEncodeDescribeUnnamedStatement() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let message = PostgresFrontendMessage.describe(.preparedStatement("")) - encoder.encode(data: message, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.describePreparedStatement("") + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 7) XCTAssertEqual(PostgresFrontendMessage.ID.describe.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(6, byteBuffer.readInteger(as: Int32.self)) diff --git a/Tests/PostgresNIOTests/New/Messages/ExecuteTests.swift b/Tests/PostgresNIOTests/New/Messages/ExecuteTests.swift index dc5e2767..834ad0dd 100644 --- a/Tests/PostgresNIOTests/New/Messages/ExecuteTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/ExecuteTests.swift @@ -5,11 +5,10 @@ import NIOCore class ExecuteTests: XCTestCase { func testEncodeExecute() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let message = PostgresFrontendMessage.execute(.init(portalName: "", maxNumberOfRows: 0)) - encoder.encode(data: message, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.execute(portalName: "", maxNumberOfRows: 0) + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 10) // 1 (id) + 4 (length) + 1 (empty null terminated string) + 4 (count) XCTAssertEqual(PostgresFrontendMessage.ID.execute.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(9, byteBuffer.readInteger(as: Int32.self)) // length diff --git a/Tests/PostgresNIOTests/New/Messages/ParseTests.swift b/Tests/PostgresNIOTests/New/Messages/ParseTests.swift index 723ad1e6..9f81e4e4 100644 --- a/Tests/PostgresNIOTests/New/Messages/ParseTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/ParseTests.swift @@ -3,18 +3,19 @@ import NIOCore @testable import PostgresNIO class ParseTests: XCTestCase { - func testEncode() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let parse = PostgresFrontendMessage.Parse( - preparedStatementName: "test", - query: "SELECT version()", - parameters: [.bool, .int8, .bytea, .varchar, .text, .uuid, .json, .jsonbArray]) - let message = PostgresFrontendMessage.parse(parse) - encoder.encode(data: message, out: &byteBuffer) + let preparedStatementName = "test" + let query = "SELECT version()" + let parameters: [PostgresDataType] = [.bool, .int8, .bytea, .varchar, .text, .uuid, .json, .jsonbArray] + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.parse( + preparedStatementName: preparedStatementName, + query: query, + parameters: parameters + ) + var byteBuffer = encoder.flushBuffer() - let length: Int = 1 + 4 + (parse.preparedStatementName.count + 1) + (parse.query.count + 1) + 2 + parse.parameters.count * 4 + let length: Int = 1 + 4 + (preparedStatementName.count + 1) + (query.count + 1) + 2 + parameters.count * 4 // 1 id // + 4 length @@ -24,17 +25,11 @@ class ParseTests: XCTestCase { XCTAssertEqual(byteBuffer.readableBytes, length) XCTAssertEqual(byteBuffer.readInteger(as: UInt8.self), PostgresFrontendMessage.ID.parse.rawValue) XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(length - 1)) - XCTAssertEqual(byteBuffer.readNullTerminatedString(), parse.preparedStatementName) - XCTAssertEqual(byteBuffer.readNullTerminatedString(), parse.query) - XCTAssertEqual(byteBuffer.readInteger(as: UInt16.self), UInt16(parse.parameters.count)) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.bool.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.int8.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.bytea.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.varchar.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.text.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.uuid.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.json.rawValue) - XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), PostgresDataType.jsonbArray.rawValue) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), preparedStatementName) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), query) + XCTAssertEqual(byteBuffer.readInteger(as: UInt16.self), UInt16(parameters.count)) + for dataType in parameters { + XCTAssertEqual(byteBuffer.readInteger(as: UInt32.self), dataType.rawValue) + } } - } diff --git a/Tests/PostgresNIOTests/New/Messages/PasswordTests.swift b/Tests/PostgresNIOTests/New/Messages/PasswordTests.swift index 7572d382..4a4833d2 100644 --- a/Tests/PostgresNIOTests/New/Messages/PasswordTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/PasswordTests.swift @@ -5,11 +5,11 @@ import NIOCore class PasswordTests: XCTestCase { func testEncodePassword() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) // md522d085ed8dc3377968dc1c1a40519a2a = "abc123" with salt 1, 2, 3, 4 - let message = PostgresFrontendMessage.password(.init(value: "md522d085ed8dc3377968dc1c1a40519a2a")) - encoder.encode(data: message, out: &byteBuffer) + let password = "md522d085ed8dc3377968dc1c1a40519a2a" + encoder.password(password.utf8) + var byteBuffer = encoder.flushBuffer() let expectedLength = 41 // 1 (id) + 4 (length) + 35 (string) + 1 (null termination) diff --git a/Tests/PostgresNIOTests/New/Messages/SASLInitialResponseTests.swift b/Tests/PostgresNIOTests/New/Messages/SASLInitialResponseTests.swift index 08b3097d..90aa6b34 100644 --- a/Tests/PostgresNIOTests/New/Messages/SASLInitialResponseTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/SASLInitialResponseTests.swift @@ -4,15 +4,14 @@ import NIOCore class SASLInitialResponseTests: XCTestCase { - func testEncodeWithData() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let sasl = PostgresFrontendMessage.SASLInitialResponse( - saslMechanism: "hello", initialData: [0, 1, 2, 3, 4, 5, 6, 7]) - let message = PostgresFrontendMessage.saslInitialResponse(sasl) - encoder.encode(data: message, out: &byteBuffer) + func testEncode() { + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + let saslMechanism = "hello" + let initialData: [UInt8] = [0, 1, 2, 3, 4, 5, 6, 7] + encoder.saslInitialResponse(mechanism: saslMechanism, bytes: initialData) + var byteBuffer = encoder.flushBuffer() - let length: Int = 1 + 4 + (sasl.saslMechanism.count + 1) + 4 + sasl.initialData.count + let length: Int = 1 + 4 + (saslMechanism.count + 1) + 4 + initialData.count // 1 id // + 4 length @@ -23,21 +22,20 @@ class SASLInitialResponseTests: XCTestCase { XCTAssertEqual(byteBuffer.readableBytes, length) XCTAssertEqual(byteBuffer.readInteger(as: UInt8.self), PostgresFrontendMessage.ID.saslInitialResponse.rawValue) XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(length - 1)) - XCTAssertEqual(byteBuffer.readNullTerminatedString(), sasl.saslMechanism) - XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(sasl.initialData.count)) - XCTAssertEqual(byteBuffer.readBytes(length: sasl.initialData.count), sasl.initialData) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), saslMechanism) + XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(initialData.count)) + XCTAssertEqual(byteBuffer.readBytes(length: initialData.count), initialData) XCTAssertEqual(byteBuffer.readableBytes, 0) } func testEncodeWithoutData() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let sasl = PostgresFrontendMessage.SASLInitialResponse( - saslMechanism: "hello", initialData: []) - let message = PostgresFrontendMessage.saslInitialResponse(sasl) - encoder.encode(data: message, out: &byteBuffer) + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + let saslMechanism = "hello" + let initialData: [UInt8] = [] + encoder.saslInitialResponse(mechanism: saslMechanism, bytes: initialData) + var byteBuffer = encoder.flushBuffer() - let length: Int = 1 + 4 + (sasl.saslMechanism.count + 1) + 4 + sasl.initialData.count + let length: Int = 1 + 4 + (saslMechanism.count + 1) + 4 + initialData.count // 1 id // + 4 length @@ -48,8 +46,9 @@ class SASLInitialResponseTests: XCTestCase { XCTAssertEqual(byteBuffer.readableBytes, length) XCTAssertEqual(byteBuffer.readInteger(as: UInt8.self), PostgresFrontendMessage.ID.saslInitialResponse.rawValue) XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(length - 1)) - XCTAssertEqual(byteBuffer.readNullTerminatedString(), sasl.saslMechanism) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), saslMechanism) XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(-1)) + XCTAssertEqual(byteBuffer.readBytes(length: initialData.count), initialData) XCTAssertEqual(byteBuffer.readableBytes, 0) } } diff --git a/Tests/PostgresNIOTests/New/Messages/SASLResponseTests.swift b/Tests/PostgresNIOTests/New/Messages/SASLResponseTests.swift index e148420f..cdb0f10b 100644 --- a/Tests/PostgresNIOTests/New/Messages/SASLResponseTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/SASLResponseTests.swift @@ -5,28 +5,26 @@ import NIOCore class SASLResponseTests: XCTestCase { func testEncodeWithData() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let sasl = PostgresFrontendMessage.SASLResponse(data: [0, 1, 2, 3, 4, 5, 6, 7]) - let message = PostgresFrontendMessage.saslResponse(sasl) - encoder.encode(data: message, out: &byteBuffer) - - let length: Int = 1 + 4 + (sasl.data.count) + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + let data: [UInt8] = [0, 1, 2, 3, 4, 5, 6, 7] + encoder.saslResponse(data) + var byteBuffer = encoder.flushBuffer() + + let length: Int = 1 + 4 + (data.count) XCTAssertEqual(byteBuffer.readableBytes, length) XCTAssertEqual(byteBuffer.readInteger(as: UInt8.self), PostgresFrontendMessage.ID.saslResponse.rawValue) XCTAssertEqual(byteBuffer.readInteger(as: Int32.self), Int32(length - 1)) - XCTAssertEqual(byteBuffer.readBytes(length: sasl.data.count), sasl.data) + XCTAssertEqual(byteBuffer.readBytes(length: data.count), data) XCTAssertEqual(byteBuffer.readableBytes, 0) } func testEncodeWithoutData() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let sasl = PostgresFrontendMessage.SASLResponse(data: []) - let message = PostgresFrontendMessage.saslResponse(sasl) - encoder.encode(data: message, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + let data: [UInt8] = [] + encoder.saslResponse(data) + var byteBuffer = encoder.flushBuffer() + let length: Int = 1 + 4 XCTAssertEqual(byteBuffer.readableBytes, length) diff --git a/Tests/PostgresNIOTests/New/Messages/SSLRequestTests.swift b/Tests/PostgresNIOTests/New/Messages/SSLRequestTests.swift index 9a973f2b..e9e6af81 100644 --- a/Tests/PostgresNIOTests/New/Messages/SSLRequestTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/SSLRequestTests.swift @@ -5,16 +5,14 @@ import NIOCore class SSLRequestTests: XCTestCase { func testSSLRequest() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - let request = PostgresFrontendMessage.SSLRequest() - let message = PostgresFrontendMessage.sslRequest(request) - encoder.encode(data: message, out: &byteBuffer) + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.ssl() + var byteBuffer = encoder.flushBuffer() let byteBufferLength = Int32(byteBuffer.readableBytes) XCTAssertEqual(byteBufferLength, byteBuffer.readInteger()) - XCTAssertEqual(request.code, byteBuffer.readInteger()) - + XCTAssertEqual(PostgresFrontendMessage.SSLRequest.requestCode, byteBuffer.readInteger()) + XCTAssertEqual(byteBuffer.readableBytes, 0) } diff --git a/Tests/PostgresNIOTests/New/Messages/StartupTests.swift b/Tests/PostgresNIOTests/New/Messages/StartupTests.swift index 08a9ee21..5af3bf34 100644 --- a/Tests/PostgresNIOTests/New/Messages/StartupTests.swift +++ b/Tests/PostgresNIOTests/New/Messages/StartupTests.swift @@ -4,45 +4,69 @@ import NIOCore class StartupTests: XCTestCase { - func testStartupMessage() { - let encoder = PSQLFrontendMessageEncoder() + func testStartupMessageWithDatabase() { + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + var byteBuffer = ByteBuffer() + + let user = "test" + let database = "abc123" + + encoder.startup(user: user, database: database, options: []) + byteBuffer = encoder.flushBuffer() + + let byteBufferLength = Int32(byteBuffer.readableBytes) + XCTAssertEqual(byteBufferLength, byteBuffer.readInteger()) + XCTAssertEqual(PostgresFrontendMessage.Startup.versionThree, byteBuffer.readInteger()) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "user") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "test") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "database") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "abc123") + XCTAssertEqual(byteBuffer.readInteger(), UInt8(0)) + + XCTAssertEqual(byteBuffer.readableBytes, 0) + } + + func testStartupMessageWithoutDatabase() { + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + var byteBuffer = ByteBuffer() + + let user = "test" + + encoder.startup(user: user, database: nil, options: []) + byteBuffer = encoder.flushBuffer() + + let byteBufferLength = Int32(byteBuffer.readableBytes) + XCTAssertEqual(byteBufferLength, byteBuffer.readInteger()) + XCTAssertEqual(PostgresFrontendMessage.Startup.versionThree, byteBuffer.readInteger()) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "user") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "test") + XCTAssertEqual(byteBuffer.readInteger(), UInt8(0)) + + XCTAssertEqual(byteBuffer.readableBytes, 0) + } + + func testStartupMessageWithAdditionalOptions() { + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) var byteBuffer = ByteBuffer() - let replicationValues: [PostgresFrontendMessage.Startup.Parameters.Replication] = [ - .`true`, - .`false`, - .database - ] + let user = "test" + let database = "abc123" - for replication in replicationValues { - let parameters = PostgresFrontendMessage.Startup.Parameters( - user: "test", - database: "abc123", - options: "some options", - replication: replication - ) - - let startup = PostgresFrontendMessage.Startup.versionThree(parameters: parameters) - let message = PostgresFrontendMessage.startup(startup) - encoder.encode(data: message, out: &byteBuffer) - - let byteBufferLength = Int32(byteBuffer.readableBytes) - XCTAssertEqual(byteBufferLength, byteBuffer.readInteger()) - XCTAssertEqual(startup.protocolVersion, byteBuffer.readInteger()) - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "user") - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "test") - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "database") - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "abc123") - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "options") - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "some options") - if replication != .false { - XCTAssertEqual(byteBuffer.readNullTerminatedString(), "replication") - XCTAssertEqual(byteBuffer.readNullTerminatedString(), replication.stringValue) - } - XCTAssertEqual(byteBuffer.readInteger(), UInt8(0)) - - XCTAssertEqual(byteBuffer.readableBytes, 0) - } + encoder.startup(user: user, database: database, options: [("some", "options")]) + byteBuffer = encoder.flushBuffer() + + let byteBufferLength = Int32(byteBuffer.readableBytes) + XCTAssertEqual(byteBufferLength, byteBuffer.readInteger()) + XCTAssertEqual(PostgresFrontendMessage.Startup.versionThree, byteBuffer.readInteger()) + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "user") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "test") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "database") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "abc123") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "some") + XCTAssertEqual(byteBuffer.readNullTerminatedString(), "options") + XCTAssertEqual(byteBuffer.readInteger(), UInt8(0)) + + XCTAssertEqual(byteBuffer.readableBytes, 0) } } diff --git a/Tests/PostgresNIOTests/New/PSQLBackendMessageTests.swift b/Tests/PostgresNIOTests/New/PSQLBackendMessageTests.swift index 10e8503a..195c7fb4 100644 --- a/Tests/PostgresNIOTests/New/PSQLBackendMessageTests.swift +++ b/Tests/PostgresNIOTests/New/PSQLBackendMessageTests.swift @@ -256,11 +256,12 @@ class PSQLBackendMessageTests: XCTestCase { } func testDebugDescription() { + let salt: UInt32 = 0x00_01_02_03 XCTAssertEqual("\(PostgresBackendMessage.authentication(.ok))", ".authentication(.ok)") XCTAssertEqual("\(PostgresBackendMessage.authentication(.kerberosV5))", ".authentication(.kerberosV5)") - XCTAssertEqual("\(PostgresBackendMessage.authentication(.md5(salt: (0, 1, 2, 3))))", - ".authentication(.md5(salt: (0, 1, 2, 3)))") + XCTAssertEqual("\(PostgresBackendMessage.authentication(.md5(salt: salt)))", + ".authentication(.md5(salt: \(salt)))") XCTAssertEqual("\(PostgresBackendMessage.authentication(.plaintext))", ".authentication(.plaintext)") XCTAssertEqual("\(PostgresBackendMessage.authentication(.scmCredential))", diff --git a/Tests/PostgresNIOTests/New/PSQLConnectionTests.swift b/Tests/PostgresNIOTests/New/PSQLConnectionTests.swift deleted file mode 100644 index 2a58d4f6..00000000 --- a/Tests/PostgresNIOTests/New/PSQLConnectionTests.swift +++ /dev/null @@ -1,37 +0,0 @@ -import NIOCore -import NIOPosix -import XCTest -import Logging -@testable import PostgresNIO - -class PSQLConnectionTests: XCTestCase { - - func testConnectionFailure() { - // We start a local server and close it immediately to ensure that the port - // number we try to connect to is not used by any other process. - let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) - defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } - - var tempChannel: Channel? - XCTAssertNoThrow(tempChannel = try ServerBootstrap(group: eventLoopGroup) - .bind(to: .init(ipAddress: "127.0.0.1", port: 0)).wait()) - let maybePort = tempChannel?.localAddress?.port - XCTAssertNoThrow(try tempChannel?.close().wait()) - guard let port = maybePort else { - return XCTFail("Could not get port number from temp started server") - } - - let config = PostgresConnection.Configuration( - host: "127.0.0.1", port: port, - username: "postgres", password: "abc123", database: "postgres", - tls: .disable - ) - - var logger = Logger.psqlTest - logger.logLevel = .trace - - XCTAssertThrowsError(try PostgresConnection.connect(on: eventLoopGroup.next(), configuration: config, id: 1, logger: logger).wait()) { - XCTAssertTrue($0 is PSQLError) - } - } -} diff --git a/Tests/PostgresNIOTests/New/PSQLFrontendMessageTests.swift b/Tests/PostgresNIOTests/New/PSQLFrontendMessageTests.swift index 59b69bae..33afbe0d 100644 --- a/Tests/PostgresNIOTests/New/PSQLFrontendMessageTests.swift +++ b/Tests/PostgresNIOTests/New/PSQLFrontendMessageTests.swift @@ -23,30 +23,30 @@ class PSQLFrontendMessageTests: XCTestCase { // MARK: Encoder func testEncodeFlush() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - encoder.encode(data: .flush, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.flush() + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 5) XCTAssertEqual(PostgresFrontendMessage.ID.flush.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(4, byteBuffer.readInteger(as: Int32.self)) // payload length } func testEncodeSync() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - encoder.encode(data: .sync, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.sync() + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 5) XCTAssertEqual(PostgresFrontendMessage.ID.sync.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(4, byteBuffer.readInteger(as: Int32.self)) // payload length } func testEncodeTerminate() { - let encoder = PSQLFrontendMessageEncoder() - var byteBuffer = ByteBuffer() - encoder.encode(data: .terminate, out: &byteBuffer) - + var encoder = PostgresFrontendMessageEncoder(buffer: .init()) + encoder.terminate() + var byteBuffer = encoder.flushBuffer() + XCTAssertEqual(byteBuffer.readableBytes, 5) XCTAssertEqual(PostgresFrontendMessage.ID.terminate.rawValue, byteBuffer.readInteger(as: UInt8.self)) XCTAssertEqual(4, byteBuffer.readInteger(as: Int32.self)) // payload length diff --git a/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift b/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift index f27ff060..65ca26c3 100644 --- a/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift +++ b/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift @@ -1,3 +1,4 @@ +import Atomics import NIOCore import Logging import XCTest @@ -5,69 +6,43 @@ import XCTest import NIOCore import NIOEmbedded -class PSQLRowStreamTests: XCTestCase { +final class PSQLRowStreamTests: XCTestCase { + let logger = Logger(label: "PSQLRowStreamTests") + let eventLoop = EmbeddedEventLoop() + func testEmptyStream() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "INSERT INTO foo bar;", logger: logger, promise: promise - ) - let stream = PSQLRowStream( - rowDescription: [], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .noRows(.success("INSERT 0 1")) + source: .noRows(.success(.tag("INSERT 0 1"))), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertEqual(try stream.all().wait(), []) XCTAssertEqual(stream.commandTag, "INSERT 0 1") } func testFailedStream() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "SELECT * FROM test;", logger: logger, promise: promise - ) - let stream = PSQLRowStream( - rowDescription: [], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .noRows(.failure(PSQLError.connectionClosed)) + source: .noRows(.failure(PSQLError.serverClosedConnection(underlying: nil))), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertThrowsError(try stream.all().wait()) { - XCTAssertEqual($0 as? PSQLError, .connectionClosed) + XCTAssertEqual($0 as? PSQLError, .serverClosedConnection(underlying: nil)) } } func testGetArrayAfterStreamHasFinished() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "SELECT * FROM test;", logger: logger, promise: promise - ) - let dataSource = CountingDataSource() let stream = PSQLRowStream( - rowDescription: [ - self.makeColumnDescription(name: "foo", dataType: .text, format: .binary) - ], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)], + dataSource + ), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertEqual(dataSource.hitCancel, 0) @@ -89,22 +64,15 @@ class PSQLRowStreamTests: XCTestCase { } func testGetArrayBeforeStreamHasFinished() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "SELECT * FROM test;", logger: logger, promise: promise) let dataSource = CountingDataSource() let stream = PSQLRowStream( - rowDescription: [ - self.makeColumnDescription(name: "foo", dataType: .text, format: .binary) - ], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)], + dataSource + ), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertEqual(dataSource.hitCancel, 0) @@ -139,24 +107,15 @@ class PSQLRowStreamTests: XCTestCase { } func testOnRowAfterStreamHasFinished() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "SELECT * FROM test;", logger: logger, promise: promise - ) - let dataSource = CountingDataSource() let stream = PSQLRowStream( - rowDescription: [ - self.makeColumnDescription(name: "foo", dataType: .text, format: .binary) - ], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)], + dataSource + ), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertEqual(dataSource.hitCancel, 0) @@ -170,12 +129,12 @@ class PSQLRowStreamTests: XCTestCase { XCTAssertEqual(dataSource.hitDemand, 0) // attach consumer - var counter = 0 + let counter = ManagedAtomic(0) let future = stream.onRow { row in - XCTAssertEqual(try row.decode(String.self, context: .default), "\(counter)") - counter += 1 + let expected = counter.loadThenWrappingIncrement(ordering: .relaxed) + XCTAssertEqual(try row.decode(String.self, context: .default), "\(expected)") } - XCTAssertEqual(counter, 2) + XCTAssertEqual(counter.load(ordering: .relaxed), 2) XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertNoThrow(try future.wait()) @@ -183,30 +142,23 @@ class PSQLRowStreamTests: XCTestCase { } func testOnRowThrowsErrorOnInitialBatch() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "SELECT * FROM test;", logger: logger, promise: promise - ) - let dataSource = CountingDataSource() let stream = PSQLRowStream( - rowDescription: [ - self.makeColumnDescription(name: "foo", dataType: .text, format: .binary) - ], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)], + dataSource + ), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertEqual(dataSource.hitCancel, 0) stream.receive([ [ByteBuffer(string: "0")], - [ByteBuffer(string: "1")] + [ByteBuffer(string: "1")], + [ByteBuffer(string: "2")], + [ByteBuffer(string: "3")], ]) stream.receive(completion: .success("SELECT 2")) @@ -214,15 +166,15 @@ class PSQLRowStreamTests: XCTestCase { XCTAssertEqual(dataSource.hitDemand, 0) // attach consumer - var counter = 0 + let counter = ManagedAtomic(0) let future = stream.onRow { row in - XCTAssertEqual(try row.decode(String.self, context: .default), "\(counter)") - if counter == 1 { - throw OnRowError(row: counter) + let expected = counter.loadThenWrappingIncrement(ordering: .relaxed) + XCTAssertEqual(try row.decode(String.self, context: .default), "\(expected)") + if expected == 1 { + throw OnRowError(row: expected) } - counter += 1 } - XCTAssertEqual(counter, 1) + XCTAssertEqual(counter.load(ordering: .relaxed), 2) // one more than where we excited, because we already incremented XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertThrowsError(try future.wait()) { @@ -230,26 +182,16 @@ class PSQLRowStreamTests: XCTestCase { } } - func testOnRowBeforeStreamHasFinished() { - let logger = Logger(label: "test") - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - - let queryContext = ExtendedQueryContext( - query: "SELECT * FROM test;", logger: logger, promise: promise - ) - let dataSource = CountingDataSource() let stream = PSQLRowStream( - rowDescription: [ - self.makeColumnDescription(name: "foo", dataType: .text, format: .binary) - ], - queryContext: queryContext, - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)], + dataSource + ), + eventLoop: self.eventLoop, + logger: self.logger ) - promise.succeed(stream) XCTAssertEqual(dataSource.hitDemand, 0) XCTAssertEqual(dataSource.hitCancel, 0) @@ -261,26 +203,26 @@ class PSQLRowStreamTests: XCTestCase { XCTAssertEqual(dataSource.hitDemand, 0, "Before we have a consumer demand is not signaled") // attach consumer - var counter = 0 + let counter = ManagedAtomic(0) let future = stream.onRow { row in - XCTAssertEqual(try row.decode(String.self, context: .default), "\(counter)") - counter += 1 + let expected = counter.loadThenWrappingIncrement(ordering: .relaxed) + XCTAssertEqual(try row.decode(String.self, context: .default), "\(expected)") } - XCTAssertEqual(counter, 2) + XCTAssertEqual(counter.load(ordering: .relaxed), 2) XCTAssertEqual(dataSource.hitDemand, 1) stream.receive([ [ByteBuffer(string: "2")], [ByteBuffer(string: "3")] ]) - XCTAssertEqual(counter, 4) + XCTAssertEqual(counter.load(ordering: .relaxed), 4) XCTAssertEqual(dataSource.hitDemand, 2) stream.receive([ [ByteBuffer(string: "4")], [ByteBuffer(string: "5")] ]) - XCTAssertEqual(counter, 6) + XCTAssertEqual(counter.load(ordering: .relaxed), 6) XCTAssertEqual(dataSource.hitDemand, 3) stream.receive(completion: .success("SELECT 6")) diff --git a/Tests/PostgresNIOTests/New/PostgresChannelHandlerTests.swift b/Tests/PostgresNIOTests/New/PostgresChannelHandlerTests.swift index 7ab0ce30..206f38a3 100644 --- a/Tests/PostgresNIOTests/New/PostgresChannelHandlerTests.swift +++ b/Tests/PostgresNIOTests/New/PostgresChannelHandlerTests.swift @@ -6,19 +6,28 @@ import NIOEmbedded @testable import PostgresNIO class PostgresChannelHandlerTests: XCTestCase { - + + var eventLoop: EmbeddedEventLoop! + + override func setUp() { + super.setUp() + self.eventLoop = EmbeddedEventLoop() + } + // MARK: Startup func testHandlerAddedWithoutSSL() { let config = self.testConnectionConfiguration() - let handler = PostgresChannelHandler(configuration: config, configureSSLCallback: nil) + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop, configureSSLCallback: nil) let embedded = EmbeddedChannel(handlers: [ ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), ReverseMessageToByteHandler(PSQLBackendMessageEncoder()), handler - ]) - defer { XCTAssertNoThrow(try embedded.finish()) } - + ], loop: self.eventLoop) + defer { + XCTAssertNoThrow({ try embedded.finish() }) + } + var maybeMessage: PostgresFrontendMessage? XCTAssertNoThrow(embedded.connect(to: try .init(ipAddress: "0.0.0.0", port: 5432), promise: nil)) XCTAssertNoThrow(maybeMessage = try embedded.readOutbound(as: PostgresFrontendMessage.self)) @@ -28,9 +37,8 @@ class PostgresChannelHandlerTests: XCTestCase { XCTAssertEqual(startup.parameters.user, config.username) XCTAssertEqual(startup.parameters.database, config.database) - XCTAssertEqual(startup.parameters.options, nil) - XCTAssertEqual(startup.parameters.replication, .false) - + XCTAssert(startup.parameters.options.isEmpty) + XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.authentication(.ok))) XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.backendKeyData(.init(processID: 1234, secretKey: 5678)))) XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.readyForQuery(.idle))) @@ -40,24 +48,18 @@ class PostgresChannelHandlerTests: XCTestCase { var config = self.testConnectionConfiguration() XCTAssertNoThrow(config.tls = .require(try NIOSSLContext(configuration: .makeClientConfiguration()))) var addSSLCallbackIsHit = false - let handler = PostgresChannelHandler(configuration: config) { channel in + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop) { channel, _ in addSSLCallbackIsHit = true } let embedded = EmbeddedChannel(handlers: [ ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), ReverseMessageToByteHandler(PSQLBackendMessageEncoder()), handler - ]) - - var maybeMessage: PostgresFrontendMessage? + ], loop: self.eventLoop) + XCTAssertNoThrow(embedded.connect(to: try .init(ipAddress: "0.0.0.0", port: 5432), promise: nil)) - XCTAssertNoThrow(maybeMessage = try embedded.readOutbound(as: PostgresFrontendMessage.self)) - guard case .sslRequest(let request) = maybeMessage else { - return XCTFail("Unexpected message") - } - - XCTAssertEqual(request.code, 80877103) - + XCTAssertEqual(.sslRequest, try embedded.readOutbound(as: PostgresFrontendMessage.self)) + XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.sslSupported)) // a NIOSSLHandler has been added, after it SSL had been negotiated @@ -82,7 +84,7 @@ class PostgresChannelHandlerTests: XCTestCase { var config = self.testConnectionConfiguration() XCTAssertNoThrow(config.tls = .require(try NIOSSLContext(configuration: .makeClientConfiguration()))) var addSSLCallbackIsHit = false - let handler = PostgresChannelHandler(configuration: config) { channel in + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop) { channel, _ in addSSLCallbackIsHit = true } let eventHandler = TestEventHandler() @@ -90,16 +92,10 @@ class PostgresChannelHandlerTests: XCTestCase { ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), handler, eventHandler - ]) + ], loop: self.eventLoop) - var maybeMessage: PostgresFrontendMessage? XCTAssertNoThrow(embedded.connect(to: try .init(ipAddress: "0.0.0.0", port: 5432), promise: nil)) - XCTAssertNoThrow(maybeMessage = try embedded.readOutbound(as: PostgresFrontendMessage.self)) - guard case .sslRequest(let request) = maybeMessage else { - return XCTFail("Unexpected message") - } - - XCTAssertEqual(request.code, 80877103) + XCTAssertEqual(.sslRequest, try embedded.readOutbound(as: PostgresFrontendMessage.self)) var responseBuffer = ByteBuffer() responseBuffer.writeInteger(UInt8(ascii: "S")) @@ -118,7 +114,7 @@ class PostgresChannelHandlerTests: XCTestCase { func testSSLUnsupportedClosesConnection() throws { let config = self.testConnectionConfiguration(tls: .require(try NIOSSLContext(configuration: .makeClientConfiguration()))) - let handler = PostgresChannelHandler(configuration: config) { channel in + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop) { channel, _ in XCTFail("This callback should never be exectuded") throw PSQLError.sslUnsupported } @@ -126,15 +122,15 @@ class PostgresChannelHandlerTests: XCTestCase { ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), ReverseMessageToByteHandler(PSQLBackendMessageEncoder()), handler - ]) + ], loop: self.eventLoop) let eventHandler = TestEventHandler() - try embedded.pipeline.addHandler(eventHandler, position: .last).wait() + try embedded.pipeline.syncOperations.addHandler(eventHandler, position: .last) embedded.connect(to: try .init(ipAddress: "0.0.0.0", port: 5432), promise: nil) XCTAssertTrue(embedded.isActive) // read the ssl request message - XCTAssertEqual(try embedded.readOutbound(as: PostgresFrontendMessage.self), .sslRequest(.init())) + XCTAssertEqual(try embedded.readOutbound(as: PostgresFrontendMessage.self), .sslRequest) try embedded.writeInbound(PostgresBackendMessage.sslUnsupported) // the event handler should have seen an error @@ -154,22 +150,22 @@ class PostgresChannelHandlerTests: XCTestCase { database: config.database ) let state = ConnectionStateMachine(.waitingToStartAuthentication) - let handler = PostgresChannelHandler(configuration: config, state: state, configureSSLCallback: nil) + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop, state: state, configureSSLCallback: nil) let embedded = EmbeddedChannel(handlers: [ ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), - ReverseMessageToByteHandler(PSQLBackendMessageEncoder()), handler - ]) - + ], loop: self.eventLoop) + embedded.triggerUserOutboundEvent(PSQLOutgoingEvent.authenticate(authContext), promise: nil) XCTAssertEqual(try embedded.readOutbound(as: PostgresFrontendMessage.self), .startup(.versionThree(parameters: authContext.toStartupParameters()))) + let salt: UInt32 = 0x00_01_02_03 + + let encoder = PSQLBackendMessageEncoder() + var byteBuffer = ByteBuffer() + encoder.encode(data: .authentication(.md5(salt: salt)), out: &byteBuffer) + XCTAssertNoThrow(try embedded.writeInbound(byteBuffer)) - XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.authentication(.md5(salt: (0,1,2,3))))) - - var message: PostgresFrontendMessage? - XCTAssertNoThrow(message = try embedded.readOutbound(as: PostgresFrontendMessage.self)) - - XCTAssertEqual(message, .password(.init(value: "md522d085ed8dc3377968dc1c1a40519a2a"))) + XCTAssertEqual(try embedded.readOutbound(as: PostgresFrontendMessage.self), .password(.init(value: "md522d085ed8dc3377968dc1c1a40519a2a"))) } func testRunAuthenticateCleartext() { @@ -181,24 +177,57 @@ class PostgresChannelHandlerTests: XCTestCase { database: config.database ) let state = ConnectionStateMachine(.waitingToStartAuthentication) - let handler = PostgresChannelHandler(configuration: config, state: state, configureSSLCallback: nil) + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop, state: state, configureSSLCallback: nil) let embedded = EmbeddedChannel(handlers: [ ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), ReverseMessageToByteHandler(PSQLBackendMessageEncoder()), handler - ]) - + ], loop: self.eventLoop) + embedded.triggerUserOutboundEvent(PSQLOutgoingEvent.authenticate(authContext), promise: nil) XCTAssertEqual(try embedded.readOutbound(as: PostgresFrontendMessage.self), .startup(.versionThree(parameters: authContext.toStartupParameters()))) XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.authentication(.plaintext))) - - var message: PostgresFrontendMessage? - XCTAssertNoThrow(message = try embedded.readOutbound(as: PostgresFrontendMessage.self)) - - XCTAssertEqual(message, .password(.init(value: password))) + XCTAssertEqual(try embedded.readOutbound(as: PostgresFrontendMessage.self), .password(.init(value: password))) } - + + func testHandlerThatSendsMultipleWrongMessages() { + let config = self.testConnectionConfiguration() + let handler = PostgresChannelHandler(configuration: config, eventLoop: self.eventLoop, configureSSLCallback: nil) + let embedded = EmbeddedChannel(handlers: [ + ReverseByteToMessageHandler(PSQLFrontendMessageDecoder()), + handler + ], loop: self.eventLoop) + + var maybeMessage: PostgresFrontendMessage? + XCTAssertNoThrow(embedded.connect(to: try .init(ipAddress: "0.0.0.0", port: 5432), promise: nil)) + XCTAssertNoThrow(maybeMessage = try embedded.readOutbound(as: PostgresFrontendMessage.self)) + guard case .startup(let startup) = maybeMessage else { + return XCTFail("Unexpected message") + } + + XCTAssertEqual(startup.parameters.user, config.username) + XCTAssertEqual(startup.parameters.database, config.database) + XCTAssert(startup.parameters.options.isEmpty) + XCTAssertEqual(startup.parameters.replication, .false) + + var buffer = ByteBuffer() + buffer.writeMultipleIntegers(UInt8(ascii: "R"), UInt32(8), Int32(0)) + buffer.writeMultipleIntegers(UInt8(ascii: "K"), UInt32(12), Int32(1234), Int32(5678)) + buffer.writeMultipleIntegers(UInt8(ascii: "Z"), UInt32(5), UInt8(ascii: "I")) + XCTAssertNoThrow(try embedded.writeInbound(buffer)) + XCTAssertTrue(embedded.isActive) + + buffer.clear() + buffer.writeMultipleIntegers(UInt8(ascii: "Z"), UInt32(5), UInt8(ascii: "I")) + buffer.writeMultipleIntegers(UInt8(ascii: "Z"), UInt32(5), UInt8(ascii: "I")) + buffer.writeMultipleIntegers(UInt8(ascii: "Z"), UInt32(5), UInt8(ascii: "I")) + buffer.writeMultipleIntegers(UInt8(ascii: "Z"), UInt32(5), UInt8(ascii: "I")) + + XCTAssertThrowsError(try embedded.writeInbound(buffer)) + XCTAssertFalse(embedded.isActive) + } + // MARK: Helpers func testConnectionConfiguration( @@ -246,3 +275,14 @@ class TestEventHandler: ChannelInboundHandler { self.events.append(psqlEvent) } } + +extension AuthContext { + func toStartupParameters() -> PostgresFrontendMessage.Startup.Parameters { + PostgresFrontendMessage.Startup.Parameters( + user: self.username, + database: self.database, + options: self.additionalParameters, + replication: .false + ) + } +} diff --git a/Tests/PostgresNIOTests/New/PostgresConnectionTests.swift b/Tests/PostgresNIOTests/New/PostgresConnectionTests.swift new file mode 100644 index 00000000..d0f8e2b0 --- /dev/null +++ b/Tests/PostgresNIOTests/New/PostgresConnectionTests.swift @@ -0,0 +1,771 @@ +import NIOCore +import NIOPosix +import NIOEmbedded +import XCTest +import Logging +@testable import PostgresNIO + +class PostgresConnectionTests: XCTestCase { + + let logger = Logger(label: "PostgresConnectionTests") + + func testConnectionFailure() { + // We start a local server and close it immediately to ensure that the port + // number we try to connect to is not used by any other process. + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + + var tempChannel: Channel? + XCTAssertNoThrow(tempChannel = try ServerBootstrap(group: eventLoopGroup) + .bind(to: .init(ipAddress: "127.0.0.1", port: 0)).wait()) + let maybePort = tempChannel?.localAddress?.port + XCTAssertNoThrow(try tempChannel?.close().wait()) + guard let port = maybePort else { + return XCTFail("Could not get port number from temp started server") + } + + let config = PostgresConnection.Configuration( + host: "127.0.0.1", port: port, + username: "postgres", password: "abc123", database: "postgres", + tls: .disable + ) + + var logger = Logger.psqlTest + logger.logLevel = .trace + + XCTAssertThrowsError(try PostgresConnection.connect(on: eventLoopGroup.next(), configuration: config, id: 1, logger: logger).wait()) { + XCTAssertTrue($0 is PSQLError) + } + } + + func testOptionsAreSentOnTheWire() async throws { + let eventLoop = NIOAsyncTestingEventLoop() + let channel = try await NIOAsyncTestingChannel(loop: eventLoop) { channel in + try channel.pipeline.syncOperations.addHandlers(ReverseByteToMessageHandler(PSQLFrontendMessageDecoder())) + try channel.pipeline.syncOperations.addHandlers(ReverseMessageToByteHandler(PSQLBackendMessageEncoder())) + } + try await channel.connect(to: .makeAddressResolvingHost("localhost", port: 5432)) + + let configuration = { + var config = PostgresConnection.Configuration( + establishedChannel: channel, + username: "username", + password: "postgres", + database: "database" + ) + config.options.additionalStartupParameters = [ + ("DateStyle", "ISO, MDY"), + ("application_name", "postgres-nio-test"), + ("server_encoding", "UTF8"), + ("integer_datetimes", "on"), + ("client_encoding", "UTF8"), + ("TimeZone", "Etc/UTC"), + ("is_superuser", "on"), + ("server_version", "13.1 (Debian 13.1-1.pgdg100+1)"), + ("session_authorization", "postgres"), + ("IntervalStyle", "postgres"), + ("standard_conforming_strings", "on") + ] + return config + }() + + async let connectionPromise = PostgresConnection.connect(on: eventLoop, configuration: configuration, id: 1, logger: .psqlTest) + let message = try await channel.waitForOutboundWrite(as: PostgresFrontendMessage.self) + XCTAssertEqual(message, .startup(.versionThree(parameters: .init(user: "username", database: "database", options: configuration.options.additionalStartupParameters, replication: .false)))) + try await channel.writeInbound(PostgresBackendMessage.authentication(.ok)) + try await channel.writeInbound(PostgresBackendMessage.backendKeyData(.init(processID: 1234, secretKey: 5678))) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + let connection = try await connectionPromise + try await connection.close() + } + + func testSimpleListen() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + let events = try await connection.listen("foo") + for try await event in events { + XCTAssertEqual(event.payload, "wooohooo") + break + } + } + + let listenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(listenMessage.parse.query, #"LISTEN "foo";"#) + + try await channel.writeInbound(PostgresBackendMessage.parseComplete) + try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: []))) + try await channel.writeInbound(PostgresBackendMessage.noData) + try await channel.writeInbound(PostgresBackendMessage.bindComplete) + try await channel.writeInbound(PostgresBackendMessage.commandComplete("LISTEN")) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + try await channel.writeInbound(PostgresBackendMessage.notification(.init(backendPID: 12, channel: "foo", payload: "wooohooo"))) + + let unlistenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(unlistenMessage.parse.query, #"UNLISTEN "foo";"#) + + try await channel.writeInbound(PostgresBackendMessage.parseComplete) + try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: []))) + try await channel.writeInbound(PostgresBackendMessage.noData) + try await channel.writeInbound(PostgresBackendMessage.bindComplete) + try await channel.writeInbound(PostgresBackendMessage.commandComplete("UNLISTEN")) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + switch await taskGroup.nextResult()! { + case .success: + break + case .failure(let failure): + XCTFail("Unexpected error: \(failure)") + } + } + } + + func testSimpleListenDoesNotUnlistenIfThereIsAnotherSubscriber() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { taskGroup in + taskGroup.addTask { + let events = try await connection.listen("foo") + for try await event in events { + XCTAssertEqual(event.payload, "wooohooo") + break + } + } + + taskGroup.addTask { + let events = try await connection.listen("foo") + var counter = 0 + loop: for try await event in events { + defer { counter += 1 } + switch counter { + case 0: + XCTAssertEqual(event.payload, "wooohooo") + case 1: + XCTAssertEqual(event.payload, "wooohooo2") + break loop + default: + XCTFail("Unexpected message: \(event)") + } + } + } + + let listenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(listenMessage.parse.query, #"LISTEN "foo";"#) + + try await channel.writeInbound(PostgresBackendMessage.parseComplete) + try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: []))) + try await channel.writeInbound(PostgresBackendMessage.noData) + try await channel.writeInbound(PostgresBackendMessage.bindComplete) + try await channel.writeInbound(PostgresBackendMessage.commandComplete("LISTEN")) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + try await channel.writeInbound(PostgresBackendMessage.notification(.init(backendPID: 12, channel: "foo", payload: "wooohooo"))) + try await channel.writeInbound(PostgresBackendMessage.notification(.init(backendPID: 12, channel: "foo", payload: "wooohooo2"))) + + let unlistenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(unlistenMessage.parse.query, #"UNLISTEN "foo";"#) + + try await channel.writeInbound(PostgresBackendMessage.parseComplete) + try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: []))) + try await channel.writeInbound(PostgresBackendMessage.noData) + try await channel.writeInbound(PostgresBackendMessage.bindComplete) + try await channel.writeInbound(PostgresBackendMessage.commandComplete("UNLISTEN")) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + switch await taskGroup.nextResult()! { + case .success: + break + case .failure(let failure): + XCTFail("Unexpected error: \(failure)") + } + } + } + + func testSimpleListenConnectionDrops() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { [logger] taskGroup in + taskGroup.addTask { + let events = try await connection.listen("foo") + var iterator = events.makeAsyncIterator() + let first = try await iterator.next() + XCTAssertEqual(first?.payload, "wooohooo") + do { + _ = try await iterator.next() + XCTFail("Did not expect to not throw") + } catch { + logger.error("error", metadata: ["error": "\(error)"]) + } + } + + let listenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(listenMessage.parse.query, #"LISTEN "foo";"#) + + try await channel.writeInbound(PostgresBackendMessage.parseComplete) + try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: []))) + try await channel.writeInbound(PostgresBackendMessage.noData) + try await channel.writeInbound(PostgresBackendMessage.bindComplete) + try await channel.writeInbound(PostgresBackendMessage.commandComplete("LISTEN")) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + try await channel.writeInbound(PostgresBackendMessage.notification(.init(backendPID: 12, channel: "foo", payload: "wooohooo"))) + struct MyWeirdError: Error {} + channel.pipeline.fireErrorCaught(MyWeirdError()) + + switch await taskGroup.nextResult()! { + case .success: + break + case .failure(let failure): + XCTFail("Unexpected error: \(failure)") + } + } + } + + func testCloseGracefullyClosesWhenInternalQueueIsEmpty() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + try await withThrowingTaskGroup(of: Void.self) { [logger] taskGroup async throws -> () in + for _ in 1...2 { + taskGroup.addTask { + let rows = try await connection.query("SELECT 1;", logger: logger) + var iterator = rows.decode(Int.self).makeAsyncIterator() + let first = try await iterator.next() + XCTAssertEqual(first, 1) + let second = try await iterator.next() + XCTAssertNil(second) + } + } + + for i in 0...1 { + let listenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(listenMessage.parse.query, "SELECT 1;") + + if i == 0 { + taskGroup.addTask { + try await connection.closeGracefully() + } + } + + try await channel.writeInbound(PostgresBackendMessage.parseComplete) + try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: []))) + let intDescription = RowDescription.Column( + name: "", + tableOID: 0, + columnAttributeNumber: 0, + dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary + ) + try await channel.writeInbound(PostgresBackendMessage.rowDescription(.init(columns: [intDescription]))) + try await channel.testingEventLoop.executeInContext { channel.read() } + try await channel.writeInbound(PostgresBackendMessage.bindComplete) + try await channel.testingEventLoop.executeInContext { channel.read() } + try await channel.writeInbound(PostgresBackendMessage.dataRow([Int(1)])) + try await channel.testingEventLoop.executeInContext { channel.read() } + try await channel.writeInbound(PostgresBackendMessage.commandComplete("SELECT 1 1")) + try await channel.testingEventLoop.executeInContext { channel.read() } + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + } + + let terminate = try await channel.waitForOutboundWrite(as: PostgresFrontendMessage.self) + XCTAssertEqual(terminate, .terminate) + try await channel.closeFuture.get() + XCTAssertEqual(channel.isActive, false) + + while let taskResult = await taskGroup.nextResult() { + switch taskResult { + case .success: + break + case .failure(let failure): + XCTFail("Unexpected error: \(failure)") + } + } + } + } + + func testCloseClosesImmediatly() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { [logger] taskGroup async throws -> () in + for _ in 1...2 { + taskGroup.addTask { + try await connection.query("SELECT 1;", logger: logger) + } + } + + let listenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(listenMessage.parse.query, "SELECT 1;") + + async let close: () = connection.close() + + try await channel.closeFuture.get() + XCTAssertEqual(channel.isActive, false) + + try await close + + while let taskResult = await taskGroup.nextResult() { + switch taskResult { + case .success: + XCTFail("Expected queries to fail") + case .failure(let failure): + guard let error = failure as? PSQLError else { + return XCTFail("Unexpected error type: \(failure)") + } + XCTAssertEqual(error.code, .clientClosedConnection) + } + } + } + } + + func testIfServerJustClosesTheErrorReflectsThat() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + let logger = self.logger + + async let response = try await connection.query("SELECT 1;", logger: logger) + + let listenMessage = try await channel.waitForUnpreparedRequest() + XCTAssertEqual(listenMessage.parse.query, "SELECT 1;") + + try await channel.testingEventLoop.executeInContext { channel.pipeline.fireChannelInactive() } + try await channel.testingEventLoop.executeInContext { channel.pipeline.fireChannelUnregistered() } + + do { + _ = try await response + XCTFail("Expected to throw") + } catch { + XCTAssertEqual((error as? PSQLError)?.code, .serverClosedConnection) + } + + // retry on same connection + + do { + _ = try await connection.query("SELECT 1;", logger: self.logger) + XCTFail("Expected to throw") + } catch { + XCTAssertEqual((error as? PSQLError)?.code, .serverClosedConnection) + } + } + + struct TestPrepareStatement: PostgresPreparedStatement { + static let sql = "SELECT datname FROM pg_stat_activity WHERE state = $1" + typealias Row = String + + var state: String + + func makeBindings() -> PostgresBindings { + var bindings = PostgresBindings() + bindings.append(.init(string: self.state)) + return bindings + } + + func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row { + try row.decode(Row.self) + } + } + + func testPreparedStatement() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "active") + let result = try await connection.execute(preparedStatement, logger: .psqlTest) + var rows = 0 + for try await database in result { + rows += 1 + XCTAssertEqual("test_database", database) + } + XCTAssertEqual(rows, 1) + } + + let prepareRequest = try await channel.waitForPrepareRequest() + XCTAssertEqual(prepareRequest.parse.query, "SELECT datname FROM pg_stat_activity WHERE state = $1") + XCTAssertEqual(prepareRequest.parse.parameters.first, .text) + guard case .preparedStatement(let name) = prepareRequest.describe else { + fatalError("Describe should contain a prepared statement") + } + XCTAssertEqual(name, String(reflecting: TestPrepareStatement.self)) + + try await channel.sendPrepareResponse( + parameterDescription: .init(dataTypes: [ + PostgresDataType.text + ]), + rowDescription: .init(columns: [ + .init( + name: "datname", + tableOID: 12222, + columnAttributeNumber: 2, + dataType: .name, + dataTypeSize: 64, + dataTypeModifier: -1, + format: .text + ) + ]) + ) + + let preparedRequest = try await channel.waitForPreparedRequest() + XCTAssertEqual(preparedRequest.bind.preparedStatementName, String(reflecting: TestPrepareStatement.self)) + XCTAssertEqual(preparedRequest.bind.parameters.count, 1) + XCTAssertEqual(preparedRequest.bind.resultColumnFormats, [.binary]) + + try await channel.sendPreparedResponse( + dataRows: [ + ["test_database"] + ], + commandTag: TestPrepareStatement.sql + ) + } + } + + func testWeDontCrashOnUnexpectedChannelEvents() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + enum MyEvent { + case pleaseDontCrash + } + channel.pipeline.fireUserInboundEventTriggered(MyEvent.pleaseDontCrash) + try await connection.close() + } + + func testSerialExecutionOfSamePreparedStatement() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in + // Send the same prepared statement twice, but with different parameters. + // Send one first and wait to send the other request until preparation is complete + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "active") + let result = try await connection.execute(preparedStatement, logger: .psqlTest) + var rows = 0 + for try await database in result { + rows += 1 + XCTAssertEqual("test_database", database) + } + XCTAssertEqual(rows, 1) + } + + let prepareRequest = try await channel.waitForPrepareRequest() + XCTAssertEqual(prepareRequest.parse.query, "SELECT datname FROM pg_stat_activity WHERE state = $1") + XCTAssertEqual(prepareRequest.parse.parameters.first, .text) + guard case .preparedStatement(let name) = prepareRequest.describe else { + fatalError("Describe should contain a prepared statement") + } + XCTAssertEqual(name, String(reflecting: TestPrepareStatement.self)) + + try await channel.sendPrepareResponse( + parameterDescription: .init(dataTypes: [ + PostgresDataType.text + ]), + rowDescription: .init(columns: [ + .init( + name: "datname", + tableOID: 12222, + columnAttributeNumber: 2, + dataType: .name, + dataTypeSize: 64, + dataTypeModifier: -1, + format: .text + ) + ]) + ) + + let preparedRequest1 = try await channel.waitForPreparedRequest() + var buffer = preparedRequest1.bind.parameters[0]! + let parameter1 = buffer.readString(length: buffer.readableBytes)! + XCTAssertEqual(parameter1, "active") + try await channel.sendPreparedResponse( + dataRows: [ + ["test_database"] + ], + commandTag: TestPrepareStatement.sql + ) + + // Now that the statement has been prepared and executed, send another request that will only get executed + // without preparation + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "idle") + let result = try await connection.execute(preparedStatement, logger: .psqlTest) + var rows = 0 + for try await database in result { + rows += 1 + XCTAssertEqual("test_database", database) + } + XCTAssertEqual(rows, 1) + } + + let preparedRequest2 = try await channel.waitForPreparedRequest() + buffer = preparedRequest2.bind.parameters[0]! + let parameter2 = buffer.readString(length: buffer.readableBytes)! + XCTAssertEqual(parameter2, "idle") + try await channel.sendPreparedResponse( + dataRows: [ + ["test_database"] + ], + commandTag: TestPrepareStatement.sql + ) + // Ensure we received and responded to both the requests + let parameters = [parameter1, parameter2] + XCTAssert(parameters.contains("active")) + XCTAssert(parameters.contains("idle")) + } + } + + func testStatementPreparationOnlyHappensOnceWithConcurrentRequests() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in + // Send the same prepared statement twice, but with different parameters. + // Let them race to tests that requests and responses aren't mixed up + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "active") + let result = try await connection.execute(preparedStatement, logger: .psqlTest) + var rows = 0 + for try await database in result { + rows += 1 + XCTAssertEqual("test_database_active", database) + } + XCTAssertEqual(rows, 1) + } + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "idle") + let result = try await connection.execute(preparedStatement, logger: .psqlTest) + var rows = 0 + for try await database in result { + rows += 1 + XCTAssertEqual("test_database_idle", database) + } + XCTAssertEqual(rows, 1) + } + + // The channel deduplicates prepare requests, we're going to see only one of them + let prepareRequest = try await channel.waitForPrepareRequest() + XCTAssertEqual(prepareRequest.parse.query, "SELECT datname FROM pg_stat_activity WHERE state = $1") + XCTAssertEqual(prepareRequest.parse.parameters.first, .text) + guard case .preparedStatement(let name) = prepareRequest.describe else { + fatalError("Describe should contain a prepared statement") + } + XCTAssertEqual(name, String(reflecting: TestPrepareStatement.self)) + + try await channel.sendPrepareResponse( + parameterDescription: .init(dataTypes: [ + PostgresDataType.text + ]), + rowDescription: .init(columns: [ + .init( + name: "datname", + tableOID: 12222, + columnAttributeNumber: 2, + dataType: .name, + dataTypeSize: 64, + dataTypeModifier: -1, + format: .text + ) + ]) + ) + + // Now both the tasks have their statements prepared. + // We should see both of their execute requests coming in, the order is nondeterministic + let preparedRequest1 = try await channel.waitForPreparedRequest() + var buffer = preparedRequest1.bind.parameters[0]! + let parameter1 = buffer.readString(length: buffer.readableBytes)! + try await channel.sendPreparedResponse( + dataRows: [ + ["test_database_\(parameter1)"] + ], + commandTag: TestPrepareStatement.sql + ) + let preparedRequest2 = try await channel.waitForPreparedRequest() + buffer = preparedRequest2.bind.parameters[0]! + let parameter2 = buffer.readString(length: buffer.readableBytes)! + try await channel.sendPreparedResponse( + dataRows: [ + ["test_database_\(parameter2)"] + ], + commandTag: TestPrepareStatement.sql + ) + // Ensure we received and responded to both the requests + let parameters = [parameter1, parameter2] + XCTAssert(parameters.contains("active")) + XCTAssert(parameters.contains("idle")) + } + } + + func testStatementPreparationFailure() async throws { + let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel() + + try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in + // Send the same prepared statement twice, but with different parameters. + // Send one first and wait to send the other request until preparation is complete + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "active") + do { + _ = try await connection.execute(preparedStatement, logger: .psqlTest) + XCTFail("Was supposed to fail") + } catch { + XCTAssert(error is PSQLError) + } + } + + let prepareRequest = try await channel.waitForPrepareRequest() + XCTAssertEqual(prepareRequest.parse.query, "SELECT datname FROM pg_stat_activity WHERE state = $1") + XCTAssertEqual(prepareRequest.parse.parameters.first, .text) + guard case .preparedStatement(let name) = prepareRequest.describe else { + fatalError("Describe should contain a prepared statement") + } + XCTAssertEqual(name, String(reflecting: TestPrepareStatement.self)) + + // Respond with an error taking care to return a SQLSTATE that isn't + // going to get the connection closed. + try await channel.writeInbound(PostgresBackendMessage.error(.init(fields: [ + .sqlState : "26000" // invalid_sql_statement_name + ]))) + try await channel.testingEventLoop.executeInContext { channel.read() } + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + try await channel.testingEventLoop.executeInContext { channel.read() } + + + // Send another requests with the same prepared statement, which should fail straight + // away without any interaction with the server + taskGroup.addTask { + let preparedStatement = TestPrepareStatement(state: "idle") + do { + _ = try await connection.execute(preparedStatement, logger: .psqlTest) + XCTFail("Was supposed to fail") + } catch { + XCTAssert(error is PSQLError) + } + } + } + } + + func makeTestConnectionWithAsyncTestingChannel() async throws -> (PostgresConnection, NIOAsyncTestingChannel) { + let eventLoop = NIOAsyncTestingEventLoop() + let channel = try await NIOAsyncTestingChannel(loop: eventLoop) { channel in + try channel.pipeline.syncOperations.addHandlers(ReverseByteToMessageHandler(PSQLFrontendMessageDecoder())) + try channel.pipeline.syncOperations.addHandlers(ReverseMessageToByteHandler(PSQLBackendMessageEncoder())) + } + try await channel.connect(to: .makeAddressResolvingHost("localhost", port: 5432)) + + let configuration = PostgresConnection.Configuration( + establishedChannel: channel, + username: "username", + password: "postgres", + database: "database" + ) + + let logger = self.logger + async let connectionPromise = PostgresConnection.connect(on: eventLoop, configuration: configuration, id: 1, logger: logger) + let message = try await channel.waitForOutboundWrite(as: PostgresFrontendMessage.self) + XCTAssertEqual(message, .startup(.versionThree(parameters: .init(user: "username", database: "database", options: [], replication: .false)))) + try await channel.writeInbound(PostgresBackendMessage.authentication(.ok)) + try await channel.writeInbound(PostgresBackendMessage.backendKeyData(.init(processID: 1234, secretKey: 5678))) + try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + + let connection = try await connectionPromise + + self.addTeardownBlock { + try await connection.close() + } + + return (connection, channel) + } +} + +extension NIOAsyncTestingChannel { + + func waitForUnpreparedRequest() async throws -> UnpreparedRequest { + let parse = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let describe = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let bind = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let execute = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let sync = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + + guard case .parse(let parse) = parse, + case .describe(let describe) = describe, + case .bind(let bind) = bind, + case .execute(let execute) = execute, + case .sync = sync + else { + fatalError() + } + + return UnpreparedRequest(parse: parse, describe: describe, bind: bind, execute: execute) + } + + func waitForPrepareRequest() async throws -> PrepareRequest { + let parse = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let describe = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let sync = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + + guard case .parse(let parse) = parse, + case .describe(let describe) = describe, + case .sync = sync + else { + fatalError("Unexpected message") + } + + return PrepareRequest(parse: parse, describe: describe) + } + + func sendPrepareResponse( + parameterDescription: PostgresBackendMessage.ParameterDescription, + rowDescription: RowDescription + ) async throws { + try await self.writeInbound(PostgresBackendMessage.parseComplete) + try await self.testingEventLoop.executeInContext { self.read() } + try await self.writeInbound(PostgresBackendMessage.parameterDescription(parameterDescription)) + try await self.testingEventLoop.executeInContext { self.read() } + try await self.writeInbound(PostgresBackendMessage.rowDescription(rowDescription)) + try await self.testingEventLoop.executeInContext { self.read() } + try await self.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + try await self.testingEventLoop.executeInContext { self.read() } + } + + func waitForPreparedRequest() async throws -> PreparedRequest { + let bind = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let execute = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + let sync = try await self.waitForOutboundWrite(as: PostgresFrontendMessage.self) + + guard case .bind(let bind) = bind, + case .execute(let execute) = execute, + case .sync = sync + else { + fatalError() + } + + return PreparedRequest(bind: bind, execute: execute) + } + + func sendPreparedResponse( + dataRows: [DataRow], + commandTag: String + ) async throws { + try await self.writeInbound(PostgresBackendMessage.bindComplete) + try await self.testingEventLoop.executeInContext { self.read() } + for dataRow in dataRows { + try await self.writeInbound(PostgresBackendMessage.dataRow(dataRow)) + } + try await self.testingEventLoop.executeInContext { self.read() } + try await self.writeInbound(PostgresBackendMessage.commandComplete(commandTag)) + try await self.testingEventLoop.executeInContext { self.read() } + try await self.writeInbound(PostgresBackendMessage.readyForQuery(.idle)) + try await self.testingEventLoop.executeInContext { self.read() } + } +} + +struct UnpreparedRequest { + var parse: PostgresFrontendMessage.Parse + var describe: PostgresFrontendMessage.Describe + var bind: PostgresFrontendMessage.Bind + var execute: PostgresFrontendMessage.Execute +} + +struct PrepareRequest { + var parse: PostgresFrontendMessage.Parse + var describe: PostgresFrontendMessage.Describe +} + +struct PreparedRequest { + var bind: PostgresFrontendMessage.Bind + var execute: PostgresFrontendMessage.Execute +} diff --git a/Tests/PostgresNIOTests/New/PostgresErrorTests.swift b/Tests/PostgresNIOTests/New/PostgresErrorTests.swift index 639d6b5e..33df5439 100644 --- a/Tests/PostgresNIOTests/New/PostgresErrorTests.swift +++ b/Tests/PostgresNIOTests/New/PostgresErrorTests.swift @@ -2,6 +2,55 @@ import XCTest import NIOCore +final class PSQLErrorTests: XCTestCase { + func testPostgresBindingsDescription() { + let testBinds1 = PostgresBindings(capacity: 0) + var testBinds2 = PostgresBindings(capacity: 1) + var testBinds3 = PostgresBindings(capacity: 2) + testBinds2.append(1, context: .default) + testBinds3.appendUnprotected(1, context: .default) + testBinds3.appendUnprotected("foo", context: .default) + testBinds3.append("secret", context: .default) + + XCTAssertEqual(String(describing: testBinds1), "[]") + XCTAssertEqual(String(reflecting: testBinds1), "[]") + XCTAssertEqual(String(describing: testBinds2), "[****]") + XCTAssertEqual(String(reflecting: testBinds2), "[(****; BIGINT; format: binary)]") + XCTAssertEqual(String(describing: testBinds3), #"[1, "foo", ****]"#) + XCTAssertEqual(String(reflecting: testBinds3), #"[(1; BIGINT; format: binary), ("foo"; TEXT; format: binary), (****; TEXT; format: binary)]"#) + } + + func testPostgresQueryDescription() { + let testBinds1 = PostgresBindings(capacity: 0) + var testBinds2 = PostgresBindings(capacity: 1) + testBinds2.append(1, context: .default) + let testQuery1 = PostgresQuery(unsafeSQL: "TEST QUERY") + let testQuery2 = PostgresQuery(unsafeSQL: "TEST QUERY", binds: testBinds1) + let testQuery3 = PostgresQuery(unsafeSQL: "TEST QUERY", binds: testBinds2) + + XCTAssertEqual(String(describing: testQuery1), "TEST QUERY []") + XCTAssertEqual(String(reflecting: testQuery1), "PostgresQuery(sql: TEST QUERY, binds: [])") + XCTAssertEqual(String(describing: testQuery2), "TEST QUERY []") + XCTAssertEqual(String(reflecting: testQuery2), "PostgresQuery(sql: TEST QUERY, binds: [])") + XCTAssertEqual(String(describing: testQuery3), "TEST QUERY [****]") + XCTAssertEqual(String(reflecting: testQuery3), "PostgresQuery(sql: TEST QUERY, binds: [(****; BIGINT; format: binary)])") + } + + func testPSQLErrorDescription() { + var error1 = PSQLError.server(.init(fields: [.localizedSeverity: "ERROR", .severity: "ERROR", .sqlState: "00000", .message: "Test message", .detail: "More test message", .hint: "It's a test, that's your hint", .position: "1", .schemaName: "testsch", .tableName: "testtab", .columnName: "testcol", .dataTypeName: "testtyp", .constraintName: "testcon", .file: #fileID, .line: "0", .routine: #function])) + var testBinds = PostgresBindings(capacity: 1) + testBinds.append(1, context: .default) + error1.query = .init(unsafeSQL: "TEST QUERY", binds: testBinds) + + XCTAssertEqual(String(describing: error1), """ + PSQLError – Generic description to prevent accidental leakage of sensitive data. For debugging details, use `String(reflecting: error)`. + """) + XCTAssertEqual(String(reflecting: error1), """ + PSQLError(code: server, serverInfo: [sqlState: 00000, detail: More test message, file: PostgresNIOTests/PostgresErrorTests.swift, hint: It's a test, that's your hint, line: 0, message: Test message, position: 1, routine: testPSQLErrorDescription(), localizedSeverity: ERROR, severity: ERROR, columnName: testcol, dataTypeName: testtyp, constraintName: testcon, schemaName: testsch, tableName: testtab], query: PostgresQuery(sql: TEST QUERY, binds: [(****; BIGINT; format: binary)])) + """) + } +} + final class PostgresDecodingErrorTests: XCTestCase { func testPostgresDecodingErrorEquality() { let error1 = PostgresDecodingError( @@ -59,9 +108,13 @@ final class PostgresDecodingErrorTests: XCTestCase { ) // Plain description - XCTAssertEqual(String(describing: error1), "Database error") - XCTAssertEqual(String(describing: error2), "Database error") - + XCTAssertEqual(String(describing: error1), """ + PostgresDecodingError – Generic description to prevent accidental leakage of sensitive data. For debugging details, use `String(reflecting: error)`. + """) + XCTAssertEqual(String(describing: error2), """ + PostgresDecodingError – Generic description to prevent accidental leakage of sensitive data. For debugging details, use `String(reflecting: error)`. + """) + // Extended debugDescription XCTAssertEqual(String(reflecting: error1), """ PostgresDecodingError(code: typeMismatch,\ diff --git a/Tests/PostgresNIOTests/New/PostgresQueryTests.swift b/Tests/PostgresNIOTests/New/PostgresQueryTests.swift index f50d414a..4930f0c4 100644 --- a/Tests/PostgresNIOTests/New/PostgresQueryTests.swift +++ b/Tests/PostgresNIOTests/New/PostgresQueryTests.swift @@ -31,6 +31,27 @@ final class PostgresQueryTests: XCTestCase { XCTAssertEqual(query.binds.bytes, expected) } + func testStringInterpolationWithDynamicType() { + let type = PostgresDataType(16435) + let format = PostgresFormat.binary + let dynamicString = DynamicString(value: "Hello world", psqlType: type, psqlFormat: format) + + let query: PostgresQuery = """ + INSERT INTO foo (dynamicType) SET (\(dynamicString)); + """ + + XCTAssertEqual(query.sql, "INSERT INTO foo (dynamicType) SET ($1);") + + var expectedBindsBytes = ByteBuffer() + expectedBindsBytes.writeInteger(Int32(dynamicString.value.utf8.count)) + expectedBindsBytes.writeString(dynamicString.value) + + let expectedMetadata: [PostgresBindings.Metadata] = [.init(dataType: type, format: format, protected: true)] + + XCTAssertEqual(query.binds.bytes, expectedBindsBytes) + XCTAssertEqual(query.binds.metadata, expectedMetadata) + } + func testStringInterpolationWithCustomJSONEncoder() { struct Foo: Codable, PostgresCodable { var helloWorld: String @@ -89,3 +110,19 @@ final class PostgresQueryTests: XCTestCase { XCTAssertEqual(query.binds.bytes, expected) } } + +extension PostgresQueryTests { + struct DynamicString: PostgresDynamicTypeEncodable { + let value: String + + var psqlType: PostgresDataType + var psqlFormat: PostgresFormat + + func encode( + into byteBuffer: inout ByteBuffer, + context: PostgresNIO.PostgresEncodingContext + ) where JSONEncoder: PostgresJSONEncoder { + byteBuffer.writeString(value) + } + } +} diff --git a/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift b/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift index e1fdad11..9d662252 100644 --- a/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift +++ b/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift @@ -1,27 +1,27 @@ import Atomics import NIOEmbedded -import Dispatch +import NIOPosix import XCTest @testable import PostgresNIO import NIOCore import Logging final class PostgresRowSequenceTests: XCTestCase { + let logger = Logger(label: "PSQLRowStreamTests") func testBackpressureWorks() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) let rowSequence = stream.asyncSequence() XCTAssertEqual(dataSource.requestCount, 0) @@ -38,20 +38,20 @@ final class PostgresRowSequenceTests: XCTestCase { XCTAssertNil(empty) } + func testCancellationWorksWhileIterating() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) let rowSequence = stream.asyncSequence() XCTAssertEqual(dataSource.requestCount, 0) @@ -60,7 +60,7 @@ final class PostgresRowSequenceTests: XCTestCase { var counter = 0 for try await row in rowSequence { - XCTAssertEqual(try row.decode(Int.self, context: .default), counter) + XCTAssertEqual(try row.decode(Int.self), counter) counter += 1 if counter == 64 { @@ -72,19 +72,18 @@ final class PostgresRowSequenceTests: XCTestCase { } func testCancellationWorksBeforeIterating() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) let rowSequence = stream.asyncSequence() XCTAssertEqual(dataSource.requestCount, 0) @@ -99,19 +98,18 @@ final class PostgresRowSequenceTests: XCTestCase { } func testDroppingTheSequenceCancelsTheSource() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) var rowSequence: PostgresRowSequence? = stream.asyncSequence() rowSequence = nil @@ -121,19 +119,18 @@ final class PostgresRowSequenceTests: XCTestCase { } func testStreamBasedOnCompletedQuery() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) let rowSequence = stream.asyncSequence() let dataRows: [DataRow] = (0..<128).map { [ByteBuffer(integer: Int64($0))] } @@ -142,7 +139,7 @@ final class PostgresRowSequenceTests: XCTestCase { var counter = 0 for try await row in rowSequence { - XCTAssertEqual(try row.decode(Int.self, context: .default), counter) + XCTAssertEqual(try row.decode(Int.self), counter) counter += 1 } @@ -150,19 +147,18 @@ final class PostgresRowSequenceTests: XCTestCase { } func testStreamIfInitializedWithAllData() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) let dataRows: [DataRow] = (0..<128).map { [ByteBuffer(integer: Int64($0))] } stream.receive(dataRows) @@ -172,7 +168,7 @@ final class PostgresRowSequenceTests: XCTestCase { var counter = 0 for try await row in rowSequence { - XCTAssertEqual(try row.decode(Int.self, context: .default), counter) + XCTAssertEqual(try row.decode(Int.self), counter) counter += 1 } @@ -180,21 +176,20 @@ final class PostgresRowSequenceTests: XCTestCase { } func testStreamIfInitializedWithError() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) - stream.receive(completion: .failure(PSQLError.connectionClosed)) + stream.receive(completion: .failure(PSQLError.serverClosedConnection(underlying: nil))) let rowSequence = stream.asyncSequence() @@ -205,37 +200,36 @@ final class PostgresRowSequenceTests: XCTestCase { } XCTFail("Expected that an error was thrown before.") } catch { - XCTAssertEqual(error as? PSQLError, .connectionClosed) + XCTAssertEqual(error as? PSQLError, .serverClosedConnection(underlying: nil)) } } func testSucceedingRowContinuationsWorks() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let eventLoop = NIOSingletons.posixEventLoopGroup.next() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), eventLoop: eventLoop, - rowSource: .stream(dataSource) + logger: self.logger ) - promise.succeed(stream) - let rowSequence = stream.asyncSequence() + let rowSequence = try await eventLoop.submit { stream.asyncSequence() }.get() var rowIterator = rowSequence.makeAsyncIterator() - DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(1)) { + eventLoop.scheduleTask(in: .seconds(1)) { let dataRows: [DataRow] = (0..<1).map { [ByteBuffer(integer: Int64($0))] } stream.receive(dataRows) } let row1 = try await rowIterator.next() - XCTAssertEqual(try row1?.decode(Int.self, context: .default), 0) + XCTAssertEqual(try row1?.decode(Int.self), 0) - DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(1)) { + eventLoop.scheduleTask(in: .seconds(1)) { stream.receive(completion: .success("SELECT 1")) } @@ -244,57 +238,55 @@ final class PostgresRowSequenceTests: XCTestCase { } func testFailingRowContinuationsWorks() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let eventLoop = NIOSingletons.posixEventLoopGroup.next() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), eventLoop: eventLoop, - rowSource: .stream(dataSource) + logger: self.logger ) - promise.succeed(stream) - let rowSequence = stream.asyncSequence() + let rowSequence = try await eventLoop.submit { stream.asyncSequence() }.get() var rowIterator = rowSequence.makeAsyncIterator() - DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(1)) { + eventLoop.scheduleTask(in: .seconds(1)) { let dataRows: [DataRow] = (0..<1).map { [ByteBuffer(integer: Int64($0))] } stream.receive(dataRows) } let row1 = try await rowIterator.next() - XCTAssertEqual(try row1?.decode(Int.self, context: .default), 0) + XCTAssertEqual(try row1?.decode(Int.self), 0) - DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(1)) { - stream.receive(completion: .failure(PSQLError.connectionClosed)) + eventLoop.scheduleTask(in: .seconds(1)) { + stream.receive(completion: .failure(PSQLError.serverClosedConnection(underlying: nil))) } do { _ = try await rowIterator.next() XCTFail("Expected that an error was thrown before.") } catch { - XCTAssertEqual(error as? PSQLError, .connectionClosed) + XCTAssertEqual(error as? PSQLError, .serverClosedConnection(underlying: nil)) } } func testAdaptiveRowBufferShrinksAndGrows() async throws { - let eventLoop = EmbeddedEventLoop() - let promise = eventLoop.makePromise(of: PSQLRowStream.self) - let logger = Logger(label: "test") let dataSource = MockRowDataSource() + let embeddedEventLoop = EmbeddedEventLoop() let stream = PSQLRowStream( - rowDescription: [ - .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) - ], - queryContext: .init(query: "SELECT * FROM foo", logger: logger, promise: promise), - eventLoop: eventLoop, - rowSource: .stream(dataSource) + source: .stream( + [ + .init(name: "test", tableOID: 0, columnAttributeNumber: 0, dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary) + ], + dataSource + ), + eventLoop: embeddedEventLoop, + logger: self.logger ) - promise.succeed(stream) let initialDataRows: [DataRow] = (0..(_ value: T) throws -> Data where T : Encodable { - self.didEncode = true + self.counter.wrappingIncrement(ordering: .relaxed) return try JSONEncoder().encode(value) } } @@ -21,14 +22,16 @@ class PostgresJSONCodingTests: XCTestCase { var bar: Int } let customJSONEncoder = CustomJSONEncoder() + XCTAssertEqual(customJSONEncoder.counter.load(ordering: .relaxed), 0) PostgresNIO._defaultJSONEncoder = customJSONEncoder XCTAssertNoThrow(try PostgresData(json: Object(foo: 1, bar: 2))) - XCTAssert(customJSONEncoder.didEncode) + XCTAssertEqual(customJSONEncoder.counter.load(ordering: .relaxed), 1) let customJSONBEncoder = CustomJSONEncoder() + XCTAssertEqual(customJSONBEncoder.counter.load(ordering: .relaxed), 0) PostgresNIO._defaultJSONEncoder = customJSONBEncoder XCTAssertNoThrow(try PostgresData(json: Object(foo: 1, bar: 2))) - XCTAssert(customJSONBEncoder.didEncode) + XCTAssertEqual(customJSONBEncoder.counter.load(ordering: .relaxed), 1) } // https://github.com/vapor/postgres-nio/issues/126 @@ -38,9 +41,9 @@ class PostgresJSONCodingTests: XCTestCase { PostgresNIO._defaultJSONDecoder = previousDefaultJSONDecoder } final class CustomJSONDecoder: PostgresJSONDecoder { - var didDecode = false + let counter = ManagedAtomic(0) func decode(_ type: T.Type, from data: Data) throws -> T where T : Decodable { - self.didDecode = true + self.counter.wrappingIncrement(ordering: .relaxed) return try JSONDecoder().decode(type, from: data) } } @@ -49,13 +52,15 @@ class PostgresJSONCodingTests: XCTestCase { var bar: Int } let customJSONDecoder = CustomJSONDecoder() + XCTAssertEqual(customJSONDecoder.counter.load(ordering: .relaxed), 0) PostgresNIO._defaultJSONDecoder = customJSONDecoder XCTAssertNoThrow(try PostgresData(json: Object(foo: 1, bar: 2)).json(as: Object.self)) - XCTAssert(customJSONDecoder.didDecode) + XCTAssertEqual(customJSONDecoder.counter.load(ordering: .relaxed), 1) let customJSONBDecoder = CustomJSONDecoder() + XCTAssertEqual(customJSONBDecoder.counter.load(ordering: .relaxed), 0) PostgresNIO._defaultJSONDecoder = customJSONBDecoder XCTAssertNoThrow(try PostgresData(json: Object(foo: 1, bar: 2)).json(as: Object.self)) - XCTAssert(customJSONBDecoder.didDecode) + XCTAssertEqual(customJSONBDecoder.counter.load(ordering: .relaxed), 1) } } diff --git a/docker-compose.yml b/docker-compose.yml index 600bdc99..3eff4249 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -10,6 +10,12 @@ x-shared-config: &shared_config - 5432:5432 services: + psql-16: + image: postgres:16 + <<: *shared_config + psql-15: + image: postgres:15 + <<: *shared_config psql-14: image: postgres:14 <<: *shared_config