fix(proxy): refuse pgbouncer older than 1.19, whose unix socket cannot vouch for the tcp port

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/db/pgbouncer.py

@@ -28,6 +28,7 @@ from __future__ import annotations
 
 import atexit
 import os
+import re
 import shlex
 import shutil
 import socket
@@ -56,6 +57,8 @@ PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0
 PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0
 PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0
 PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody"
+PGBOUNCER_MIN_VERSION: Final = (1, 19)
+PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)")
 PGBOUNCER_TOKEN_AUTH_CONFLICT: Final = (
     f"the in-container pgbouncer cannot be combined with {IAM_TOKEN_DB_AUTH_ENV_VAR} or "
     f"{AZURE_POSTGRESQL_AUTH_ENV_VAR}: each worker rotates the database password on its own schedule and the pooler "
@@ -272,6 +275,25 @@ def unix_socket_path(runtime_dir: Path, port: int) -> Path:
     return runtime_dir / f".s.PGSQL.{port}"
 
 
+def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError:
+    """``(major, minor)`` from ``<binary> --version``.
+
+    Readiness relies on PgBouncer exiting when it cannot bind its TCP port,
+    which it does from 1.19 on. Older releases log a warning and serve the unix
+    socket alone, so their socket would vouch for a port held by someone else.
+    """
+    try:
+        output: Final = subprocess.run(
+            (binary, "--version"), capture_output=True, text=True, check=False, timeout=10
+        ).stdout
+    except (OSError, subprocess.TimeoutExpired) as run_error:
+        return PgBouncerError(f"could not run {binary!r} --version: {run_error}")
+    found: Final = PGBOUNCER_VERSION_PATTERN.search(output)
+    if found is None:
+        return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}")
+    return int(found[1]), int(found[2])
+
+
 def _end(process: subprocess.Popen[bytes]) -> None:
     if process.poll() is not None:
         return
@@ -296,7 +318,8 @@ class PgBouncerProcess:
     A connect probe of ``port`` cannot tell the child from another process
     that grabbed the port after the availability check, so readiness also
     needs ``socket_path``: the unix socket PgBouncer creates in the private
-    runtime directory, which it only does once every TCP listener is bound.
+    runtime directory, which it only does once every TCP listener is bound
+    (PgBouncer 1.19 or newer, see ``pgbouncer_version``).
     """
 
     def __init__(
@@ -429,6 +452,15 @@ def start_in_container_pgbouncer(
     """
     if token_auth_enabled:
         return PgBouncerError(PGBOUNCER_TOKEN_AUTH_CONFLICT)
+    version: Final = pgbouncer_version(settings.binary)
+    if isinstance(version, PgBouncerError):
+        return version
+    if version < PGBOUNCER_MIN_VERSION:
+        return PgBouncerError(
+            f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy "
+            f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} "
+            "or newer is required"
+        )
     runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-"))
     atexit.register(shutil.rmtree, runtime_dir, ignore_errors=True)
     run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None

tests/test_litellm/proxy/db/test_pgbouncer.py

@@ -21,6 +21,7 @@ from litellm.proxy.db.pgbouncer import (
     PgBouncerPlan,
     PgBouncerProcess,
     PgBouncerSettings,
+    pgbouncer_version,
     plan_pgbouncer,
     start_in_container_pgbouncer,
     unix_socket_path,
@@ -170,12 +171,14 @@ def _fake_pooler(
     exit_immediately: bool = False,
     port_file: Path | None = None,
     bind_delay_seconds: float = 0.0,
+    version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable",
 ) -> Path:
     """An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.<port>`` in the socket dir.
 
     Port and socket dir come from the ini it is given, else from ``port`` and
     ``tmp_path``. With ``port_file`` each start reads the port from that file
     instead. ``bind_delay_seconds`` holds the bind back, like a slow start.
+    ``--version`` prints ``version_banner``.
     """
     script: Final = tmp_path / "fake-pgbouncer"
     script.write_text(
@@ -183,6 +186,9 @@ def _fake_pooler(
             f"""\
             #!{sys.executable}
             import configparser, os, pathlib, select, socket, sys, time
+            if sys.argv[1:] == ["--version"]:
+                print({version_banner!r})
+                sys.exit(0)
             if {exit_immediately!r}:
                 sys.exit(3)
             ini = configparser.ConfigParser()
@@ -481,6 +487,40 @@ class TestStartInContainerPgBouncer:
         assert "AZURE_POSTGRESQL_AUTH" in outcome.reason
         assert not _listening(port)
 
+    def test_a_pgbouncer_that_survives_a_failed_tcp_bind_is_refused_without_starting(self, tmp_path: Path):
+        port: Final = _free_port()
+        binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.18.1\nlibevent 2.1.12-stable")
+        settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
+        outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm")
+        assert isinstance(outcome, PgBouncerError)
+        assert "PgBouncer 1.18" in outcome.reason
+        assert "1.19" in outcome.reason
+        assert not _listening(port)
+
+    def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path):
+        port: Final = _free_port()
+        binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0")
+        settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
+        assert start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") == (
+            f"postgresql://app:pw@127.0.0.1:{port}/litellm?pgbouncer=true"
+        )
+        assert _listening(port)
+
+
+class TestPgBouncerVersion:
+    def test_reads_major_and_minor_from_the_banner(self, tmp_path: Path):
+        assert pgbouncer_version(str(_fake_pooler(tmp_path, _free_port()))) == (1, 25)
+
+    def test_a_binary_that_cannot_run_is_reported(self, tmp_path: Path):
+        outcome: Final = pgbouncer_version(str(tmp_path / "missing-pgbouncer"))
+        assert isinstance(outcome, PgBouncerError)
+        assert "missing-pgbouncer" in outcome.reason
+
+    def test_a_banner_without_a_version_is_reported(self, tmp_path: Path):
+        outcome: Final = pgbouncer_version(str(_fake_pooler(tmp_path, _free_port(), version_banner="something else")))
+        assert isinstance(outcome, PgBouncerError)
+        assert "something else" in outcome.reason
+
 
 class TestPgBouncerSettings:
     def test_reads_the_litellm_pgbouncer_env_vars(self, monkeypatch: pytest.MonkeyPatch):