From 6ed4a9a64f6f9ff822d78dcc95c52c7e07ff8911 Mon Sep 17 00:00:00 2001 From: Jacob Quinn Date: Sat, 26 Sep 2026 09:11:06 -0600 Subject: [PATCH 1/2] Add scoped connection pooling with session reset --- docs/src/index.md | 43 ++++ docs/src/migration.md | 6 +- src/MySQL.jl | 1 + src/connection.jl | 4 + src/pool.jl | 239 ++++++++++++++++++ test/mysql_trim_workload.jl | 58 ++++- test/pool_workflow.jl | 144 +++++++++++ test/protocol/live_tests.jl | 2 + test/protocol/pool_tests.jl | 472 ++++++++++++++++++++++++++++++++++++ test/protocol/runtests.jl | 1 + 10 files changed, 962 insertions(+), 8 deletions(-) create mode 100644 src/pool.jl create mode 100644 test/pool_workflow.jl create mode 100644 test/protocol/pool_tests.jl diff --git a/docs/src/index.md b/docs/src/index.md index 6d9d928..718fd9f 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -78,6 +78,46 @@ DBInterface.close!(stmt) `MySQL.load(table, conn, name)` creates a table from a Tables.jl source's schema and inserts its rows in batches. Close the connection with `DBInterface.close!(conn)`. +## Scoped connection pooling + +`MySQL.ConnectionPool` creates connections on demand, up to `limit`. Use +`MySQL.with_connection` to keep related operations on one session: + +```julia +pool = MySQL.ConnectionPool("localhost", "user", "password"; db="mydb", limit=8) +try + id = MySQL.with_connection(pool) do conn + DBInterface.execute(conn, "INSERT INTO users (name) VALUES ('alice')") + first(DBInterface.execute(conn, "SELECT LAST_INSERT_ID()"))[1] + end +finally + close(pool) +end +``` + +Before a connection is reused, the pool drains unread results, resets the server session +(including rolling back unfinished transactions), invalidates old statement and cursor +generations, restores the configured database and utf8mb4 settings, and repeats +`init_command`. Use that option only for repeatable session setup. These steps add network +round trips; pooling avoids a new connection and authentication, not cleanup. With no +configured `db`, a session that selected a database is replaced because a reset cannot +clear that selection. Servers without session-reset support also require replacement. +Consume results inside the callback to observe SQL errors; unread responses are discarded +by the driver's ordinary draining behavior. + +Finish all work on the borrowed connection before the callback returns, including work +started in other tasks. Do not keep using the connection, prepared statements, or +unfinished streamed results afterward. Owned buffered results and materialized values +remain readable. The connection is not a revocable lease: an escaped prepared statement +may still auto-reprepare under the ordinary driver rules. For nested work, pass the +already borrowed connection; another lease uses separate capacity and can block. + +An uncaught callback error retires the connection and is rethrown unchanged. A reset +failure also retires it and propagates; callbacks are never replayed. A local open check +does not prove the server is reachable. `close(pool)` rejects new leases, wakes waiting +callers, and closes idle connections. It lets already acquired leases finish, even if +their callbacks start after a racing close, then closes their sessions on return. + ## Connection options `DBInterface.connect(MySQL.Connection, host, user, password; kw...)` accepts the keywords @@ -257,6 +297,9 @@ MySQL.Connection MySQL.Statement MySQL.Cursor MySQL.ConnectOptions +MySQL.ConnectionPool +MySQL.with_connection +Base.close(::MySQL.ConnectionPool) DBInterface.connect DBInterface.close! DBInterface.execute diff --git a/docs/src/migration.md b/docs/src/migration.md index c9fed22..efd5bfd 100644 --- a/docs/src/migration.md +++ b/docs/src/migration.md @@ -198,8 +198,10 @@ schema-aware sinks make. The runtime-schema conveniences (`Tables.columntable`, `row.name` access, `MySQL.load`) build columns from runtime `Type` values and are not statically resolvable; use them from regular Julia. A custom `local_infile_handler` (and the `IO` it returns) is dispatched dynamically and works in a trimmed binary only if its -methods were compiled in. The trim workload also exercises abandoned-connection reaping, -but leaves all connection deadlines disabled; it does not validate timeout behavior. +methods were compiled in. The trim workload also exercises abandoned-connection reaping +and sequential scoped pool leases with statically known callbacks, session cleanup, +callback-error retirement, and closed-pool rejection. It leaves all connection deadlines +disabled; it does not validate timeout behavior. `connect_timeout` hangs in trimmed executables with Reseau ≤ 1.4.1 — its deadline-armed dial parked on tasks a trimmed build never ran, fixed in [Reseau #151](https://github.com/JuliaServices/Reseau.jl/pull/151). diff --git a/src/MySQL.jl b/src/MySQL.jl index c422eaa..71ef4f5 100644 --- a/src/MySQL.jl +++ b/src/MySQL.jl @@ -39,6 +39,7 @@ include("connection.jl") include("cursor.jl") include("statement.jl") include("load.jl") +include("pool.jl") include("precompile.jl") end # module diff --git a/src/connection.jl b/src/connection.jl index 1c0e6b8..e78344a 100644 --- a/src/connection.jl +++ b/src/connection.jl @@ -56,6 +56,10 @@ like `host`/`user`/`password`. """ function DBInterface.connect(::Type{Connection}, host::AbstractString, user::AbstractString, passwd::Union{AbstractString, Nothing}=nothing; db::Union{AbstractString, Nothing}=nothing, port::Union{Integer, Nothing}=nothing, kw...) opts = ConnectOptions(strip_scheme(host), user, passwd; db=db, port=port, kw...) + return connection(opts) +end + +function connection(opts::ConnectOptions) h = connect(opts) results = ResultOptions(; zero_dates=opts.zero_dates, time_type=opts.time_type) return Connection( diff --git a/src/pool.jl b/src/pool.jl new file mode 100644 index 0000000..132301b --- /dev/null +++ b/src/pool.jl @@ -0,0 +1,239 @@ +""" + MySQL.ConnectionPool(host, user, password=nothing; limit=10, connection_options...) + +A pool of at most `limit` connections, created lazily. Connection options are resolved +once, as for `DBInterface.connect`. Borrow a connection with [`MySQL.with_connection`](@ref) +and close the pool with `close(pool)`. + +Before reuse, each session is reset, its configured database and utf8mb4 settings are +restored, and `init_command` is run again. Use `init_command` only for repeatable session +setup, not one-time writes. Cleanup adds network round trips and uses the configured +read/write timeouts. With no configured database, cleanup also checks `DATABASE()` and +replaces a session that selected a database, since reset cannot clear that selection. +""" +mutable struct ConnectionPool + options::ConnectOptions + condition::Threads.Condition + idle::Vector{Connection} + total::Int + limit::Int + open::Bool +end + +function ConnectionPool(host::AbstractString, user::AbstractString, password::Union{Nothing, AbstractString}=nothing; limit::Int=10, kw...) + limit > 0 || throw(ArgumentError("connection pool limit must be positive")) + opts = ConnectOptions(strip_scheme(host), user, password; kw...) + return ConnectionPool(opts, Threads.Condition(), Connection[], 0, limit, true) +end + +Base.isopen(pool::ConnectionPool) = lock(() -> pool.open, pool.condition) + +function Base.show(io::IO, pool::ConnectionPool) + state = lock(pool.condition) do + (pool.open, pool.total, length(pool.idle)) + end + print(io, "MySQL.ConnectionPool(", state[1] ? "open" : "closed", + ", limit=", pool.limit, ", connections=", state[2], ", idle=", state[3], ")") + return nothing +end + +@noinline closed_pool() = throw(InvalidStateException("MySQL connection pool is closed", :closed)) + +function release_reservation!(pool::ConnectionPool) + lock(pool.condition) do + pool.total -= 1 + notify(pool.condition; all=false) + end + return nothing +end + +# Keep the reservation until close finishes: retiring sessions still count against limit. +function retire_connection!(pool::ConnectionPool, conn::Connection) + try + DBInterface.close!(conn) + finally + release_reservation!(pool) + end + return nothing +end + +function acquire_connection(pool::ConnectionPool) + while true + conn = nothing + owned = false + try + lock(pool.condition) + try + while pool.open && isempty(pool.idle) && pool.total >= pool.limit + wait(pool.condition) + end + pool.open || closed_pool() + if isempty(pool.idle) + pool.total += 1 + else + conn = pop!(pool.idle) + end + owned = true + finally + unlock(pool.condition) + end + conn === nothing && (conn = connection(pool.options)) + if !isopen(conn) + owned = false + retire_connection!(pool, conn) + continue + end + isopen(pool) || closed_pool() + owned = false + return conn + catch + if owned + if conn === nothing + release_reservation!(pool) + else + try + retire_connection!(pool, conn) + catch + end + end + end + rethrow() + end + end +end + +function return_connection!(pool::ConnectionPool, conn::Connection) + pooled = false + try + lock(pool.condition) + try + if pool.open + push!(pool.idle, conn) + pooled = true + notify(pool.condition; all=false) + end + finally + unlock(pool.condition) + end + finally + pooled || retire_connection!(pool, conn) + end + return nothing +end + +# COM_RESET_CONNECTION releases server statements; client generations also have to move +# before later finalizers can park old IDs or a new borrower can prepare another statement. +function reset_pooled_connection!(conn::Connection) + return lock(conn.lock) do + h = conn.handle + h === nothing && return false + s = h.session + isopen(s) || return false + drain_pending!(conn) + response = try + P.reset_connection!(s) + P.read_command_response!(s; kind=P.CMD_SIMPLE) + catch err + # Servers without COM_RESET_CONNECTION can only be reused by reconnecting. + err isa P.Error && err.errno == 1047 && s.phase == P.READY && return false + rethrow() + end + if !(response isa P.OKPacket) || s.phase != P.READY + throw(P.fault!(s, P.ProtocolError("connection reset did not return one final OK"))) + end + invalidate_cursors!(conn) + discard_parked_statements!(conn) + conn.transaction_owner = nothing + conn.buffered_bytes = 0 + opts = conn.options + if !isempty(opts.db) + P.init_db!(s, opts.db) + selected = P.read_command_response!(s; kind=P.CMD_SIMPLE) + if !(selected isa P.OKPacket) || s.phase != P.READY + throw(P.fault!(s, P.ProtocolError("database restoration did not return one final OK"))) + end + conn.db = opts.db + end + h.bootstrapped = bootstrap_charset!(s, response) + opts.init_command === nothing || run_init_command!(s, opts) + if isempty(opts.db) + row = first(DBInterface.execute(conn, "SELECT DATABASE()")) + ismissing(Tables.getcolumn(row, Union{Missing, String}, 1, Symbol("DATABASE()"))) || return false + end + return true + end +end + +""" + MySQL.with_connection(f, pool) + MySQL.with_connection(pool) do conn + # Related operations use this one session. + end + +Call `f(conn)` with an exclusively borrowed connection, then clean up its session before +returning `f`'s value. An uncaught callback error retires the connection and is rethrown +unchanged. Reset failures also retire the connection and propagate; callbacks are never +replayed. A local open check does not guarantee that the server is still reachable. +Consume results inside the callback to observe SQL errors; unread responses are discarded +with the driver's ordinary draining behavior. + +Finish all work using `conn` before returning. Connections, statements and unfinished +streamed results must not be used outside the callback. Owned buffered results and +materialized values remain readable. The Connection is not a revocable lease: escaped +statements can still auto-reprepare under the ordinary driver rules. Nested leases take +separate capacity and can block; pass the borrowed connection into nested work instead. +""" +function with_connection(f, pool::ConnectionPool) + conn = nothing + completed = false + reusable = false + try + conn = acquire_connection(pool) + result = f(conn) + reusable = isopen(pool) && reset_pooled_connection!(conn) + completed = true + return result + finally + if conn === nothing + # Acquisition either took no reservation or retired it before throwing. + elseif completed + reusable ? return_connection!(pool, conn) : retire_connection!(pool, conn) + else + try + retire_connection!(pool, conn) + catch + # Preserve the callback or reset exception and its original backtrace. + end + end + end +end + +""" + close(pool::MySQL.ConnectionPool) + +Reject new leases, wake waiting callers, and close idle connections. Already acquired +leases are not interrupted, even if their callbacks start after a racing close; their +connections close on return. Connections still being created or reset also close instead +of entering the closed pool. Repeated close is harmless. +All idle closes are attempted, with failures reported in a `CompositeException`. +""" +function Base.close(pool::ConnectionPool) + idle = lock(pool.condition) do + pool.open || return Connection[] + pool.open = false + detached = pool.idle + pool.idle = Connection[] + notify(pool.condition; all=true) + detached + end + errors = Any[] + for conn in idle + try + retire_connection!(pool, conn) + catch err + push!(errors, err) + end + end + isempty(errors) || throw(CompositeException(errors)) + return nothing +end diff --git a/test/mysql_trim_workload.jl b/test/mysql_trim_workload.jl index 766a085..43ee5db 100644 --- a/test/mysql_trim_workload.jl +++ b/test/mysql_trim_workload.jl @@ -1,6 +1,7 @@ # JuliaC --trim=safe workload: drives the main MySQL.jl entrypoints — connect (handshake, # auth, charset bootstrap), text-protocol execute (buffered and streaming), prepared -# statements (binary protocol), one-shot parameterized execute, ping, escape, and close — +# statements (binary protocol), one-shot parameterized execute, scoped connection pools, +# ping, escape, and close — # against an in-process scripted server on a loopback Reseau TCP listener, so the # executable needs no database and runs on every platform. # @@ -174,13 +175,16 @@ function serve_connection!(conn)::Nothing cmd = payload[1] if cmd == P.COM_QUIT return nothing - elseif cmd == P.COM_PING + elseif cmd == P.COM_PING || cmd == P.COM_RESET_CONNECTION || cmd == P.COM_INIT_DB send_packet(conn, 1, ok_payload()) elseif cmd == P.COM_STMT_CLOSE # no response elseif cmd == P.COM_QUERY sql = String(payload[2:end]) - if occursin("FROM people", sql) + if sql == "SELECT DATABASE()" + send_resultset(conn, Vector{UInt8}[coldef("DATABASE()")], + Vector{UInt8}[text_row(Union{Missing, String}[missing])]) + elseif occursin("FROM people", sql) send_resultset(conn, people_cols(), people_rows()) else # SET NAMES, CREATE TABLE, START TRANSACTION, COMMIT, INSERT, ... @@ -330,7 +334,47 @@ function exercise_reaper(port::Int)::Nothing return nothing end -function run_trim_workload()::Nothing +function exercise_pool(port::Int, repetitions::Int)::Nothing + for database in ("trim", "") + pool = MySQL.ConnectionPool("127.0.0.1", "root", "secret"; + port, db=database, limit=1, ssl_mode=:disabled, init_command="SET @configured=41") + try + for _ in 1:repetitions + owned = MySQL.with_connection(pool) do conn + check(MySQL.ping(conn), "pooled ping") + cursor = DBInterface.execute(conn, "SELECT id, name, score, joined FROM people")::MySQL.TextCursor{true} + check(Tables.getcolumn(first(cursor), Int32, 1, :id) == 1, "pooled typed row") + cursor + end + check(length(owned) == 3, "pooled owned result") + check(pool.total == 1 && length(pool.idle) == 1, "pooled reuse") + end + marker = ErrorException("pool callback marker") + caught = false + try + MySQL.with_connection(pool) do conn + check(MySQL.ping(conn), "retiring pooled ping") + repetitions > 0 && throw(marker) + end + catch err + caught = err === marker + end + check(caught && pool.total == 0, "pool callback retirement") + finally + close(pool) + end + caught = false + try + MySQL.with_connection(MySQL.ping, pool) + catch err + caught = err isa InvalidStateException + end + check(caught && !isopen(pool), "closed pool rejects lease") + end + return nothing +end + +function run_trim_workload(repetitions::Int)::Nothing listener = TCP.listen(TCP.loopback_addr(0)) laddr = TCP.addr(listener)::TCP.SocketAddrV4 port = Int(laddr.port) @@ -340,6 +384,7 @@ function run_trim_workload()::Nothing schedule(server_task) try run_workload(port) + exercise_pool(port, repetitions) exercise_reaper(port) finally close(listener) @@ -351,8 +396,9 @@ function run_trim_workload()::Nothing end function (@main)(args::Vector{String})::Cint - _ = args - run_trim_workload() + repetitions = isempty(args) ? 2 : parse(Int, args[1]) + check(1 <= repetitions <= 4, "pool repetitions must be between 1 and 4") + run_trim_workload(repetitions) Core.println("mysql trim workload passed") return 0 end diff --git a/test/pool_workflow.jl b/test/pool_workflow.jl new file mode 100644 index 0000000..4dd8994 --- /dev/null +++ b/test/pool_workflow.jl @@ -0,0 +1,144 @@ +# Public pool behavior on both real server families, using only disposable test databases. +function run_pool_workflow(admin) + @testset "Scoped connection pool workflow" begin + opts = admin.options + query(conn, sql) = DBInterface.execute(conn, sql) + scalar(conn, sql) = first(query(conn, sql))[1] + database, other = "pool_review", "pool_review_other" + query(admin, "CREATE DATABASE IF NOT EXISTS $database") + query(admin, "CREATE DATABASE IF NOT EXISTS $other") + query(admin, "DROP TABLE IF EXISTS $database.items") + query(admin, "CREATE TABLE $database.items (id BIGINT PRIMARY KEY AUTO_INCREMENT, label VARCHAR(40))") + pool = MySQL.ConnectionPool(opts.host, opts.user, opts.password; + port=opts.port, db=database, limit=1, init_command="SET @pool_initialized=41", + connect_timeout=10, read_timeout=10, write_timeout=10) + identity = Ref{UInt32}() + initial_isolation = Ref{String}() + isolation = MySQL.server_kind(admin) == :mariadb ? "@@tx_isolation" : "@@transaction_isolation" + stale = Ref{MySQL.Statement}() + unfinished = Ref{MySQL.TextCursor{false}}() + retained = nothing + try + retained = MySQL.with_connection(pool) do conn + identity[] = MySQL.connection_id(conn) + initial_isolation[] = String(scalar(conn, "SELECT $isolation")) + @test scalar(conn, "SELECT DATABASE()") == database + @test scalar(conn, "SELECT @pool_initialized") == 41 + insert = query(conn, "INSERT INTO items (label) VALUES ('kept')") + @test scalar(conn, "SELECT LAST_INSERT_ID()") == DBInterface.lastrowid(insert) + owned = query(conn, "SELECT 'owned across leases' AS value") + query(conn, "CREATE TEMPORARY TABLE pool_temporary (id INT)") + @test scalar(conn, "SELECT GET_LOCK('mysql_pool_review', 0)") == 1 + query(conn, "PREPARE pool_sql_statement FROM 'SELECT 43'") + stale[] = DBInterface.prepare(conn, "SELECT 44") + query(conn, "SET @pool_initialized=99, @pool_leaked=1") + query(conn, "SET SESSION TRANSACTION ISOLATION LEVEL READ UNCOMMITTED") + query(conn, "SET autocommit=0") + query(conn, "INSERT INTO items (label) VALUES ('rolled back')") + query(conn, "USE $other") + query(conn, "SET NAMES latin1") + unfinished[] = DBInterface.execute(conn, "SELECT 1 UNION ALL SELECT 2"; mysql_store_result=false) + @test first(first(unfinished[])) == 1 + owned + end + @test_throws MySQL.Protocol.ProtocolError iterate(unfinished[]) + @test scalar(admin, "SELECT IS_FREE_LOCK('mysql_pool_review')") == 1 + MySQL.with_connection(pool) do conn + @test MySQL.connection_id(conn) == identity[] + @test scalar(conn, "SELECT DATABASE()") == database + @test scalar(conn, "SELECT @pool_initialized") == 41 + @test ismissing(scalar(conn, "SELECT @pool_leaked")) + @test scalar(conn, "SELECT @@autocommit") == 1 + @test scalar(conn, "SELECT $isolation") == initial_isolation[] + @test Tuple(first(query(conn, "SELECT @@character_set_client, @@character_set_connection, @@character_set_results"))) == + ("utf8mb4", "utf8mb4", "utf8mb4") + @test scalar(conn, "SELECT COUNT(*) FROM items WHERE label='rolled back'") == 0 + @test scalar(conn, "SELECT COUNT(*) FROM items WHERE label='kept'") == 1 + @test_throws MySQL.Error query(conn, "SELECT * FROM pool_temporary") + @test_throws MySQL.Error query(conn, "EXECUTE pool_sql_statement") + # Closing an old-generation Julia statement must not close a new server ID. + current = DBInterface.prepare(conn, "SELECT 45") + try + @test stale[].generation < current.generation + DBInterface.close!(stale[]) + @test first(first(DBInterface.execute(current))) == 45 + finally + DBInterface.close!(current) + end + @test scalar(conn, "SELECT 'héllo 🐘'") == "héllo 🐘" + end + @test first(first(retained)) == "owned across leases" + marker = ErrorException("pool callback failure") + result = try + MySQL.with_connection(pool) do conn + query(conn, "START TRANSACTION") + query(conn, "INSERT INTO items (label) VALUES ('callback error')") + throw(marker) + end + catch err + err + end + @test result === marker + @test pool.total == 0 + MySQL.with_connection(pool) do conn + @test MySQL.connection_id(conn) != identity[] + @test scalar(conn, "SELECT COUNT(*) FROM items WHERE label='callback error'") == 0 + @test scalar(conn, "SELECT @pool_initialized") == 41 + # The transaction helper works within the scoped session. + DBInterface.transaction(conn) do + query(conn, "INSERT INTO items (label) VALUES ('transaction')") + end + end + @test scalar(admin, "SELECT COUNT(*) FROM $database.items WHERE label='transaction'") == 1 + finally + isassigned(stale) && DBInterface.close!(stale[]) + close(pool) + end + @test !isopen(pool) && pool.total == 0 + + # With no configured database, unchanged sessions are reusable; USE requires retirement. + no_database = MySQL.ConnectionPool(opts.host, opts.user, opts.password; + port=opts.port, limit=1, connect_timeout=10, read_timeout=10, write_timeout=10) + try + a = MySQL.with_connection(no_database) do conn + @test ismissing(scalar(conn, "SELECT DATABASE()")) + MySQL.connection_id(conn) + end + b = MySQL.with_connection(MySQL.connection_id, no_database) + @test a == b + MySQL.with_connection(no_database) do conn + query(conn, "USE $other") + end + @test no_database.total == 0 + MySQL.with_connection(no_database) do conn + @test MySQL.connection_id(conn) != a + @test ismissing(scalar(conn, "SELECT DATABASE()")) + end + finally + close(no_database) + end + + # Concurrent callers get distinct sessions and keep INSERT/LAST_INSERT_ID together. + concurrent = MySQL.ConnectionPool(opts.host, opts.user, opts.password; + port=opts.port, db=database, limit=2, connect_timeout=10, read_timeout=10, write_timeout=10) + entered, proceed = Channel{UInt32}(2), Channel{Nothing}(2) + jobs = [Threads.@spawn MySQL.with_connection(concurrent) do conn + put!(entered, MySQL.connection_id(conn)) + take!(proceed) + inserted = query(conn, "INSERT INTO items (label) VALUES ('concurrent')") + last = scalar(conn, "SELECT LAST_INSERT_ID()") + @test last == DBInterface.lastrowid(inserted) + last + end for _ in 1:2] + try + @test take!(entered) != take!(entered) + put!(proceed, nothing) + put!(proceed, nothing) + @test fetch(jobs[1]) != fetch(jobs[2]) + finally + close(concurrent) + end + query(admin, "DROP DATABASE $database") + query(admin, "DROP DATABASE $other") + end +end diff --git a/test/protocol/live_tests.jl b/test/protocol/live_tests.jl index c2646d3..b9a1c37 100644 --- a/test/protocol/live_tests.jl +++ b/test/protocol/live_tests.jl @@ -5,6 +5,7 @@ include(joinpath(@__DIR__, "..", "behavior_manifest.jl")) using .BehaviorManifest include(joinpath(@__DIR__, "leak_soak.jl")) include(joinpath(@__DIR__, "..", "user_workflow.jl")) +include(joinpath(@__DIR__, "..", "pool_workflow.jl")) const LIVE_IMAGES = split(get(ENV, "MYSQL_NATIVE_IMAGES", "mysql:8.4,mariadb:11.4"), ',') const ROOT_PW = "native-secret" @@ -143,6 +144,7 @@ function run_live_lane(ref::String; soak::Bool=false, manifest::Bool=false) conn = DBInterface.connect(MySQL.Connection, "127.0.0.1", "root", ROOT_PW; port=port) try run_user_workflow(conn) + run_pool_workflow(conn) finally DBInterface.close!(conn) end diff --git a/test/protocol/pool_tests.jl b/test/protocol/pool_tests.jl new file mode 100644 index 0000000..5cf00be --- /dev/null +++ b/test/protocol/pool_tests.jl @@ -0,0 +1,472 @@ +# Multiple ordinary wire-protocol connections; no pool-specific factory or reset hooks. +function with_pool_peer(f, handler) + listener = Reseau.TCP.listen(Reseau.TCP.loopback_addr(0)) + port = Int(Reseau.TCP.addr(listener).port) + children = Task[] + accepted = AcceptedCounter(0) + server = Threads.@spawn begin + while true + c = try + Reseau.TCP.accept(listener) + catch + break + end + number = increment!(accepted) + push!(children, Threads.@spawn try + handler(c, number) + finally + close(c) + end) + end + end + try + return f(port, accepted) + finally + close(listener) + wait(server) + foreach(fetch, children) + end +end + +function pool_peer_session(c; fault=:none, selected=false, reset_entered=nothing, reset_continue=nothing, + statement_closes=UInt32[]) + reset = false + next_statement = UInt32(77) + statements = Set{UInt32}() + plain_peer_connect!(c; caps=MYSQL8_SERVER_CAPS & ~P.CLIENT_SSL, after=c -> begin + while true + _, cmd, data = try + read_command(c) + catch err + err isa EOFError && break + rethrow() + end + cmd == P.COM_QUIT && break + if cmd == P.COM_RESET_CONNECTION + reset = true + empty!(statements) + next_statement = UInt32(77) + if reset_entered !== nothing + put!(reset_entered, nothing) + take!(reset_continue) + end + fault === :disconnect && break + if fault === :reset || fault === :unsupported + send_err(c, 1, fault === :unsupported ? 1047 : 1205, "reset failed") + elseif fault === :malformed + send_packet(c, 1, UInt8[0]) + elseif fault === :more + send_ok(c, 1; more=true) + else + send_ok(c, 1) + end + elseif cmd == P.COM_INIT_DB + reset && fault === :database ? send_err(c, 1, 1049, "database failed") : send_ok(c, 1) + elseif cmd == P.COM_PING + send_ok(c, 1) + elseif cmd == P.COM_STMT_PREPARE + push!(statements, next_statement) + send_prepare_ok(c, 1, next_statement, P.ColumnDef[], P.ColumnDef[]) + next_statement += UInt32(1) + elseif cmd == P.COM_STMT_CLOSE || cmd == P.COM_STMT_EXECUTE + id = UInt32(data[1]) | (UInt32(data[2]) << 8) | (UInt32(data[3]) << 16) | (UInt32(data[4]) << 24) + if cmd == P.COM_STMT_CLOSE + push!(statement_closes, id) + delete!(statements, id) + else + id in statements ? send_ok(c, 1; affected=1) : send_err(c, 1, 1243, "unknown statement") + end + elseif cmd == P.COM_QUERY + sql = String(data) + if reset && ((fault === :charset && sql == "SET NAMES utf8mb4") || + (fault === :init && sql == "SET @configured=41")) + send_err(c, 1, 1205, "setup failed") + elseif sql == "SELECT DATABASE()" + row = selected ? vcat(UInt8[5], codeunits("other")) : UInt8[0xfb] + send_resultset(c, 1, [coldef("db")], [row]) + elseif sql == "SELECT 7" + send_resultset(c, 1, [coldef("n"; type=P.MYSQL_TYPE_LONG)], [UInt8[1, 0x37]]) + elseif sql == "SELECT streamed" + send_resultset(c, 1, [coldef("n"; type=P.MYSQL_TYPE_LONG)], + [UInt8[1, 0x37], UInt8[1, 0x38], UInt8[1, 0x39]]) + elseif sql == "SELECT broken" || sql == "SELECT unread error" + send_packet(c, 1, column_count(1)) + send_packet(c, 2, coldef("n"; type=P.MYSQL_TYPE_LONG)) + # A wrong sequence is a transport/protocol drain failure. A well-formed + # SQL error is discarded by the driver's existing drain semantics. + send_err(c, sql == "SELECT broken" ? 6 : 3, 1205, "late row error") + elseif sql == "multiple" + send_ok(c, 1; status=P.SERVER_STATUS_AUTOCOMMIT | P.SERVER_MORE_RESULTS_EXISTS) + send_ok(c, 2) + else + send_ok(c, 1) + end + else + error("unexpected pool command $cmd") + end + end + end) +end + +test_pool(port; kw...) = MySQL.ConnectionPool("127.0.0.1", "root", "pw"; + port, ssl_mode=:disabled, connect_timeout=5, read_timeout=5, write_timeout=5, kw...) + +pool_throw_callback(_, marker) = throw(marker) + +# Observe the actual blocked condition wait, rather than assuming a task has run after a sleep. +function wait_for_pool_waiters(pool, n) + @test timedwait(() -> lock(() -> length(pool.condition.waitq) >= n, pool.condition), + 5; pollint=0.001) == :ok +end + +@testset "Scoped connection pools" begin + @test_throws ArgumentError MySQL.ConnectionPool("h", "u"; limit=0) + @test_throws ArgumentError MySQL.ConnectionPool("h", "u"; limit=-1) + @test_throws ArgumentError MySQL.ConnectionPool("h", "u"; unknown_option=true) + lazy = MySQL.ConnectionPool("unused.invalid", "u", "not-for-display"; limit=3) + @test lazy.total == 0 && isempty(lazy.idle) && isopen(lazy) + @test !occursin("not-for-display", sprint(show, lazy)) + close(lazy) + close(lazy) + @test !isopen(lazy) + @test_throws InvalidStateException MySQL.with_connection(identity, lazy) + + @testset "Reuse, streamed draining, multiple results and owned output" begin + with_pool_peer((c, _) -> pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=1, multi_statements=true) + try + first_conn = Ref{MySQL.Connection}() + streamed = Ref{MySQL.TextCursor{false}}() + generation = Ref(0) + output = MySQL.with_connection(pool) do conn + first_conn[] = conn + generation[] = @atomic conn.generation + @test !islocked(pool.condition) + DBInterface.execute(conn, "SELECT 7") + end + @test first(first(output)) == 7 + @test (@atomic first_conn[].generation) > generation[] + @test (@atomic first_conn[].active_token) == 0 + @test MySQL.with_connection(pool) do conn + @test conn === first_conn[] + streamed[] = DBInterface.execute(conn, "SELECT streamed"; mysql_store_result=false) + first(first(streamed[])) + end == 7 + @test_throws P.ProtocolError iterate(streamed[]) + @test MySQL.with_connection(pool) do conn + DBInterface.execute(conn, "multiple") + :complete + end === :complete + @test (@atomic accepted.value) == 1 + @test pool.total == 1 && length(pool.idle) == 1 + finally + close(pool) + end + @test pool.total == 0 + end + end + + @testset "Reset and restoration failures retire, without replay" begin + for fault in (:reset, :malformed, :more, :disconnect, :database, :charset, :init, :unsupported) + with_pool_peer((c, _) -> pool_peer_session(c; fault)) do port, accepted + pool = test_pool(port; db="app", limit=1, init_command="SET @configured=41") + seen = Ref{MySQL.Connection}() + calls = Ref(0) + try + result = try + MySQL.with_connection(pool) do conn + calls[] += 1 + seen[] = conn + 31 + end + catch err + err + end + @test fault === :unsupported ? result == 31 : result isa Exception + @test calls[] == 1 && (@atomic accepted.value) == 1 + @test !isopen(seen[]) && pool.total == 0 && isempty(pool.idle) + finally + close(pool) + end + end + end + end + + @testset "Reset invalidates parked and late statement closes" begin + closes = UInt32[] + with_pool_peer((c, _) -> pool_peer_session(c; statement_closes=closes)) do port, _ + pool = test_pool(port; db="app", limit=1) + old = Ref{MySQL.Statement}() + try + MySQL.with_connection(pool) do conn + old[] = DBInterface.prepare(conn, "UPDATE old") + parked = DBInterface.prepare(conn, "UPDATE parked") + DBInterface.close!(parked) + @test conn.stmts_to_close === parked.reap + end + MySQL.with_connection(pool) do conn + @test conn.stmts_to_close === nothing + current = DBInterface.prepare(conn, "UPDATE current") + @test old[].statement_id == current.statement_id == 77 + @test old[].generation < current.generation + DBInterface.close!(old[]) + @test DBInterface.execute(current).rows_affected == 1 + @test isempty(closes) + DBInterface.close!(current) + @test MySQL.ping(conn) + @test closes == UInt32[77] + end + finally + isassigned(old) && DBInterface.close!(old[]) + close(pool) + end + end + end + + @testset "Pending protocol faults retire; abandoned SQL errors drain" begin + with_pool_peer((c, _) -> pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=1) + calls = Ref(0) + try + result = try + MySQL.with_connection(pool) do conn + calls[] += 1 + DBInterface.execute(conn, "SELECT broken"; mysql_store_result=false) + :finished + end + catch err + err + end + @test result isa P.ProtocolError + @test calls[] == 1 && pool.total == 0 + @test MySQL.with_connection(MySQL.ping, pool) + @test (@atomic accepted.value) == 2 + @test MySQL.with_connection(pool) do conn + DBInterface.execute(conn, "SELECT unread error"; mysql_store_result=false) + :discarded + end === :discarded + @test pool.total == 1 && (@atomic accepted.value) == 2 + finally + close(pool) + end + end + end + + @testset "Database-less sessions retire only when a database is selected" begin + for selected in (false, true) + with_pool_peer((c, _) -> pool_peer_session(c; selected)) do port, accepted + pool = test_pool(port; limit=1) + try + for _ in 1:2 + @test MySQL.with_connection(_ -> 11, pool) == 11 + end + @test (@atomic accepted.value) == (selected ? 2 : 1) + @test pool.total == (selected ? 0 : 1) + finally + close(pool) + end + end + end + end + + @testset "Callback errors, explicit connection close and failed creation" begin + for marker in (ErrorException("callback"), InterruptException()) + with_pool_peer((c, _) -> pool_peer_session(c; fault=:reset)) do port, _ + pool = test_pool(port; db="app", limit=1) + seen = Ref{MySQL.Connection}() + backtrace = Ref{Any}() + result = try + MySQL.with_connection(pool) do conn + seen[] = conn + pool_throw_callback(conn, marker) + end + catch err + backtrace[] = catch_backtrace() + err + end + @test result === marker + @test any(frame -> frame.func === :pool_throw_callback, stacktrace(backtrace[])) + @test !isopen(seen[]) && pool.total == 0 + close(pool) + end + end + with_pool_peer((c, n) -> n == 1 ? send_err(c, 0, 1045, "login failed") : pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=1) + try + @test_throws P.Error MySQL.with_connection(identity, pool) + @test pool.total == 0 + @test MySQL.with_connection(pool) do conn + close(conn) + 12 + end == 12 + @test pool.total == 0 + @test MySQL.with_connection(_ -> 13, pool) == 13 + @test (@atomic accepted.value) == 3 + finally + close(pool) + end + end + with_pool_peer((c, _) -> pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=1) + conn = MySQL.with_connection(identity, pool) + # Defensively reject an idle transport already known to be closed. + close(conn) + replacement = MySQL.with_connection(identity, pool) + @test replacement !== conn && pool.total == 1 + @test (@atomic accepted.value) == 2 + @test MySQL.with_connection(pool) do borrowed + close(pool) + @test pool.total == 1 + MySQL.ping(borrowed) + end + @test pool.total == 0 && !isopen(replacement) + end + end + + @testset "Exclusive capacity and shutdown" begin + with_pool_peer((c, _) -> pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=2) + entered, leave = Channel{MySQL.Connection}(4), Channel{Nothing}(4) + tasks = [Threads.@spawn MySQL.with_connection(pool) do conn + put!(entered, conn) + take!(leave) + MySQL.ping(conn) + end for _ in 1:4] + try + a, b = take!(entered), take!(entered) + @test a !== b + wait_for_pool_waiters(pool, 2) + @test pool.total == 2 && (@atomic accepted.value) == 2 + put!(leave, nothing) + c = take!(entered) + @test c === a || c === b + put!(leave, nothing) + d = take!(entered) + @test d === a || d === b + put!(leave, nothing) + put!(leave, nothing) + @test all(fetch, tasks) + @test length(pool.idle) == 2 && pool.total == 2 + finally + close(pool) + end + @test pool.total == 0 + end + with_pool_peer((c, _) -> pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=1) + entered, leave = Channel{Nothing}(1), Channel{Nothing}(1) + holder = Threads.@spawn MySQL.with_connection(pool) do conn + put!(entered, nothing) + take!(leave) + MySQL.ping(conn) + end + take!(entered) + waiter = Threads.@spawn try + MySQL.with_connection(_ -> error("closed lease entered"), pool) + catch err + err + end + wait_for_pool_waiters(pool, 1) + close(pool) + @test fetch(waiter) isa InvalidStateException + @test pool.total == 1 && !isopen(pool) + put!(leave, nothing) + @test fetch(holder) + @test pool.total == 0 && (@atomic accepted.value) == 1 + end + end + + @testset "Close during creation or reset" begin + for phase in (:creation, :reset) + entered, proceed = Channel{Nothing}(1), Channel{Nothing}(1) + handler = (c, _) -> begin + if phase === :creation + put!(entered, nothing) + take!(proceed) + pool_peer_session(c) + else + pool_peer_session(c; reset_entered=entered, reset_continue=proceed) + end + end + with_pool_peer(handler) do port, _ + pool = test_pool(port; db="app", limit=1) + called = Ref(false) + work = Threads.@spawn try + MySQL.with_connection(pool) do _ + called[] = true + 29 + end + catch err + err + end + take!(entered) + @test !islocked(pool.condition) + close(pool) + @test pool.total == 1 + put!(proceed, nothing) + result = fetch(work) + @test phase === :creation ? result isa InvalidStateException : result == 29 + @test called[] == (phase === :reset) + @test pool.total == 0 && isempty(pool.idle) + end + end + end + + @testset "Cooperative interrupted waiter leaves capacity intact" begin + with_pool_peer((c, _) -> pool_peer_session(c)) do port, accepted + pool = test_pool(port; db="app", limit=1) + entered, leave = Channel{Nothing}(1), Channel{Nothing}(1) + holder = Threads.@spawn MySQL.with_connection(pool) do conn + put!(entered, nothing) + take!(leave) + MySQL.ping(conn) + end + take!(entered) + called = Ref(false) + waiter = Threads.@spawn try + MySQL.with_connection(pool) do _ + called[] = true + end + catch err + err + end + try + wait_for_pool_waiters(pool, 1) + marker = InterruptException() + lock(pool.condition) do + # Condition's error delivery wakes its registered waiter safely. + notify(pool.condition, marker; error=true) + end + @test fetch(waiter) === marker + @test !called[] && pool.total == 1 && isempty(pool.idle) + put!(leave, nothing) + @test fetch(holder) + @test MySQL.with_connection(MySQL.ping, pool) + @test pool.total == 1 && (@atomic accepted.value) == 1 + finally + close(pool) + end + end + end + + @testset "Retirement retains capacity and releases the bookkeeping lock" begin + with_pool_peer((c, _) -> pool_peer_session(c)) do port, _ + pool = test_pool(port; db="app", limit=1) + conn = MySQL.with_connection(identity, pool) + locked, release = Channel{Nothing}(1), Channel{Nothing}(1) + locker = Threads.@spawn lock(conn.lock) do + put!(locked, nothing) + take!(release) + end + take!(locked) + closing = Threads.@spawn close(pool) + @test timedwait(() -> !isopen(pool), 5; pollint=0.001) == :ok + @test pool.total == 1 && !istaskdone(closing) + @test !islocked(pool.condition) + put!(release, nothing) + fetch(locker) + fetch(closing) + @test pool.total == 0 && !isopen(conn) + end + end +end diff --git a/test/protocol/runtests.jl b/test/protocol/runtests.jl index 02c5733..867f5d9 100644 --- a/test/protocol/runtests.jl +++ b/test/protocol/runtests.jl @@ -27,6 +27,7 @@ empty!(P.COVERAGE) include("native_tests.jl") include("cursor_tests.jl") include("binary_tests.jl") + include("pool_tests.jl") include("sharedvalues_tests.jl") include("fuzz_tests.jl") include("perf_tests.jl") From 9fc59d24b5696b3e7f8533f302e6bcc3af00e087 Mon Sep 17 00:00:00 2001 From: Jacob Quinn Date: Sat, 26 Sep 2026 09:25:57 -0600 Subject: [PATCH 2/2] Observe pool waiters through condition notifications --- test/protocol/pool_tests.jl | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/test/protocol/pool_tests.jl b/test/protocol/pool_tests.jl index 5cf00be..6ae2d29 100644 --- a/test/protocol/pool_tests.jl +++ b/test/protocol/pool_tests.jl @@ -113,9 +113,10 @@ test_pool(port; kw...) = MySQL.ConnectionPool("127.0.0.1", "root", "pw"; pool_throw_callback(_, marker) = throw(marker) -# Observe the actual blocked condition wait, rather than assuming a task has run after a sleep. -function wait_for_pool_waiters(pool, n) - @test timedwait(() -> lock(() -> length(pool.condition.waitq) >= n, pool.condition), +# Notify reports how many callers were actually waiting. Normal notifications make +# them recheck capacity; error delivery is observed in the same operation. +function wait_for_pool_waiters(pool, n; value=nothing, error=false) + @test timedwait(() -> lock(() -> notify(pool.condition, value; all=true, error) >= n, pool.condition), 5; pollint=0.001) == :ok end @@ -431,12 +432,8 @@ end err end try - wait_for_pool_waiters(pool, 1) marker = InterruptException() - lock(pool.condition) do - # Condition's error delivery wakes its registered waiter safely. - notify(pool.condition, marker; error=true) - end + wait_for_pool_waiters(pool, 1; value=marker, error=true) @test fetch(waiter) === marker @test !called[] && pool.total == 1 && isempty(pool.idle) put!(leave, nothing)