chore: improve startup and shutdown for background tasks

Author: StefanFl

backend/application/background_tasks/management/__init__.py

@@ -0,0 +1,26 @@
+import logging
+from typing import Any
+
+from django.core.management.base import BaseCommand
+from huey.contrib.djhuey import HUEY as huey
+from huey.contrib.stats import HueyInflight
+
+logger = logging.getLogger("secobserve.background_tasks")
+
+
+class Command(BaseCommand):
+    help = "Delete stale in-flight entries of the Huey statistics, left over from an ungraceful shutdown."
+
+    def handle(self, *args: Any, **options: Any) -> None:
+        # The consumer and the web server share one container and there is only one backend
+        # container per pod, so nothing of this queue can legitimately be in flight during startup.
+        # Entries still present are leftovers of tasks that never emitted a terminal signal, for
+        # example because the container was killed while they were executing. Neither the huey
+        # consumer nor the statistics themselves clean them up, so they would be reported as
+        # running for hours by the background task statistics.
+        if getattr(huey, "_stats", None) is None:
+            logger.info("Huey statistics are not enabled, no in-flight entries to flush")
+            return
+
+        deleted = HueyInflight.delete().where(HueyInflight.queue == huey.name).execute()
+        logger.info("Flushed %s stale in-flight entries of queue %s", deleted, huey.name)

backend/application/background_tasks/management/commands/__init__.py

@@ -0,0 +1,61 @@
+from unittest.mock import MagicMock, patch
+
+import peewee
+from django.core.management import call_command
+from django.test import TestCase
+from huey.contrib.stats import HueyInflight
+
+COMMAND = "flush_huey_inflight"
+MODULE = "application.background_tasks.management.commands.flush_huey_inflight"
+
+
+class TestFlushHueyInflightCommand(TestCase):
+    """Tests for the `flush_huey_inflight` management command.
+
+    The command deletes the entries of the `huey_inflight` table for its own queue. They are
+    written by the Huey statistics when a task starts and are only deleted when the task emits a
+    terminal signal, which does not happen when the container is killed while the task is running.
+    """
+
+    def setUp(self) -> None:
+        # The Huey statistics use peewee, not the Django ORM, so the table cannot be created by a
+        # Django migration and has to be set up here.
+        self.database = peewee.SqliteDatabase(":memory:")
+        self.bind_ctx = self.database.bind_ctx([HueyInflight])
+        self.bind_ctx.__enter__()
+        self.database.create_tables([HueyInflight])
+
+        self.huey = MagicMock()
+        self.huey.name = "secobserve"
+        self.huey._stats = MagicMock()
+
+    def tearDown(self) -> None:
+        self.database.drop_tables([HueyInflight])
+        self.bind_ctx.__exit__(None, None, None)
+        self.database.close()
+
+    def test_stale_entries_of_own_queue_are_deleted(self) -> None:
+        HueyInflight.create(task_id="task_1", queue="secobserve", task="module.Task_1", started=1000.0)
+        HueyInflight.create(task_id="task_2", queue="secobserve", task="module.Task_2", started=2000.0)
+        HueyInflight.create(task_id="task_3", queue="other_queue", task="module.Task_3", started=3000.0)
+
+        with patch(f"{MODULE}.huey", self.huey):
+            call_command(COMMAND)
+
+        remaining = [entry.task_id for entry in HueyInflight.select()]
+        self.assertEqual(["task_3"], remaining)
+
+    def test_no_entries_is_not_an_error(self) -> None:
+        with patch(f"{MODULE}.huey", self.huey):
+            call_command(COMMAND)
+
+        self.assertEqual(0, HueyInflight.select().count())
+
+    def test_nothing_is_deleted_when_statistics_are_disabled(self) -> None:
+        HueyInflight.create(task_id="task_1", queue="secobserve", task="module.Task_1", started=1000.0)
+        self.huey._stats = None
+
+        with patch(f"{MODULE}.huey", self.huey):
+            call_command(COMMAND)
+
+        self.assertEqual(1, HueyInflight.select().count())

backend/application/background_tasks/management/commands/flush_huey_inflight.py

@@ -40,9 +40,26 @@ User.objects.create_superuser(
 EOD
 fi
 
+python manage.py flush_huey_inflight
 python manage.py run_huey --flush-locks &
+HUEY_PID=$!
+RUNSERVER_PID=""
+
+# This script is PID 1 of the container, but a non-interactive shell does not forward signals to
+# its children. Without the trap neither the development server nor the huey consumer would ever
+# see the SIGTERM of a container shutdown and both would be killed after the grace period. The huey
+# consumer needs the SIGTERM to finish its running tasks and to mark them as no longer in flight.
+trap 'kill -TERM $HUEY_PID $RUNSERVER_PID 2>/dev/null' TERM INT
 
 python manage.py register_parsers
 python manage.py initial_license_load
 
-python manage.py runserver_plus 0.0.0.0:8000
+python manage.py runserver_plus 0.0.0.0:8000 &
+RUNSERVER_PID=$!
+
+# "wait" returns a non-zero exit code when it is interrupted by a signal, which would abort the
+# script because of "set -o errexit". Waiting for the huey consumer after the development server
+# has terminated gives it the chance to shut down gracefully before this script and with it the
+# container exits.
+wait "$RUNSERVER_PID" || true
+wait "$HUEY_PID" || true

backend/unittests/background_tasks/management/__init__.py

@@ -46,7 +46,16 @@ fi
 
 python /app/manage.py collectstatic --noinput
 
+python manage.py flush_huey_inflight
 python manage.py run_huey --flush-locks &
+HUEY_PID=$!
+GUNICORN_PID=""
+
+# This script is PID 1 of the container, but a non-interactive shell does not forward signals to
+# its children. Without the trap neither gunicorn nor the huey consumer would ever see the SIGTERM
+# of a container shutdown and both would be killed after the grace period. The huey consumer needs
+# the SIGTERM to finish its running tasks and to mark them as no longer in flight.
+trap 'kill -TERM $HUEY_PID $GUNICORN_PID 2>/dev/null' TERM INT
 
 python manage.py register_parsers
 python manage.py initial_license_load
@@ -66,4 +75,11 @@ if [ -z "${GUNICORN_LIMIT_REQUEST_FIELD_SIZE}" ]; then
     GUNICORN_LIMIT_REQUEST_FIELD_SIZE=16380
 fi
 
-gunicorn config.wsgi --no-control-socket --bind 0.0.0.0:5000 --timeout 1200 --chdir=/app -w $GUNICORN_WORKERS -k gthread --threads $GUNICORN_THREADS --limit-request-field_size $GUNICORN_LIMIT_REQUEST_FIELD_SIZE
+gunicorn config.wsgi --no-control-socket --bind 0.0.0.0:5000 --timeout 1200 --chdir=/app -w $GUNICORN_WORKERS -k gthread --threads $GUNICORN_THREADS --limit-request-field_size $GUNICORN_LIMIT_REQUEST_FIELD_SIZE &
+GUNICORN_PID=$!
+
+# "wait" returns a non-zero exit code when it is interrupted by a signal, which would abort the
+# script because of "set -o errexit". Waiting for the huey consumer after gunicorn has terminated
+# gives it the chance to shut down gracefully before this script and with it the container exits.
+wait "$GUNICORN_PID" || true
+wait "$HUEY_PID" || true