celery_backend/config: Enable to set Sentry DSN per task type

Add a task_prerun celery signal handler in order to set Sentry DSN based
on task name. The mapping between a task name and its DSN must be stored
in configuration under a "task_sentry_dsn" key.

For this feature to work, no SWH_SENTRY_DSN environment variable should
be defined as it overrides the sentry_dsn value passed to init_sentry
function.

Related to swh/meta#4949.

swh/scheduler/celery_backend/config.py

@@ -12,7 +12,7 @@ from typing import Any, Dict, Optional
 import urllib.parse
 
 from celery import Celery
-from celery.signals import celeryd_after_setup, setup_logging, worker_init
+from celery.signals import celeryd_after_setup, setup_logging, task_prerun, worker_init
 from celery.utils.log import ColorFormatter
 from celery.worker.control import Panel
 from kombu import Exchange, Queue
@@ -163,19 +163,23 @@ def setup_queues_and_tasks(sender, instance, **kwargs):
     instance.app.conf["worker_name"] = sender
 
 
-@worker_init.connect
-@_print_errors
-def on_worker_init(*args, **kwargs):
+def _init_sentry(sentry_dsn):
     try:
         from sentry_sdk.integrations.celery import CeleryIntegration
     except ImportError:
         integrations = []
     else:
         integrations = [CeleryIntegration()]
-    sentry_dsn = None  # will be set in `init_sentry` function
     init_sentry(sentry_dsn, integrations=integrations)
 
 
+@worker_init.connect
+@_print_errors
+def on_worker_init(*args, **kwargs):
+    sentry_dsn = None  # will be set in `init_sentry` function
+    _init_sentry(sentry_dsn)
+
+
 @Panel.register
 def monotonic(state):
     """Get the current value for the monotonic clock"""
@@ -373,5 +377,12 @@ def build_app(config=None):
 
 app = build_app(CONFIG)
 
+
+@task_prerun.connect
+def celery_task_received(task_id, task, *args, **kwargs):
+    sentry_dsn = CONFIG.get("task_sentry_dsn", {}).get(task.name)
+    _init_sentry(sentry_dsn)
+
+
 # XXX for BW compat
 Celery.get_queue_length = get_queue_length

swh/scheduler/tests/test_celery_tasks.py

@@ -5,6 +5,7 @@
 
 """Module in charge of testing the scheduler runner module"""
 
+import copy
 from itertools import count
 from time import sleep
 
@@ -12,6 +13,7 @@ from celery.result import AsyncResult, GroupResult
 from kombu import Exchange, Queue
 import pytest
 
+from swh.scheduler.celery_backend import config
 from swh.scheduler.celery_backend.runner import run_ready_tasks
 from swh.scheduler.tests.tasks import (
     TASK_ADD,
@@ -281,3 +283,24 @@ def test_statsd_with_status(
         "swh_task_success_count:1|c|"
         "#task:swh.scheduler.tests.tasks.echo,worker:unknown worker"
     )
+
+
+def test_sentry_dispatch(
+    swh_scheduler_celery_app, swh_scheduler_celery_worker, sentry_events, mocker
+):
+    task_dsn = {
+        TASK_PING: "https://public@sentry.softwareheritage.org/1",
+        TASK_ERROR: "https://public@sentry.softwareheritage.org/2",
+    }
+    updated_config = copy.deepcopy(config.CONFIG)
+    updated_config.update({"task_sentry_dsn": task_dsn})
+    mocker.patch.object(config, "CONFIG", updated_config)
+
+    init_sentry = mocker.patch.object(config, "init_sentry")
+
+    for task in (TASK_PING, TASK_ERROR):
+        res = swh_scheduler_celery_app.send_task(task)
+        assert res
+        res.wait(propagate=False)
+        init_sentry.assert_called()
+        assert task_dsn[task] in init_sentry.call_args_list[-1][0]