root: celery refactor (#6095)

* root: celery refactor cleanup deprecation messages by configuring celery with a single object run celery as django management command Signed-off-by: Jens Langhammer <jens@goauthentik.io> * improve debug experience Signed-off-by: Jens Langhammer <jens@goauthentik.io> * fix lint Signed-off-by: Jens Langhammer <jens@goauthentik.io> * add debugpy to dev dependencies Signed-off-by: Jens Langhammer <jens@goauthentik.io> * fix task_always_eager Signed-off-by: Jens Langhammer <jens@goauthentik.io> --------- Signed-off-by: Jens Langhammer <jens@goauthentik.io>
2023-06-28 16:44:50 +02:00 · 2023-06-28 16:44:50 +02:00 · a987846c76
parent 35e2b648ba
commit a987846c76
13 changed files with 188 additions and 46 deletions
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@ -0,0 +1,27 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: PDB attach Server",
+            "type": "python",
+            "request": "attach",
+            "connect": {
+                "host": "localhost",
+                "port": 6800
+            },
+            "justMyCode": true,
+            "django": true
+        },
+        {
+            "name": "Python: PDB attach Worker",
+            "type": "python",
+            "request": "attach",
+            "connect": {
+                "host": "localhost",
+                "port": 6900
+            },
+            "justMyCode": true,
+            "django": true
+        },
+    ]
+}
--- a/authentik/admin/api/workers.py
+++ b/authentik/admin/api/workers.py
@ -19,7 +19,7 @@ class WorkerView(APIView):
    def get(self, request: Request) -> Response:
        """Get currently connected worker count."""
        count = len(CELERY_APP.control.ping(timeout=0.5))
-        # In debug we run with `CELERY_TASK_ALWAYS_EAGER`, so tasks are ran on the main process
+        # In debug we run with `task_always_eager`, so tasks are ran on the main process
        if settings.DEBUG:  # pragma: no cover
            count += 1
        return Response({"count": count})
--- a/authentik/core/management/commands/worker.py
+++ b/authentik/core/management/commands/worker.py
@ -0,0 +1,40 @@
+"""Run worker"""
+from sys import exit as sysexit
+from tempfile import tempdir
+
+from celery.apps.worker import Worker
+from django.core.management.base import BaseCommand
+from django.db import close_old_connections
+from structlog.stdlib import get_logger
+
+from authentik.lib.config import CONFIG
+from authentik.root.celery import CELERY_APP
+
+LOGGER = get_logger()
+
+
+class Command(BaseCommand):
+    """Run worker"""
+
+    def handle(self, **options):
+        close_old_connections()
+        if CONFIG.y_bool("remote_debug"):
+            import debugpy
+
+            debugpy.listen(("0.0.0.0", 6900))  # nosec
+        worker: Worker = CELERY_APP.Worker(
+            no_color=False,
+            quiet=True,
+            optimization="fair",
+            max_tasks_per_child=1,
+            autoscale=(3, 1),
+            task_events=True,
+            beat=True,
+            schedule_filename=f"{tempdir}/celerybeat-schedule",
+            queues=["authentik", "authentik_scheduled", "authentik_events"],
+        )
+        for task in CELERY_APP.tasks:
+            LOGGER.debug("Registered task", task=task)
+
+        worker.start()
+        sysexit(worker.exitcode)
--- a/authentik/events/monitored_tasks.py
+++ b/authentik/events/monitored_tasks.py
@ -41,6 +41,7 @@ class TaskResult:

    def with_error(self, exc: Exception) -> "TaskResult":
        """Since errors might not always be pickle-able, set the traceback"""
+        # TODO: Mark exception somehow so that is rendered as <pre> in frontend
        self.messages.append(exception_to_string(exc))
        return self

--- a/authentik/lib/default.yml
+++ b/authentik/lib/default.yml
@ -26,6 +26,7 @@ redis:
  cache_timeout_reputation: 300

 debug: false
+remote_debug: false

 log_level: info

--- a/authentik/root/celery.py
+++ b/authentik/root/celery.py
@ -130,11 +130,7 @@ class LivenessProbe(bootsteps.StartStopStep):
        HEARTBEAT_FILE.touch()


-# Using a string here means the worker doesn't have to serialize
-# the configuration object to child processes.
-# - namespace='CELERY' means all celery-related configuration keys
-#   should have a `CELERY_` prefix.
-CELERY_APP.config_from_object(settings, namespace="CELERY")
+CELERY_APP.config_from_object(settings.CELERY)

 # Load task modules from all registered Django app configs.
 CELERY_APP.autodiscover_tasks()
--- a/authentik/root/settings.py
+++ b/authentik/root/settings.py
@ -182,13 +182,13 @@ REST_FRAMEWORK = {
    },
 }

-REDIS_PROTOCOL_PREFIX = "redis://"
-REDIS_CELERY_TLS_REQUIREMENTS = ""
+_redis_protocol_prefix = "redis://"
+_redis_celery_tls_requirements = ""
 if CONFIG.y_bool("redis.tls", False):
-    REDIS_PROTOCOL_PREFIX = "rediss://"
-    REDIS_CELERY_TLS_REQUIREMENTS = f"?ssl_cert_reqs={CONFIG.y('redis.tls_reqs')}"
+    _redis_protocol_prefix = "rediss://"
+    _redis_celery_tls_requirements = f"?ssl_cert_reqs={CONFIG.y('redis.tls_reqs')}"
 _redis_url = (
-    f"{REDIS_PROTOCOL_PREFIX}:"
+    f"{_redis_protocol_prefix}:"
    f"{quote_plus(CONFIG.y('redis.password'))}@{quote_plus(CONFIG.y('redis.host'))}:"
    f"{int(CONFIG.y('redis.port'))}"
 )
@ -326,12 +326,11 @@ USE_TZ = True

 LOCALE_PATHS = ["./locale"]

-# Celery settings
-# Add a 10 minute timeout to all Celery tasks.
-CELERY_TASK_SOFT_TIME_LIMIT = 600
-CELERY_WORKER_MAX_TASKS_PER_CHILD = 50
-CELERY_WORKER_CONCURRENCY = 2
-CELERY_BEAT_SCHEDULE = {
+CELERY = {
+    "task_soft_time_limit": 600,
+    "worker_max_tasks_per_child": 50,
+    "worker_concurrency": 2,
+    "beat_schedule": {
        "clean_expired_models": {
            "task": "authentik.core.tasks.clean_expired_models",
            "schedule": crontab(minute="2-59/5"),
@ -342,11 +341,12 @@ CELERY_BEAT_SCHEDULE = {
            "schedule": crontab(minute="9-59/5"),
            "options": {"queue": "authentik_scheduled"},
        },
+    },
+    "task_create_missing_queues": True,
+    "task_default_queue": "authentik",
+    "broker_url": f"{_redis_url}/{CONFIG.y('redis.db')}{_redis_celery_tls_requirements}",
+    "result_backend": f"{_redis_url}/{CONFIG.y('redis.db')}{_redis_celery_tls_requirements}",
 }
-CELERY_TASK_CREATE_MISSING_QUEUES = True
-CELERY_TASK_DEFAULT_QUEUE = "authentik"
-CELERY_BROKER_URL = f"{_redis_url}/{CONFIG.y('redis.db')}{REDIS_CELERY_TLS_REQUIREMENTS}"
-CELERY_RESULT_BACKEND = f"{_redis_url}/{CONFIG.y('redis.db')}{REDIS_CELERY_TLS_REQUIREMENTS}"

 # Sentry integration
 env = get_env()
@ -455,7 +455,7 @@ _DISALLOWED_ITEMS = [
    "INSTALLED_APPS",
    "MIDDLEWARE",
    "AUTHENTICATION_BACKENDS",
-    "CELERY_BEAT_SCHEDULE",
+    "CELERY",
 ]


@ -466,7 +466,7 @@ def _update_settings(app_path: str):
        INSTALLED_APPS.extend(getattr(settings_module, "INSTALLED_APPS", []))
        MIDDLEWARE.extend(getattr(settings_module, "MIDDLEWARE", []))
        AUTHENTICATION_BACKENDS.extend(getattr(settings_module, "AUTHENTICATION_BACKENDS", []))
-        CELERY_BEAT_SCHEDULE.update(getattr(settings_module, "CELERY_BEAT_SCHEDULE", {}))
+        CELERY["beat_schedule"].update(getattr(settings_module, "CELERY_BEAT_SCHEDULE", {}))
        for _attr in dir(settings_module):
            if not _attr.startswith("__") and _attr not in _DISALLOWED_ITEMS:
                globals()[_attr] = getattr(settings_module, _attr)
@ -482,7 +482,7 @@ for _app in INSTALLED_APPS:
 _update_settings("data.user_settings")

 if DEBUG:
-    CELERY_TASK_ALWAYS_EAGER = True
+    CELERY["task_always_eager"] = True
    os.environ[ENV_GIT_HASH_KEY] = "dev"
    INSTALLED_APPS.append("silk")
    SILKY_PYTHON_PROFILER = True
--- a/authentik/root/test_runner.py
+++ b/authentik/root/test_runner.py
@ -30,7 +30,7 @@ class PytestTestRunner:  # pragma: no cover
            self.args.append(f"--randomly-seed={kwargs['randomly_seed']}")

        settings.TEST = True
-        settings.CELERY_TASK_ALWAYS_EAGER = True
+        settings.CELERY["task_always_eager"] = True
        CONFIG.y_set("avatars", "none")
        CONFIG.y_set("geoip", "tests/GeoLite2-City-Test.mmdb")
        CONFIG.y_set("blueprints_dir", "./blueprints")
--- a/docker-compose.override.yml
+++ b/docker-compose.override.yml
@ -0,0 +1,36 @@
+# This file is used for development and debugging, and should not be used for production instances
+
+version: '3.5'
+
+services:
+  flower:
+    image: ${AUTHENTIK_IMAGE:-ghcr.io/goauthentik/server}:${AUTHENTIK_TAG:-2023.5.4}
+    restart: unless-stopped
+    command: worker-status
+    environment:
+      AUTHENTIK_REDIS__HOST: redis
+      AUTHENTIK_POSTGRESQL__HOST: postgresql
+      AUTHENTIK_POSTGRESQL__USER: ${PG_USER:-authentik}
+      AUTHENTIK_POSTGRESQL__NAME: ${PG_DB:-authentik}
+      AUTHENTIK_POSTGRESQL__PASSWORD: ${PG_PASS}
+    env_file:
+      - .env
+    ports:
+      - "9001:9000"
+    depends_on:
+      - postgresql
+      - redis
+  server:
+    environment:
+      AUTHENTIK_REMOTE_DEBUG: "true"
+      PYDEVD_THREAD_DUMP_ON_WARN_EVALUATION_TIMEOUT: "true"
+    ports:
+      - 6800:6800
+  worker:
+    environment:
+      CELERY_RDB_HOST: "0.0.0.0"
+      CELERY_RDBSIG: "1"
+      AUTHENTIK_REMOTE_DEBUG: "true"
+      PYDEVD_THREAD_DUMP_ON_WARN_EVALUATION_TIMEOUT: "true"
+    ports:
+      - 6900:6900
--- a/lifecycle/ak
+++ b/lifecycle/ak
@ -54,6 +54,16 @@ function cleanup {
    rm -f ${MODE_FILE}
 }

+function prepare_debug {
+    pip install --no-cache-dir -r /requirements-dev.txt
+    touch /unittest.xml
+    chown authentik:authentik /unittest.xml
+}
+
+if [[ "${AUTHENTIK_REMOTE_DEBUG}" == "true" ]]; then
+    prepare_debug
+fi
+
 if [[ "$1" == "server" ]]; then
    wait_for_db
    set_mode "server"
@ -67,7 +77,7 @@ if [[ "$1" == "server" ]]; then
 elif [[ "$1" == "worker" ]]; then
    wait_for_db
    set_mode "worker"
-    check_if_root "celery -A authentik.root.celery worker -Ofair --max-tasks-per-child=1 --autoscale 3,1 -E -B -s /tmp/celerybeat-schedule -Q authentik,authentik_scheduled,authentik_events"
+    check_if_root "python -m manage worker"
 elif [[ "$1" == "worker-status" ]]; then
    wait_for_db
    celery -A authentik.root.celery flower \
@ -75,9 +85,7 @@ elif [[ "$1" == "worker-status" ]]; then
 elif [[ "$1" == "bash" ]]; then
    /bin/bash
 elif [[ "$1" == "test-all" ]]; then
-    pip install --no-cache-dir -r /requirements-dev.txt
-    touch /unittest.xml
-    chown authentik:authentik /unittest.xml
+    prepare_debug
    check_if_root "python -m manage test authentik"
 elif [[ "$1" == "healthcheck" ]]; then
    run_authentik healthcheck $(cat $MODE_FILE)
--- a/lifecycle/gunicorn.conf.py
+++ b/lifecycle/gunicorn.conf.py
@ -157,3 +157,8 @@ if not CONFIG.y_bool("disable_startup_analytics", False):
        # pylint: disable=broad-exception-caught
        except Exception:  # nosec
            pass
+
+if CONFIG.y_bool("remote_debug"):
+    import debugpy
+
+    debugpy.listen(("0.0.0.0", 6800))  # nosec
--- a/poetry.lock
+++ b/poetry.lock
@ -1066,6 +1066,33 @@ twisted = {version = ">=22.4", extras = ["tls"]}
 [package.extras]
 tests = ["django", "hypothesis", "pytest", "pytest-asyncio"]

+[[package]]
+name = "debugpy"
+version = "1.6.7"
+description = "An implementation of the Debug Adapter Protocol for Python"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "debugpy-1.6.7-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:b3e7ac809b991006ad7f857f016fa92014445085711ef111fdc3f74f66144096"},
+    {file = "debugpy-1.6.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e3876611d114a18aafef6383695dfc3f1217c98a9168c1aaf1a02b01ec7d8d1e"},
+    {file = "debugpy-1.6.7-cp310-cp310-win32.whl", hash = "sha256:33edb4afa85c098c24cc361d72ba7c21bb92f501104514d4ffec1fb36e09c01a"},
+    {file = "debugpy-1.6.7-cp310-cp310-win_amd64.whl", hash = "sha256:ed6d5413474e209ba50b1a75b2d9eecf64d41e6e4501977991cdc755dc83ab0f"},
+    {file = "debugpy-1.6.7-cp37-cp37m-macosx_10_15_x86_64.whl", hash = "sha256:38ed626353e7c63f4b11efad659be04c23de2b0d15efff77b60e4740ea685d07"},
+    {file = "debugpy-1.6.7-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:279d64c408c60431c8ee832dfd9ace7c396984fd7341fa3116aee414e7dcd88d"},
+    {file = "debugpy-1.6.7-cp37-cp37m-win32.whl", hash = "sha256:dbe04e7568aa69361a5b4c47b4493d5680bfa3a911d1e105fbea1b1f23f3eb45"},
+    {file = "debugpy-1.6.7-cp37-cp37m-win_amd64.whl", hash = "sha256:f90a2d4ad9a035cee7331c06a4cf2245e38bd7c89554fe3b616d90ab8aab89cc"},
+    {file = "debugpy-1.6.7-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:5224eabbbeddcf1943d4e2821876f3e5d7d383f27390b82da5d9558fd4eb30a9"},
+    {file = "debugpy-1.6.7-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bae1123dff5bfe548ba1683eb972329ba6d646c3a80e6b4c06cd1b1dd0205e9b"},
+    {file = "debugpy-1.6.7-cp38-cp38-win32.whl", hash = "sha256:9cd10cf338e0907fdcf9eac9087faa30f150ef5445af5a545d307055141dd7a4"},
+    {file = "debugpy-1.6.7-cp38-cp38-win_amd64.whl", hash = "sha256:aaf6da50377ff4056c8ed470da24632b42e4087bc826845daad7af211e00faad"},
+    {file = "debugpy-1.6.7-cp39-cp39-macosx_11_0_x86_64.whl", hash = "sha256:0679b7e1e3523bd7d7869447ec67b59728675aadfc038550a63a362b63029d2c"},
+    {file = "debugpy-1.6.7-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:de86029696e1b3b4d0d49076b9eba606c226e33ae312a57a46dca14ff370894d"},
+    {file = "debugpy-1.6.7-cp39-cp39-win32.whl", hash = "sha256:d71b31117779d9a90b745720c0eab54ae1da76d5b38c8026c654f4a066b0130a"},
+    {file = "debugpy-1.6.7-cp39-cp39-win_amd64.whl", hash = "sha256:c0ff93ae90a03b06d85b2c529eca51ab15457868a377c4cc40a23ab0e4e552a3"},
+    {file = "debugpy-1.6.7-py2.py3-none-any.whl", hash = "sha256:53f7a456bc50706a0eaabecf2d3ce44c4d5010e46dfc65b6b81a518b42866267"},
+    {file = "debugpy-1.6.7.zip", hash = "sha256:c4c2f0810fa25323abfdfa36cbbbb24e5c3b1a42cb762782de64439c575d67f2"},
+]
+
 [[package]]
 name = "deepmerge"
 version = "1.1.0"
@ -1686,13 +1713,13 @@ test = ["Cython (>=0.29.24,<0.30.0)"]

 [[package]]
 name = "humanize"
-version = "4.6.0"
+version = "4.7.0"
 description = "Python humanize utilities"
 optional = false
-python-versions = ">=3.7"
+python-versions = ">=3.8"
 files = [
-    {file = "humanize-4.6.0-py3-none-any.whl", hash = "sha256:401201aca462749773f02920139f302450cb548b70489b9b4b92be39fe3c3c50"},
-    {file = "humanize-4.6.0.tar.gz", hash = "sha256:5f1f22bc65911eb1a6ffe7659bd6598e33dcfeeb904eb16ee1e705a09bf75916"},
+    {file = "humanize-4.7.0-py3-none-any.whl", hash = "sha256:df7c429c2d27372b249d3f26eb53b07b166b661326e0325793e0a988082e3889"},
+    {file = "humanize-4.7.0.tar.gz", hash = "sha256:7ca0e43e870981fa684acb5b062deb307218193bca1a01f2b2676479df849b3a"},
 ]

 [package.extras]
@ -3419,13 +3446,13 @@ wsproto = ">=0.14"

 [[package]]
 name = "twilio"
-version = "8.3.0"
+version = "8.4.0"
 description = "Twilio API client and TwiML generator"
 optional = false
 python-versions = ">=3.7.0"
 files = [
-    {file = "twilio-8.3.0-py2.py3-none-any.whl", hash = "sha256:f8f4a26e7491e015777c2c12abcc068321f12302d081fc355df486601434c311"},
-    {file = "twilio-8.3.0.tar.gz", hash = "sha256:e76543b054f09304557d9bd0f9e3c21d09ca935d88f833788d43cab1f1fb67d1"},
+    {file = "twilio-8.4.0-py2.py3-none-any.whl", hash = "sha256:56b812b4d77dabcfdf7aa02aac966065e064beabd083621940856a6ee0d060ee"},
+    {file = "twilio-8.4.0.tar.gz", hash = "sha256:23fa599223d336a19d674394535d42bd1e260f7ca350a51d02b9d902370d76ef"},
 ]

 [package.dependencies]
@ -4159,4 +4186,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "7c78d6909ba8cc5b8fb41233e2506f0b919b71e263213068e479af706a9670ce"
+content-hash = "60a0e729895ebd44235e88e0414cc64e50c41736903ea61e6fb94a542dd2bb3c"
--- a/pyproject.toml
+++ b/pyproject.toml
@ -178,6 +178,7 @@ black = "*"
 bump2version = "*"
 colorama = "*"
 coverage = { extras = ["toml"], version = "*" }
+debugpy = "*"
 django-silk = "*"
 drf-jsonschema-serializer = "*"
 importlib-metadata = "*"