diff --git a/.env.sample b/.env.sample index 25ddb198..c808010e 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 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 DEBUG_8KNOT=True REDIS_PASSWORD=1234 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 955b231d..fb20b431 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,13 @@ # 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, + enable_client_connection_checks, + env_augur_statement_timeout_ms, +) def cache_query_results( @@ -57,10 +64,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 open. + with closing( + pg.connect( + db_connection_string, + options=augur_cx_options(env_augur_statement_timeout_ms), + ) ) as 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 augur_cur.itersize = server_pagination @@ -193,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 253139be..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: @@ -49,3 +50,55 @@ env_augur_host, env_augur_port, ) + +# 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 +# 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", + ] + ) + + +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 f0e11c44..bcb64a58 100644 --- a/8Knot/db_manager/augur_manager.py +++ b/8Knot/db_manager/augur_manager.py @@ -11,6 +11,11 @@ import requests from sqlalchemy.exc import SQLAlchemyError from models import SearchItem +from cache_manager.cx_common import ( + augur_cx_options, + enable_client_connection_checks, + env_augur_engine_statement_timeout_ms, +) class AugurManager: @@ -112,9 +117,10 @@ 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, ) + salc.event.listen(engine, "connect", enable_client_connection_checks) # verify that engine works try: