From 93b6ecf4054562340a65c3fad60a0879c67c270f Mon Sep 17 00:00:00 2001 From: Caio Fonseca Date: Fri, 14 Aug 2026 14:36:59 +0100 Subject: [PATCH 1/3] Cancel augur queries when the 8Knot instance shuts down Signed-off-by: Caio Fonseca --- .env.sample | 5 ++++ 8Knot/cache_manager/cache_facade.py | 38 ++++++++++++++++++++++++++--- 8Knot/cache_manager/cx_common.py | 37 ++++++++++++++++++++++++++++ 8Knot/db_manager/augur_manager.py | 3 ++- 4 files changed, 78 insertions(+), 5 deletions(-) diff --git a/.env.sample b/.env.sample index 25ddb198..b8c3395c 100644 --- a/.env.sample +++ b/.env.sample @@ -4,6 +4,11 @@ AUGUR_PASSWORD=augur AUGUR_PORT=5432 AUGUR_SCHEMA=data,augur_data AUGUR_USERNAME=augur +# Bound how long augur keeps working on our behalf, so queries don't outlive a +# shut-down 8Knot instance. Milliseconds. +AUGUR_STATEMENT_TIMEOUT_MS=500000 +AUGUR_ENGINE_STATEMENT_TIMEOUT_MS=1800000 +AUGUR_IDLE_TX_TIMEOUT_MS=120000 DEBUG_8KNOT=True REDIS_PASSWORD=1234 diff --git a/8Knot/cache_manager/cache_facade.py b/8Knot/cache_manager/cache_facade.py index 955b231d..48124f91 100644 --- a/8Knot/cache_manager/cache_facade.py +++ b/8Knot/cache_manager/cache_facade.py @@ -22,6 +22,7 @@ """ import logging +from contextlib import closing from uuid import uuid4 import psycopg2 as pg from psycopg2.extras import execute_values @@ -32,7 +33,30 @@ # other files importing cache_facade need to know how to resolve # .cx_common- interpreter is invoked at a higher level, so relative # import required. -from .cx_common import db_cx_string, env_augur_schema, cache_cx_string +from .cx_common import ( + db_cx_string, + cache_cx_string, + augur_cx_options, + env_augur_statement_timeout_ms, +) + + +def _abort_query_when_client_disconnects(augur_conn) -> None: + """Asks the augur server to stop a running query as soon as we disappear. + + Postgres only polls the client socket during a query if this is set, so without + it a query keeps running after the worker is gone. It's postgres 14+, and a + server that rejects it is still bounded by statement_timeout, so a failure here + is worth logging but not worth failing the query over. + """ + try: + with augur_conn.cursor() as cur: + cur.execute("SET client_connection_check_interval = '10s'") + # SET is transactional- without a commit it's undone by the next rollback. + augur_conn.commit() + except pg.Error as e: + logging.warning(f"AUGUR: client disconnect checks unavailable: {e}") + augur_conn.rollback() def cache_query_results( @@ -57,10 +81,16 @@ def cache_query_results( client_pagination (int, optional): _description_. Defaults to 2000. """ logging.warning(f"{target_table} -- CQR CACHE_QUERY_RESULTS BEGIN") - with pg.connect( - db_connection_string, - options=f"-c search_path={env_augur_schema}", + # 'closing' rather than a bare 'with': psycopg2's context manager ends the + # transaction but leaves the connection- and its server-side cursor- open. + with closing( + pg.connect( + db_connection_string, + options=augur_cx_options(env_augur_statement_timeout_ms), + ) ) as augur_conn: + _abort_query_when_client_disconnects(augur_conn) + with augur_conn.cursor(name=f"{target_table}-{uuid4()}") as augur_cur: # set number of rows we want from primary db at a time augur_cur.itersize = server_pagination diff --git a/8Knot/cache_manager/cx_common.py b/8Knot/cache_manager/cx_common.py index 253139be..a601413c 100644 --- a/8Knot/cache_manager/cx_common.py +++ b/8Knot/cache_manager/cx_common.py @@ -49,3 +49,40 @@ env_augur_host, env_augur_port, ) + +# how long a single statement may run against augur before the server aborts it. +# worker queries have to finish inside celery's soft time limit (540s), otherwise +# celery kills the worker and the query is orphaned on the augur server. +env_augur_statement_timeout_ms = os.getenv("AUGUR_STATEMENT_TIMEOUT_MS", "500000") + +# the app-server's searchbar query legitimately takes much longer than a worker +# query, so it gets its own ceiling. +env_augur_engine_statement_timeout_ms = os.getenv("AUGUR_ENGINE_STATEMENT_TIMEOUT_MS", "1800000") + +# how long a session may sit inside an open transaction before augur terminates it. +# a worker killed between fetches leaves its transaction open, so this is what +# reclaims the server-side cursor when 8Knot is shut down. +env_augur_idle_tx_timeout_ms = os.getenv("AUGUR_IDLE_TX_TIMEOUT_MS", "120000") + + +def augur_cx_options(statement_timeout_ms: str) -> str: + """libpq 'options' for a connection to the augur db. + + Postgres won't stop working on our behalf just because we've disappeared: it + doesn't check the client socket while it's busy in a query, so a query outlives + the 8Knot instance that asked for it. These settings make the augur server + responsible for cleaning up after us. + + tcp_keepalives_* apply to the server's end of the socket, so augur notices a + container that no longer exists in ~90s rather than the OS default of 2 hours. + """ + return " ".join( + [ + f"-c search_path={env_augur_schema}", + f"-c statement_timeout={statement_timeout_ms}", + f"-c idle_in_transaction_session_timeout={env_augur_idle_tx_timeout_ms}", + "-c tcp_keepalives_idle=60", + "-c tcp_keepalives_interval=10", + "-c tcp_keepalives_count=3", + ] + ) diff --git a/8Knot/db_manager/augur_manager.py b/8Knot/db_manager/augur_manager.py index f0e11c44..21b27727 100644 --- a/8Knot/db_manager/augur_manager.py +++ b/8Knot/db_manager/augur_manager.py @@ -11,6 +11,7 @@ import requests from sqlalchemy.exc import SQLAlchemyError from models import SearchItem +from cache_manager.cx_common import augur_cx_options, env_augur_engine_statement_timeout_ms class AugurManager: @@ -112,7 +113,7 @@ def get_engine(self): engine = salc.create_engine( database_connection_string, - connect_args={"options": "-csearch_path={}".format(self.schema)}, + connect_args={"options": augur_cx_options(env_augur_engine_statement_timeout_ms)}, pool_pre_ping=True, ) From aaf2c8542d486160a4902fbaf4cdaabf9e7ebb8c Mon Sep 17 00:00:00 2001 From: Caio Fonseca <141309898+EngCaioFonseca@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:53:39 +0100 Subject: [PATCH 2/3] Fix docstring for client disconnect query handling Signed-off-by: Caio Fonseca --- 8Knot/cache_manager/cache_facade.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/8Knot/cache_manager/cache_facade.py b/8Knot/cache_manager/cache_facade.py index 48124f91..d2dbc0e9 100644 --- a/8Knot/cache_manager/cache_facade.py +++ b/8Knot/cache_manager/cache_facade.py @@ -42,7 +42,7 @@ def _abort_query_when_client_disconnects(augur_conn) -> None: - """Asks the augur server to stop a running query as soon as we disappear. + """Asks the server to stop a running query as soon as we disappear. Postgres only polls the client socket during a query if this is set, so without it a query keeps running after the worker is gone. It's postgres 14+, and a From 812e640c32452917b135a37484f260c1b3ff2864 Mon Sep 17 00:00:00 2001 From: Caio Fonseca Date: Mon, 31 Aug 2026 01:45:03 +0100 Subject: [PATCH 3/3] Harden augur query cancellation and skip retries on timeout Signed-off-by: Caio Fonseca --- .env.sample | 4 ++-- 8Knot/_celery.py | 12 +++++++++++- 8Knot/cache_manager/cache_facade.py | 29 ++++++----------------------- 8Knot/cache_manager/cx_common.py | 22 +++++++++++++++++++--- 8Knot/db_manager/augur_manager.py | 7 ++++++- 5 files changed, 44 insertions(+), 30 deletions(-) diff --git a/.env.sample b/.env.sample index b8c3395c..c808010e 100644 --- a/.env.sample +++ b/.env.sample @@ -4,8 +4,8 @@ AUGUR_PASSWORD=augur AUGUR_PORT=5432 AUGUR_SCHEMA=data,augur_data AUGUR_USERNAME=augur -# Bound how long augur keeps working on our behalf, so queries don't outlive a -# shut-down 8Knot instance. Milliseconds. +# Bound individual source DB statements and abandoned transactions so they do +# not continue indefinitely after worker or app-server shutdown. Milliseconds. AUGUR_STATEMENT_TIMEOUT_MS=500000 AUGUR_ENGINE_STATEMENT_TIMEOUT_MS=1800000 AUGUR_IDLE_TX_TIMEOUT_MS=120000 diff --git a/8Knot/_celery.py b/8Knot/_celery.py index d541d55f..33de3b6c 100644 --- a/8Knot/_celery.py +++ b/8Knot/_celery.py @@ -1,5 +1,6 @@ -from celery import Celery +from celery import Celery, Task from dash import CeleryManager +from psycopg2.errors import QueryCanceled import os redis_host = "{}".format(os.getenv("REDIS_SERVICE_HOST", "redis-broker")) @@ -9,10 +10,19 @@ """CREATE CELERY TASK QUEUE AND MANAGER""" + + +class EightKnotTask(Task): + """Base task that does not retry deterministic database cancellations.""" + + dont_autoretry_for = (QueryCanceled,) + + celery_app = Celery( __name__, broker=REDIS_URL, backend=REDIS_URL, + task_cls=EightKnotTask, ) celery_app.conf.update( diff --git a/8Knot/cache_manager/cache_facade.py b/8Knot/cache_manager/cache_facade.py index d2dbc0e9..fb20b431 100644 --- a/8Knot/cache_manager/cache_facade.py +++ b/8Knot/cache_manager/cache_facade.py @@ -37,28 +37,11 @@ db_cx_string, cache_cx_string, augur_cx_options, + enable_client_connection_checks, env_augur_statement_timeout_ms, ) -def _abort_query_when_client_disconnects(augur_conn) -> None: - """Asks the server to stop a running query as soon as we disappear. - - Postgres only polls the client socket during a query if this is set, so without - it a query keeps running after the worker is gone. It's postgres 14+, and a - server that rejects it is still bounded by statement_timeout, so a failure here - is worth logging but not worth failing the query over. - """ - try: - with augur_conn.cursor() as cur: - cur.execute("SET client_connection_check_interval = '10s'") - # SET is transactional- without a commit it's undone by the next rollback. - augur_conn.commit() - except pg.Error as e: - logging.warning(f"AUGUR: client disconnect checks unavailable: {e}") - augur_conn.rollback() - - def cache_query_results( db_connection_string: str, query: str, @@ -82,14 +65,14 @@ def cache_query_results( """ logging.warning(f"{target_table} -- CQR CACHE_QUERY_RESULTS BEGIN") # 'closing' rather than a bare 'with': psycopg2's context manager ends the - # transaction but leaves the connection- and its server-side cursor- open. + # transaction but leaves the connection open. with closing( pg.connect( db_connection_string, options=augur_cx_options(env_augur_statement_timeout_ms), ) ) as augur_conn: - _abort_query_when_client_disconnects(augur_conn) + enable_client_connection_checks(augur_conn) with augur_conn.cursor(name=f"{target_table}-{uuid4()}") as augur_cur: # set number of rows we want from primary db at a time @@ -223,11 +206,11 @@ def caching_wrapper(func_name: str, query: str, repolist: list[int], n_repolist_ target_table=func_name, bookkeeping_data=tuple({"cache_func": func_name, "repo_id": r} for r in repolist), ) - except Exception as e: - logging.critical(f"{func_name}_POSTGRES ERROR: {e}") + except Exception: + logging.exception(f"{func_name}_POSTGRES ERROR") # raise exception so caching function knows to restart - raise Exception(e) + raise def retrieve_from_cache( diff --git a/8Knot/cache_manager/cx_common.py b/8Knot/cache_manager/cx_common.py index a601413c..1f82414b 100644 --- a/8Knot/cache_manager/cx_common.py +++ b/8Knot/cache_manager/cx_common.py @@ -5,6 +5,7 @@ import os import logging import time +import psycopg2 as pg # credentials to access database from environment try: @@ -50,9 +51,9 @@ env_augur_port, ) -# how long a single statement may run against augur before the server aborts it. -# worker queries have to finish inside celery's soft time limit (540s), otherwise -# celery kills the worker and the query is orphaned on the augur server. +# How long each worker statement may run before the server aborts it. Named-cursor +# fetches are separate statements, so disconnect checks still enforce cleanup if +# the Celery process reaches its own time limit first. env_augur_statement_timeout_ms = os.getenv("AUGUR_STATEMENT_TIMEOUT_MS", "500000") # the app-server's searchbar query legitimately takes much longer than a worker @@ -86,3 +87,18 @@ def augur_cx_options(statement_timeout_ms: str) -> str: "-c tcp_keepalives_count=3", ] ) + + +def enable_client_connection_checks(connection, _connection_record=None) -> None: + """Enable disconnect polling when PostgreSQL and its host support it.""" + if connection.server_version < 140000: + return + + try: + with connection.cursor() as cursor: + cursor.execute("SET client_connection_check_interval = '10s'") + # SET is transactional; finish it before application queries begin. + connection.commit() + except pg.Error as error: + logging.warning(f"AUGUR: client disconnect checks unavailable: {error}") + connection.rollback() diff --git a/8Knot/db_manager/augur_manager.py b/8Knot/db_manager/augur_manager.py index 21b27727..bcb64a58 100644 --- a/8Knot/db_manager/augur_manager.py +++ b/8Knot/db_manager/augur_manager.py @@ -11,7 +11,11 @@ import requests from sqlalchemy.exc import SQLAlchemyError from models import SearchItem -from cache_manager.cx_common import augur_cx_options, env_augur_engine_statement_timeout_ms +from cache_manager.cx_common import ( + augur_cx_options, + enable_client_connection_checks, + env_augur_engine_statement_timeout_ms, +) class AugurManager: @@ -116,6 +120,7 @@ def get_engine(self): connect_args={"options": augur_cx_options(env_augur_engine_statement_timeout_ms)}, pool_pre_ping=True, ) + salc.event.listen(engine, "connect", enable_client_connection_checks) # verify that engine works try: