From 40e4705789b09a4ceda13a2788842f7615d52b1f Mon Sep 17 00:00:00 2001 From: Gabriele Battimelli Date: Wed, 8 Jul 2026 00:25:53 +0200 Subject: [PATCH] Defer DB connection until after jixia analysis completes TCP keepalives (previous fix) did not prevent the connection drop -- same failure recurred at the same spot. Keepalives are invisible to whatever is actually closing the connection (likely an application-level idle policy on Heroku's side, not a network path issue), so they can't help. Split jixia_db.load_data into run_analysis() (pure CPU work, no DB connection) and load_data() (DB writes only). database/__init__.py now runs the 30-90+ minute analysis phase first, then opens a fresh connection right before writing -- so no connection is ever left idle for more than the time between individual queries. --- database/__init__.py | 54 ++++++++++++++++++++++++++------------------ database/jixia_db.py | 40 +++++++++++++++++++------------- 2 files changed, 56 insertions(+), 38 deletions(-) diff --git a/database/__init__.py b/database/__init__.py index 27f89e8..1e76bcf 100644 --- a/database/__init__.py +++ b/database/__init__.py @@ -6,11 +6,22 @@ from jixia.structs import parse_name from .informalize import generate_informal -from .jixia_db import load_data +from .jixia_db import run_analysis, load_data from .vector_db import create_vector_db from .create_schema import create_schema +def _connect(): + return psycopg.connect( + os.environ["CONNECTION_STRING"], + autocommit=True, + keepalives=1, + keepalives_idle=30, + keepalives_interval=10, + keepalives_count=5, + ) + + def main(): parser = ArgumentParser() subparser = parser.add_subparsers() @@ -42,32 +53,31 @@ def main(): args = parser.parse_args() - # The jixia analysis phase can run 30-60+ min with no DB traffic at all. - # GitHub Actions runners sit behind Azure's outbound NAT, whose default - # idle timeout (~4 min) silently drops the connection before that phase - # ends. TCP keepalives (well under 4 min) keep the NAT mapping alive. - with psycopg.connect( - os.environ["CONNECTION_STRING"], - autocommit=True, - keepalives=1, - keepalives_idle=30, - keepalives_interval=10, - keepalives_count=5, - ) as conn: - if args.command == "schema": + if args.command == "schema": + with _connect() as conn: create_schema(conn) - elif args.command == "jixia": - # jixia runs each module with cwd=project_root, so the module file - # path must be absolute — a relative root would be resolved twice. - project = LeanProject(os.path.abspath(args.project_root)) - prefixes = [parse_name(p) for p in args.prefixes.split(",")] - load_data(project, prefixes, conn) - elif args.command == "informal": + elif args.command == "jixia": + # jixia runs each module with cwd=project_root, so the module file + # path must be absolute — a relative root would be resolved twice. + project = LeanProject(os.path.abspath(args.project_root)) + prefixes = [parse_name(p) for p in args.prefixes.split(",")] + # Run analysis (30-90+ min of pure CPU work) before opening the DB + # connection — GitHub Actions runners sit behind Azure's outbound NAT + # (~4 min idle timeout), and Heroku's own connection handling has also + # been observed to drop connections left idle that long. Connecting + # only once we're ready to write avoids relying on either surviving + # a long idle period. + analysis = run_analysis(project, prefixes) + with _connect() as conn: + load_data(project, analysis, conn) + elif args.command == "informal": + with _connect() as conn: generate_informal( conn, batch_size=args.batch_size, limit_level=args.limit_level, limit_num_per_level=args.limit_num_per_level, ) - elif args.command == "vector-db": + elif args.command == "vector-db": + with _connect() as conn: create_vector_db(conn, os.environ["CHROMA_PATH"], batch_size=args.batch_size) diff --git a/database/jixia_db.py b/database/jixia_db.py index 1601fd5..dca48c3 100644 --- a/database/jixia_db.py +++ b/database/jixia_db.py @@ -37,7 +37,29 @@ def _get_range(declaration: Declaration): return None -def load_data(project: LeanProject, prefixes: list[LeanName], conn: Connection): +def run_analysis(project: LeanProject, prefixes: list[LeanName]) -> list[tuple[Path, list[LeanName]]]: + # Pure CPU analysis, no DB connection involved; this can take 30-90+ min. + lean_sysroot = Path(os.environ["LEAN_SYSROOT"]) + lean_src = lean_sysroot / "src" / "lean" + # Each jixia worker loads the full Mathlib environment (~2-3 GB), so the + # default thread count (CPUs + 4) can exhaust memory and get the process + # OOM-killed. Cap it via JIXIA_MAX_WORKERS in memory-constrained CI. + max_workers_env = os.environ.get("JIXIA_MAX_WORKERS") + max_workers = int(max_workers_env) if max_workers_env else None + + module_lists = [] + for d in project.root, lean_src: + results = project.batch_run_jixia( + base_dir=d, + prefixes=prefixes, + plugins=["module", "declaration", "symbol"], + max_workers=max_workers, + ) + module_lists.append((d, [r[0] for r in results])) + return module_lists + + +def load_data(project: LeanProject, analysis: list[tuple[Path, list[LeanName]]], conn: Connection): def load_module(data: Iterable[LeanName], base_dir: Path): values = [] for m in data: @@ -174,22 +196,8 @@ def topological_sort(): """) with conn.cursor() as cursor: - lean_sysroot = Path(os.environ["LEAN_SYSROOT"]) - lean_src = lean_sysroot / "src" / "lean" - # Each jixia worker loads the full Mathlib environment (~2-3 GB), so the - # default thread count (CPUs + 4) can exhaust memory and get the process - # OOM-killed. Cap it via JIXIA_MAX_WORKERS in memory-constrained CI. - max_workers_env = os.environ.get("JIXIA_MAX_WORKERS") - max_workers = int(max_workers_env) if max_workers_env else None all_modules = [] - for d in project.root, lean_src: - results = project.batch_run_jixia( - base_dir=d, - prefixes=prefixes, - plugins=["module", "declaration", "symbol"], - max_workers=max_workers, - ) - modules = [r[0] for r in results] + for d, modules in analysis: load_module(modules, d) all_modules += modules