Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/pgbouncer.py: 24%
375 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""In-container PgBouncer shared by every proxy worker.
3Each uvicorn worker owns a Prisma query engine with its own pool of
4``connection_limit`` server connections, so the connections a pod holds open
5against Postgres scale as ``workers * connection_limit`` and a database with a
6fixed connection ceiling runs out of room as pods and workers are added.
8When ``LITELLM_PGBOUNCER_ENABLED`` is set, the supervisor process starts one
9PgBouncer next to the workers (no extra network hop: it listens on loopback
10inside the pod) in transaction pooling mode, points ``DATABASE_URL`` at it
11with ``pgbouncer=true`` so Prisma stops using server-side prepared statements,
12and keeps it running for the life of the proxy. Every worker's pool then
13becomes cheap client connections to PgBouncer while the upstream connection
14count is capped at ``LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS`` per pod, no matter
15how many workers run.
17Migrations and the schema diff run in the supervisor before the pooler is
18started, so they always go straight to Postgres. ``DATABASE_URL_READ_REPLICA``
19is left untouched.
21The workers never hold the upstream credential: they log in to PgBouncer as
22``litellm_pgbouncer`` with a random password made at startup, and PgBouncer
23takes the database user's password from its auth file. Under
24``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH`` that password is a
25short-lived token, so the supervisor mints a new one before it expires,
26rewrites the auth file and asks PgBouncer to reload; only new upstream
27connections authenticate, so live ones are unaffected. The pooled
28``DATABASE_URL`` then carries a static password, and the workers must not run
29their own token refresh against it: ``LITELLM_PGBOUNCER_POOLED_DATABASE_URL``
30tells them so, while a read replica keeps refreshing its own token.
31"""
33from __future__ import annotations
35import atexit
36import functools
37import os
38import re
39import secrets
40import shlex
41import shutil
42import signal
43import socket
44import subprocess
45import tempfile
46import threading
47import time
48import urllib.parse
49from collections.abc import Callable, Mapping, Sequence
50from dataclasses import dataclass
51from datetime import datetime, timezone
52from pathlib import Path
53from types import MappingProxyType
54from typing import Final
56from pydantic import Field
57from pydantic_settings import BaseSettings, SettingsConfigDict
59from litellm._logging import verbose_proxy_logger
60from litellm.proxy.db.token_auth import (
61 DatabaseTokenAuth,
62 IAMEndpoint,
63 mint_database_token,
64 parse_database_token_expiration,
65 parse_iam_endpoint_from_url,
66)
68PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_"
69PGBOUNCER_POOLED_ENV_VAR: Final = "LITELLM_PGBOUNCER_POOLED_DATABASE_URL"
70PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1"
71PGBOUNCER_POOL_USER: Final = "litellm_pgbouncer"
72PGBOUNCER_INI_NAME: Final = "pgbouncer.ini"
73PGBOUNCER_USERLIST_NAME: Final = "userlist.txt"
74PGBOUNCER_CA_NAME: Final = "server-ca.pem"
75PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0
76PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0
77PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0
78PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody"
79PGBOUNCER_MIN_VERSION: Final = (1, 19)
80PGBOUNCER_MAX_PASSWORD_BYTES: Final = 2048
81PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)")
82PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS: Final = 180.0
83PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS: Final = 600.0
84PGBOUNCER_TOKEN_RETRY_SECONDS: Final = 30.0
86# Prisma's client-side TLS params describe the hop to Postgres, which becomes
87# PgBouncer's server side. They move into ``server_tls_*`` and must not stay on
88# the loopback URL: the listener speaks plain TCP and Prisma would refuse it
89# under ``sslmode=require`` or ``channel_binding=require``.
90PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset(
91 {"sslmode", "sslcert", "sslaccept", "sslidentity", "sslpassword", "channel_binding", "gssencmode"}
92)
93POOLED_URL_DROPPED_KEYS: Final[frozenset[str]] = PRISMA_TLS_PARAM_KEYS | frozenset(("options", "pgbouncer"))
94PGBOUNCER_SSLMODES: Final[frozenset[str]] = frozenset(
95 {"disable", "allow", "prefer", "require", "verify-ca", "verify-full"}
96)
99class PgBouncerSettings(BaseSettings):
100 """``LITELLM_PGBOUNCER_*`` env vars, read once in the supervisor."""
102 model_config = SettingsConfigDict(
103 env_prefix=PGBOUNCER_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True
104 )
106 enabled: bool = False
107 port: int = Field(default=6432, ge=1, le=65535)
108 max_db_connections: int = Field(default=20, ge=1)
109 max_client_conn: int = Field(default=1000, ge=1)
110 binary: str = "pgbouncer"
113@dataclass(frozen=True, slots=True)
114class PgBouncerPlan:
115 ini: str
116 pooled_url: str
117 upstream_user: str
118 upstream_password: str | None
119 pool_password: str
120 ca_source: str | None = None
122 def userlist(self, upstream_password: str) -> str:
123 return "".join(
124 f"{_userlist_quote(user)} {_userlist_quote(password)}\n"
125 for user, password in ((self.upstream_user, upstream_password), (PGBOUNCER_POOL_USER, self.pool_password))
126 )
129@dataclass(frozen=True, slots=True)
130class PgBouncerError:
131 reason: str
134def _single_quoted(value: str) -> str:
135 """Quote for SQL and for PgBouncer's ``[databases]`` connection string: both double a literal ``'``."""
136 return "'" + value.replace("'", "''") + "'"
139def _userlist_quote(value: str) -> str:
140 return '"' + value.replace('"', '""') + '"'
143def _option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None:
144 """The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else.
146 Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``; a
147 detached ``-c`` is folded into the token that follows it first.
148 """
149 folded: Final = tuple(
150 f"-c{tokens[index + 1]}" if token == "-c" and index + 1 < len(tokens) else token
151 for index, token in enumerate(tokens)
152 if index == 0 or tokens[index - 1] != "-c"
153 )
154 settings: Final = tuple(token[2:] for token in folded if token.startswith(("-c", "--")) and "=" in token[2:])
155 return settings if len(settings) == len(folded) else None
158def _connect_query(options: str) -> str | PgBouncerError:
159 """Turn Prisma's ``options=-c name=value ...`` startup param into ``SET`` statements.
161 PgBouncer rejects any ``-c`` setting in ``options`` that is not one of the
162 handful it tracks (``statement_timeout`` and ``lock_timeout`` are not), so
163 the settings are applied to each new server connection instead. Every
164 client shares them, which is what the single ``DATABASE_URL`` gave anyway.
165 """
166 settings: Final = _option_settings(tuple(shlex.split(options)))
167 if settings is None:
168 return PgBouncerError(f"cannot translate the DATABASE_URL options {options!r} into PgBouncer settings")
169 return "; ".join(
170 f"SET {name.strip()} TO {_single_quoted(value.strip())}"
171 for name, value in (setting.split("=", 1) for setting in settings)
172 )
175def _server_tls_settings(sslmode: str, sslcert: str, sslaccept: str, ca_path: Path) -> tuple[str, ...] | PgBouncerError:
176 """``server_tls_*`` lines naming ``ca_path``, the runtime-dir copy of the bundle: the original (or the
177 0600 root pinned by ``pin_bundle_root``) is often unreadable for the user PgBouncer drops to."""
178 if sslmode not in PGBOUNCER_SSLMODES:
179 return PgBouncerError(f"unsupported sslmode {sslmode!r} on DATABASE_URL")
180 verify: Final = sslmode in ("verify-ca", "verify-full") or (sslmode == "require" and sslaccept == "strict")
181 if verify and not sslcert:
182 return PgBouncerError(
183 "DATABASE_URL asks for a verified TLS connection but names no CA bundle; "
184 "add sslcert=<ca.pem> (or sslrootcert=) so the in-container PgBouncer can verify Postgres"
185 )
186 mode: Final = "verify-full" if verify else sslmode
187 return (f"server_tls_sslmode = {mode}", *((f"server_tls_ca_file = {ca_path}",) if sslcert else ()))
190def plan_pgbouncer(
191 upstream_url: str,
192 settings: PgBouncerSettings,
193 runtime_dir: Path,
194 run_as_user: str | None,
195) -> PgBouncerPlan | PgBouncerError:
196 """Render the PgBouncer config for ``upstream_url`` and the loopback URL Prisma uses instead.
198 Params describing Prisma's own pool (``connection_limit``, ``pool_timeout``,
199 ...) stay on the pooled URL; the TLS params and ``options`` describe the hop
200 to Postgres and move into the PgBouncer config. The upstream password is
201 left out of the config on purpose: PgBouncer then takes it from the auth
202 file, which can be rewritten while it runs. ``run_as_user`` is the
203 unprivileged user PgBouncer drops to when the proxy runs as root, which
204 PgBouncer itself refuses to do.
205 """
206 parsed: Final = urllib.parse.urlsplit(upstream_url)
207 params: Final[Mapping[str, str]] = MappingProxyType(
208 dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
209 )
210 dbname: Final = urllib.parse.unquote(parsed.path.lstrip("/"))
211 username: Final = urllib.parse.unquote(parsed.username or "")
212 password: Final = None if parsed.password is None else urllib.parse.unquote(parsed.password)
213 if not parsed.hostname or not username or not dbname:
214 return PgBouncerError("DATABASE_URL must carry a host, user and database name for the in-container PgBouncer")
215 if username == PGBOUNCER_POOL_USER:
216 return PgBouncerError(
217 f"the database user cannot be named {PGBOUNCER_POOL_USER!r}: that is the user the workers log in to the "
218 "in-container PgBouncer as, and PgBouncer keeps one password per user"
219 )
220 if "sslidentity" in params:
221 return PgBouncerError("client certificates (sslidentity) are not supported with the in-container PgBouncer")
222 tls: Final = _server_tls_settings(
223 params.get("sslmode", "prefer"),
224 params.get("sslcert", ""),
225 params.get("sslaccept", ""),
226 runtime_dir / PGBOUNCER_CA_NAME,
227 )
228 if isinstance(tls, PgBouncerError):
229 return tls
230 connect_query: Final = _connect_query(params["options"]) if params.get("options") else ""
231 if isinstance(connect_query, PgBouncerError):
232 return connect_query
233 upstream: Final = " ".join(
234 (
235 f"host={_single_quoted(parsed.hostname)}",
236 f"port={parsed.port or 5432}",
237 f"dbname={_single_quoted(dbname)}",
238 f"user={_single_quoted(username)}",
239 *((f"connect_query={_single_quoted(connect_query)}",) if connect_query else ()),
240 )
241 )
242 ini: Final = "\n".join(
243 (
244 "[databases]",
245 f"{dbname} = {upstream}",
246 "",
247 "[pgbouncer]",
248 f"listen_addr = {PGBOUNCER_LISTEN_ADDR}",
249 f"listen_port = {settings.port}",
250 f"unix_socket_dir = {runtime_dir}",
251 f"auth_file = {runtime_dir / PGBOUNCER_USERLIST_NAME}",
252 "auth_type = scram-sha-256",
253 f"stats_users = {PGBOUNCER_POOL_USER}",
254 "pool_mode = transaction",
255 f"max_client_conn = {settings.max_client_conn}",
256 f"default_pool_size = {settings.max_db_connections}",
257 f"max_db_connections = {settings.max_db_connections}",
258 "ignore_startup_parameters = extra_float_digits",
259 *tls,
260 *((f"user = {run_as_user}",) if run_as_user else ()),
261 "",
262 )
263 )
264 pooled_query: Final = urllib.parse.urlencode(
265 (*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true"))
266 )
267 pool_password: Final = secrets.token_urlsafe(32)
268 pooled_url: Final = urllib.parse.urlunsplit(
269 parsed._replace(
270 netloc=f"{PGBOUNCER_POOL_USER}:{pool_password}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query
271 )
272 )
273 return PgBouncerPlan(
274 ini=ini,
275 pooled_url=pooled_url,
276 upstream_user=username,
277 upstream_password=password,
278 pool_password=pool_password,
279 ca_source=params.get("sslcert") or None,
280 )
283def pooled_database_url(upstream_url: str, settings: PgBouncerSettings) -> str | PgBouncerError:
284 """The loopback URL of a PgBouncer another container in the pod already runs for ``upstream_url``.
286 Only the container that started PgBouncer knows the pool user's password, so
287 this logs in as the upstream user, whom the auth file lists as well.
288 """
289 plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir=Path("/nonexistent"), run_as_user=None)
290 if isinstance(plan, PgBouncerError):
291 return plan
292 password: Final = urllib.parse.urlsplit(upstream_url).password or ""
293 credentials: Final = f"{urllib.parse.quote(plan.upstream_user, safe='')}:{password}"
294 return urllib.parse.urlunsplit(
295 urllib.parse.urlsplit(plan.pooled_url)._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}")
296 )
299def _write_private(path: Path, content: str, run_as_user: str | None) -> None:
300 with open(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w", encoding="utf-8") as handle:
301 handle.write(content)
302 if run_as_user is not None:
303 shutil.chown(path, user=run_as_user)
306def write_pgbouncer_ini(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError:
307 """Write the ini (mode 0600) and the CA copy, and return the ini path.
309 ``run_as_user`` is the user PgBouncer drops to when started as root; it has
310 to own the files it re-reads on reload and the socket directory.
311 """
312 ini_path: Final = runtime_dir / PGBOUNCER_INI_NAME
313 ca_path: Final = runtime_dir / PGBOUNCER_CA_NAME
314 if plan.ca_source is not None:
315 try:
316 shutil.copyfile(plan.ca_source, ca_path)
317 except OSError as error:
318 return PgBouncerError(f"cannot read the CA bundle {plan.ca_source!r} named by sslcert: {error}")
319 _write_private(ini_path, plan.ini, run_as_user)
320 if run_as_user is not None:
321 runtime_dir.chmod(0o700)
322 for path in (runtime_dir, *((ca_path,) if plan.ca_source is not None else ())):
323 shutil.chown(path, user=run_as_user)
324 return ini_path
327def write_userlist(userlist: str, runtime_dir: Path, run_as_user: str | None) -> Path:
328 """Replace the auth file in one step, so a PgBouncer starting or reloading meanwhile reads the old or the new one whole."""
329 userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME
330 staged_path: Final = runtime_dir / f".{PGBOUNCER_USERLIST_NAME}.next"
331 _write_private(staged_path, userlist, run_as_user)
332 os.replace(staged_path, userlist_path)
333 return userlist_path
336def export_pooled_database_url(pooled_url: str) -> None:
337 os.environ["DATABASE_URL"] = pooled_url
338 os.environ[PGBOUNCER_POOLED_ENV_VAR] = "true"
341def database_url_is_pooled(environ: Mapping[str, str] = os.environ) -> bool:
342 return environ.get(PGBOUNCER_POOLED_ENV_VAR) == "true"
345@dataclass(frozen=True, slots=True)
346class PgBouncerTokenSource:
347 auth: DatabaseTokenAuth
348 endpoint: IAMEndpoint
350 def mint(self) -> str:
351 """The token as Postgres expects it: ``mint_database_token`` returns it percent-encoded for a URL."""
352 return urllib.parse.unquote(mint_database_token(self.auth, self.endpoint))
354 def expires_at(self, token: str) -> datetime | None:
355 return parse_database_token_expiration(self.auth, token)
358def _utcnow() -> datetime:
359 return datetime.now(timezone.utc).replace(tzinfo=None)
362class PgBouncerTokenRefresher:
363 """Keeps the token in PgBouncer's auth file current from a daemon thread.
365 ``install`` gets each fresh token and is expected to rewrite the auth file
366 and reload PgBouncer. The next refresh is due ``buffer_seconds`` before the
367 token expires, or ``fallback_seconds`` later when the expiry cannot be read.
368 A refresh that fails leaves the previous auth file in place and is retried
369 after ``retry_seconds``: the old token stays good until it expires, so a
370 transient credential-provider error costs nothing unless it persists.
371 """
373 def __init__(
374 self,
375 source: PgBouncerTokenSource,
376 install: Callable[[str], None],
377 *,
378 buffer_seconds: float = PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS,
379 fallback_seconds: float = PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS,
380 retry_seconds: float = PGBOUNCER_TOKEN_RETRY_SECONDS,
381 now: Callable[[], datetime] = _utcnow,
382 ) -> None:
383 self._source: Final = source
384 self._install: Final = install
385 self._buffer_seconds: Final = buffer_seconds
386 self._fallback_seconds: Final = fallback_seconds
387 self._retry_seconds: Final = retry_seconds
388 self._now: Final = now
389 self._stopping: Final = threading.Event()
390 self._delay: float = 0.0
391 self._thread: threading.Thread | None = None
393 def refresh(self) -> float | PgBouncerError:
394 label: Final = self._source.auth.label
395 try:
396 token: Final = self._source.mint()
397 except Exception as mint_error:
398 return PgBouncerError(f"could not mint a {label} for the in-container pgbouncer: {mint_error!r}")
399 if len(token.encode()) >= PGBOUNCER_MAX_PASSWORD_BYTES:
400 return PgBouncerError(
401 f"the {label} is {len(token.encode())} bytes long, but PgBouncer's auth file holds passwords of at "
402 f"most {PGBOUNCER_MAX_PASSWORD_BYTES - 1} bytes"
403 )
404 try:
405 self._install(token)
406 except OSError as install_error:
407 return PgBouncerError(f"could not install the {label} into the pgbouncer auth file: {install_error}")
408 expires_at: Final = self._source.expires_at(token)
409 if expires_at is None:
410 return self._fallback_seconds
411 return max(self._retry_seconds, (expires_at - self._now()).total_seconds() - self._buffer_seconds)
413 def start(self) -> PgBouncerError | None:
414 primed: Final = self.refresh()
415 if isinstance(primed, PgBouncerError):
416 return primed
417 self._delay = primed
418 self._thread = threading.Thread(target=self._run, daemon=True, name="litellm-pgbouncer-token-refresh")
419 self._thread.start()
420 return None
422 def _run(self) -> None:
423 while not self._stopping.wait(self._delay):
424 self._delay = self._refresh_and_report()
426 def _refresh_and_report(self) -> float:
427 outcome: Final = self.refresh()
428 if isinstance(outcome, PgBouncerError):
429 verbose_proxy_logger.error(
430 "In-container pgbouncer keeps its current %s (%s); retrying in %.0fs.",
431 self._source.auth.label,
432 outcome.reason,
433 self._retry_seconds,
434 )
435 return self._retry_seconds
436 verbose_proxy_logger.info(
437 "In-container pgbouncer picked up a fresh %s; the next one is due in %.0fs.",
438 self._source.auth.label,
439 outcome,
440 )
441 return outcome
443 def stop(self) -> None:
444 self._stopping.set()
445 if self._thread is not None:
446 self._thread.join()
449def _port_open(port: int) -> bool:
450 try:
451 with socket.create_connection((PGBOUNCER_LISTEN_ADDR, port), timeout=0.5):
452 return True
453 except OSError:
454 return False
457def _unix_socket_open(path: Path) -> bool:
458 with socket.socket(socket.AF_UNIX) as probe:
459 probe.settimeout(0.5)
460 try:
461 probe.connect(str(path))
462 except OSError:
463 return False
464 return True
467def unix_socket_path(runtime_dir: Path, port: int) -> Path:
468 return runtime_dir / f".s.PGSQL.{port}"
471def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError:
472 """``(major, minor)`` from ``<binary> --version``.
474 Readiness relies on PgBouncer exiting when it cannot bind its TCP port,
475 which it does from 1.19 on. Older releases log a warning and serve the unix
476 socket alone, so their socket would vouch for a port held by someone else.
477 """
478 try:
479 output: Final = subprocess.run(
480 (binary, "--version"), capture_output=True, text=True, check=False, timeout=10
481 ).stdout
482 except (OSError, subprocess.TimeoutExpired) as run_error:
483 return PgBouncerError(f"could not run {binary!r} --version: {run_error}")
484 found: Final = PGBOUNCER_VERSION_PATTERN.search(output)
485 if found is None:
486 return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}")
487 return int(found[1]), int(found[2])
490def _end(process: subprocess.Popen[bytes]) -> None:
491 if process.poll() is not None:
492 return
493 process.terminate()
494 try:
495 process.wait(timeout=PGBOUNCER_STOP_GRACE_SECONDS)
496 except subprocess.TimeoutExpired:
497 process.kill()
498 process.wait()
501class PgBouncerProcess:
502 """Runs ``argv`` as a foreground child and restarts it whenever it exits on its own.
504 Prisma reconnects by itself after a failed query, so a PgBouncer crash
505 costs the requests in flight plus one failed query per idle pooled
506 connection the crash severed, and nothing else once the replacement is
507 listening again. A replacement that cannot be spawned, finds its port
508 taken, exits again or never starts listening is retried every
509 ``restart_delay_seconds`` until ``stop`` is called.
511 A connect probe of ``port`` cannot tell the child from another process
512 that grabbed the port after the availability check, so readiness also
513 needs ``socket_path``: the unix socket PgBouncer creates in the private
514 runtime directory, which it only does once every TCP listener is bound
515 (PgBouncer 1.19 or newer, see ``pgbouncer_version``).
516 """
518 def __init__(
519 self,
520 argv: Sequence[str],
521 port: int,
522 socket_path: Path,
523 restart_delay_seconds: float = PGBOUNCER_RESTART_DELAY_SECONDS,
524 ready_timeout_seconds: float = PGBOUNCER_READY_TIMEOUT_SECONDS,
525 ) -> None:
526 self.argv: Final = tuple(argv)
527 self.port: Final = port
528 self.socket_path: Final = socket_path
529 self.restart_delay_seconds: Final = restart_delay_seconds
530 self.ready_timeout_seconds: Final = ready_timeout_seconds
531 self._stopping: Final = threading.Event()
532 self._lock: Final = threading.Lock()
533 self._process: subprocess.Popen[bytes] | None = None
535 @property
536 def pid(self) -> int | None:
537 with self._lock:
538 return None if self._process is None else self._process.pid
540 def _spawn(self) -> subprocess.Popen[bytes] | PgBouncerError | None:
541 """Start a child, or None once ``stop`` ran; both take the lock so no child can slip in after a stop.
543 The port has to be free first: a listener that is already there would
544 pass the readiness check while the child fails to bind.
545 """
546 with self._lock:
547 if self._stopping.is_set():
548 return None
549 if _port_open(self.port):
550 return PgBouncerError(f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is already in use by another process")
551 try:
552 process: Final = subprocess.Popen(self.argv)
553 except OSError as spawn_error:
554 return PgBouncerError(f"could not start {self.argv[0]!r}: {spawn_error}")
555 self._process = process
556 return process
558 def _wait_ready(self, process: subprocess.Popen[bytes]) -> PgBouncerError | None:
559 deadline: Final = time.monotonic() + self.ready_timeout_seconds
560 while time.monotonic() < deadline:
561 if process.poll() is not None:
562 return PgBouncerError(f"pgbouncer exited with status {process.returncode} during startup")
563 if _port_open(self.port) and _unix_socket_open(self.socket_path):
564 return None
565 time.sleep(0.1)
566 if _port_open(self.port):
567 return PgBouncerError(
568 f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is served by another process, not the pgbouncer that was started"
569 )
570 return PgBouncerError(
571 f"pgbouncer did not start listening on {PGBOUNCER_LISTEN_ADDR}:{self.port} "
572 f"within {self.ready_timeout_seconds:.0f}s"
573 )
575 def start(self) -> PgBouncerError | None:
576 """Spawn PgBouncer, wait until it listens on port and unix socket, then supervise it from a daemon thread."""
577 process: Final = self._spawn()
578 if process is None:
579 return PgBouncerError("pgbouncer was stopped before it started")
580 if isinstance(process, PgBouncerError):
581 return process
582 not_ready: Final = self._wait_ready(process)
583 if not_ready is not None:
584 self.stop()
585 return not_ready
586 self._watch(process)
587 return None
589 def _watch(self, process: subprocess.Popen[bytes]) -> None:
590 threading.Thread(
591 target=self._supervise, args=(process,), daemon=True, name="litellm-pgbouncer-supervisor"
592 ).start()
594 def _supervise(self, process: subprocess.Popen[bytes]) -> None:
595 status: Final = process.wait()
596 if self._stopping.is_set():
597 return
598 verbose_proxy_logger.error(
599 "In-container pgbouncer (pid %s) exited with status %s; restarting in %.1fs.",
600 process.pid,
601 status,
602 self.restart_delay_seconds,
603 )
604 self._restart_after_delay()
606 def _restart_after_delay(self) -> None:
607 time.sleep(self.restart_delay_seconds)
608 process: Final = self._spawn()
609 if process is None:
610 return
611 if isinstance(process, PgBouncerError):
612 self._retry_restart(process.reason)
613 return
614 not_ready: Final = self._wait_ready(process)
615 if not_ready is None:
616 self._watch(process)
617 return
618 _end(process)
619 self._retry_restart(not_ready.reason)
621 def _retry_restart(self, reason: str) -> None:
622 if self._stopping.is_set():
623 return
624 verbose_proxy_logger.error(
625 "In-container pgbouncer could not be restarted (%s); retrying in %.1fs.", reason, self.restart_delay_seconds
626 )
627 threading.Thread(target=self._restart_after_delay, daemon=True, name="litellm-pgbouncer-supervisor").start()
629 def reload(self) -> None:
630 with self._lock:
631 if self._process is not None:
632 self._process.send_signal(signal.SIGHUP)
634 def stop(self) -> None:
635 with self._lock:
636 self._stopping.set()
637 process: Final = self._process
638 if process is not None:
639 _end(process)
642def install_pgbouncer_token(
643 plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None, pooler: PgBouncerProcess, token: str
644) -> None:
645 write_userlist(plan.userlist(token), runtime_dir, run_as_user)
646 pooler.reload()
649def _install_upstream_password(
650 plan: PgBouncerPlan,
651 runtime_dir: Path,
652 run_as_user: str | None,
653 pooler: PgBouncerProcess,
654 token_auth: DatabaseTokenAuth | None,
655 upstream_url: str,
656) -> PgBouncerTokenRefresher | None | PgBouncerError:
657 if token_auth is None:
658 if plan.upstream_password is None:
659 return PgBouncerError(
660 "DATABASE_URL carries no password and neither IAM_TOKEN_DB_AUTH nor AZURE_POSTGRESQL_AUTH is on, "
661 "so the in-container PgBouncer has nothing to authenticate to Postgres with"
662 )
663 write_userlist(plan.userlist(plan.upstream_password), runtime_dir, run_as_user)
664 return None
665 refresher: Final = PgBouncerTokenRefresher(
666 PgBouncerTokenSource(auth=token_auth, endpoint=parse_iam_endpoint_from_url(upstream_url)),
667 functools.partial(install_pgbouncer_token, plan, runtime_dir, run_as_user, pooler),
668 )
669 failed: Final = refresher.start()
670 if failed is not None:
671 return failed
672 return refresher
675def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]:
676 """An exit hook that does nothing in a forked child, which inherits the parent's ``atexit`` table."""
677 owner_pid: Final = os.getpid()
679 def run() -> None:
680 if os.getpid() == owner_pid:
681 action()
683 return run
686def start_in_container_pgbouncer(
687 settings: PgBouncerSettings,
688 upstream_url: str,
689 token_auth: DatabaseTokenAuth | None = None,
690 register_exit_hook: Callable[[Callable[[], None]], object] = atexit.register,
691) -> str | PgBouncerError:
692 """Start the pooler for ``upstream_url`` and return the loopback URL the workers must use.
694 The pooler lives as long as this process: it is stopped from the exit hooks
695 once the worker manager has returned, and only by the process that started
696 it (gunicorn forks its workers, so they carry the hooks too). PgBouncer
697 refuses to run as root, so a root proxy (the default image) has it drop to
698 ``nobody``. With ``token_auth`` the password on ``upstream_url`` is ignored:
699 the pooler mints its own tokens and renews them for as long as it runs.
700 """
701 version: Final = pgbouncer_version(settings.binary)
702 if isinstance(version, PgBouncerError):
703 return version
704 if version < PGBOUNCER_MIN_VERSION:
705 return PgBouncerError(
706 f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy "
707 f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} "
708 "or newer is required"
709 )
710 runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-"))
711 register_exit_hook(_only_in_this_process(lambda: shutil.rmtree(runtime_dir, ignore_errors=True)))
712 run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None
713 plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir, run_as_user)
714 if isinstance(plan, PgBouncerError):
715 return plan
716 ini_path: Final = write_pgbouncer_ini(plan, runtime_dir, run_as_user)
717 if isinstance(ini_path, PgBouncerError):
718 return ini_path
719 pooler: Final = PgBouncerProcess(
720 argv=(settings.binary, str(ini_path)),
721 port=settings.port,
722 socket_path=unix_socket_path(runtime_dir, settings.port),
723 )
724 refresher: Final = _install_upstream_password(plan, runtime_dir, run_as_user, pooler, token_auth, upstream_url)
725 if isinstance(refresher, PgBouncerError):
726 return refresher
727 failed: Final = pooler.start()
728 if failed is not None:
729 if refresher is not None:
730 refresher.stop()
731 return failed
732 register_exit_hook(_only_in_this_process(pooler.stop))
733 if refresher is not None:
734 register_exit_hook(_only_in_this_process(refresher.stop))
735 verbose_proxy_logger.info(
736 "In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections%s.",
737 pooler.pid,
738 PGBOUNCER_LISTEN_ADDR,
739 settings.port,
740 settings.max_db_connections,
741 "" if token_auth is None else f" and renewing its {token_auth.label} before each one expires",
742 )
743 return plan.pooled_url