Cancel DB queries when the 8Knot instance shuts down

.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

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(

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(

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()

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: