pre_ping=True behind transaction-mode PgBouncer doubles server transactions and moves PgBouncer's queue into checkout wait, outside pool_timeout #58

import asyncio
import csv
import io
import time

from prometheus_client import REGISTRY
from pydantic import SecretStr
from sqlalchemy import text
from sqlalchemy_foundation_kit import create_async_session_manager
from sqlalchemy_foundation_kit.contrib.metrics import PostgresMetrics
from sqlalchemy_foundation_kit.contrib.settings import BasePostgresConfig, ConnectionSettings, PoolSettings
from testcontainers.core.container import DockerContainer
from testcontainers.core.network import Network
from testcontainers.core.wait_strategies import LogMessageWaitStrategy
from testcontainers.postgres import PostgresContainer


def start_pgbouncer(network: Network, **settings: object) -> DockerContainer:
    container = DockerContainer("edoburu/pgbouncer:v1.25.2-p0").with_network(network).with_exposed_ports(5432)
    env = {"DB_HOST": "db", "DB_USER": "test", "DB_PASSWORD": "test", "DB_NAME": "test",
           "AUTH_TYPE": "scram-sha-256", "ADMIN_USERS": "test", "PGPASSWORD": "test",
           "POOL_MODE": "transaction", **{k.upper(): str(v) for k, v in settings.items()}}
    for key, value in env.items():
        container.with_env(key, value)
    container.waiting_for(LogMessageWaitStrategy("process up").with_startup_timeout(30))
    return container.start()


def bouncer_transactions(bouncer: DockerContainer) -> int:
    out = bouncer.exec(["psql", "-h", "127.0.0.1", "-U", "test", "-d", "pgbouncer", "--csv", "-c", "SHOW STATS"])
    return sum(int(r["total_xact_count"]) for r in csv.DictReader(io.StringIO(out.output.decode())) if r["database"] == "test")


def config(bouncer: DockerContainer, pre_ping: bool) -> BasePostgresConfig:
    return BasePostgresConfig(
        connection=ConnectionSettings(host=bouncer.get_container_host_ip(), port=int(bouncer.get_exposed_port(5432)),
                                      user="test", password=SecretStr("test"), database="test"),
        application_name="probe",
        pool=PoolSettings(size=2, max_overflow=0, timeout=0.2, pre_ping=pre_ping),
        use_orjson_serialization=False,
    )


def sample(prefix: str, name: str) -> float:
    return REGISTRY.get_sample_value(f"{prefix}_postgres_db_{name}") or 0.0


async def warm(manager) -> None:
    async def one() -> None:
        async with manager.get_transaction() as s:
            await s.execute(text("SELECT pg_sleep(0.05)"))
    await asyncio.gather(one(), one())  # both local connections opened and back in the pool


async def main(bouncer: DockerContainer) -> None:
    print("1. PgBouncer transactions for 10 requests, each `get_transaction()` + one SELECT")
    for pre_ping in (True, False):
        manager = create_async_session_manager(config(bouncer, pre_ping))
        await warm(manager)
        before = bouncer_transactions(bouncer)
        for _ in range(10):
            async with manager.get_transaction() as s:
                await s.execute(text("SELECT 1"))
        print(f"   pre_ping={pre_ping!s:5}: {bouncer_transactions(bouncer) - before} transactions")
        await manager.aclose()

    print("2. PgBouncer's only server connection is busy for 1 s; local pool has an idle connection; pool timeout 0.2 s")
    for pre_ping in (True, False):
        prefix = f"preping_{str(pre_ping).lower()}"
        manager = create_async_session_manager(config(bouncer, pre_ping), metrics=PostgresMetrics(prefix=prefix))
        await warm(manager)
        busy, release = asyncio.Event(), asyncio.Event()

        async def holder() -> None:
            async with manager.get_transaction() as s:
                await s.execute(text("SELECT 1"))
                busy.set()
                await release.wait()

        task = asyncio.create_task(holder())
        await busy.wait()
        asyncio.get_running_loop().call_later(1.0, release.set)
        wait_before = sample(prefix, "connection_checkout_wait_seconds_sum")
        started = time.perf_counter()
        async with manager.get_transaction() as s:
            await s.execute(text("SELECT 1"))
        took = time.perf_counter() - started
        await task
        wait = sample(prefix, "connection_checkout_wait_seconds_sum") - wait_before
        print(f"   pre_ping={pre_ping!s:5}: request took {took:.2f} s; checkout wait recorded {wait:.2f} s; "
              f"connection_timeouts_total={sample(prefix, 'connection_timeouts_total')}")
        await manager.aclose()


with Network() as network, PostgresContainer(
    "postgres:17-alpine", username="test", password="test", dbname="test"
).with_network(network).with_network_aliases("db"):
    bouncer = start_pgbouncer(network, default_pool_size=1, max_client_conn=50)
    try:
        asyncio.run(main(bouncer))
    finally:
        bouncer.stop()