Celery Beat monitoring without restarting the Beat process (#2001)

New way to instrument Celery Beat tasks:
- Patch the apply_entry function of Celery beat (see _patch_beat_apply_entry) to send the IN_PROGRESS checkin.
- Put the Sentry monitor config in the schedule_task headers.
- Retrieve task.request.get("properties") in _get_headers() to include the headers put into the schedule_task in step the previous bullet point.
- Use the success/failure/retry hooks as before.

Author: antonpirker

sentry_sdk/crons/api.py

@@ -41,16 +41,15 @@ def capture_checkin(
     monitor_config=None,
 ):
     # type: (Optional[str], Optional[str], Optional[str], Optional[float], Optional[Dict[str, Any]]) -> str
-    hub = Hub.current
-
-    check_in_id = check_in_id or uuid.uuid4().hex
     check_in_event = _create_check_in_event(
         monitor_slug=monitor_slug,
         check_in_id=check_in_id,
         status=status,
         duration_s=duration,
         monitor_config=monitor_config,
     )
+
+    hub = Hub.current
     hub.capture_event(check_in_event)
 
     return check_in_event["check_in_id"]

sentry_sdk/integrations/celery.py

@@ -1,9 +1,6 @@
 from __future__ import absolute_import
 
 import sys
-import shutil
-import functools
-import tempfile
 
 from sentry_sdk.consts import OP
 from sentry_sdk._compat import reraise
@@ -25,7 +22,6 @@
     from typing import Any
     from typing import Callable
     from typing import Dict
-    from typing import List
     from typing import Optional
     from typing import Tuple
     from typing import TypeVar
@@ -40,7 +36,7 @@
     from celery import VERSION as CELERY_VERSION
     from celery import Task, Celery
     from celery.app.trace import task_has_custom
-    from celery.beat import Service  # type: ignore
+    from celery.beat import Scheduler  # type: ignore
     from celery.exceptions import (  # type: ignore
         Ignore,
         Reject,
@@ -49,8 +45,6 @@
     )
     from celery.schedules import crontab, schedule  # type: ignore
     from celery.signals import (  # type: ignore
-        beat_init,
-        task_prerun,
         task_failure,
         task_success,
         task_retry,
@@ -68,9 +62,11 @@ class CeleryIntegration(Integration):
     def __init__(self, propagate_traces=True, monitor_beat_tasks=False):
         # type: (bool, bool) -> None
         self.propagate_traces = propagate_traces
+        self.monitor_beat_tasks = monitor_beat_tasks
 
         if monitor_beat_tasks:
-            _patch_celery_beat_tasks()
+            _patch_beat_apply_entry()
+            _setup_celery_beat_signals()
 
     @staticmethod
     def setup_once():
@@ -131,6 +127,12 @@ def apply_async(*args, **kwargs):
             ) as span:
                 with capture_internal_exceptions():
                     headers = dict(hub.iter_trace_propagation_headers(span))
+                    if integration.monitor_beat_tasks:
+                        headers.update(
+                            {
+                                "sentry-monitor-start-timestamp-s": "%.9f" % now(),
+                            }
+                        )
 
                     if headers:
                         # Note: kwargs can contain headers=None, so no setdefault!
@@ -320,12 +322,15 @@ def sentry_workloop(*args, **kwargs):
 
 def _get_headers(task):
     # type: (Task) -> Dict[str, Any]
-    headers = task.request.get("headers") or {}
+    headers = task.request.get("headers", {})
 
+    # flatten nested headers
     if "headers" in headers:
         headers.update(headers["headers"])
         del headers["headers"]
 
+    headers.update(task.request.get("properties", {}))
+
     return headers
 
 
@@ -387,123 +392,47 @@ def _get_monitor_config(celery_schedule, app):
     return monitor_config
 
 
-def _reinstall_patched_tasks(app, sender, add_updated_periodic_tasks):
-    # type: (Celery, Service, List[functools.partial[Any]]) -> None
-
-    # Stop Celery Beat
-    sender.stop()
-
-    # Update tasks to include Monitor information in headers
-    for add_updated_periodic_task in add_updated_periodic_tasks:
-        add_updated_periodic_task()
-
-    # Start Celery Beat (with new (cloned) schedule, because old one is still in use)
-    cloned_schedule = tempfile.NamedTemporaryFile(suffix="-patched-by-sentry-sdk")
-    with open(sender.schedule_filename, "rb") as original_schedule:
-        shutil.copyfileobj(original_schedule, cloned_schedule)
+def _patch_beat_apply_entry():
+    # type: () -> None
+    original_apply_entry = Scheduler.apply_entry
+
+    def sentry_apply_entry(*args, **kwargs):
+        # type: (*Any, **Any) -> None
+        scheduler, schedule_entry = args
+        app = scheduler.app
+
+        celery_schedule = schedule_entry.schedule
+        monitor_config = _get_monitor_config(celery_schedule, app)
+        monitor_name = schedule_entry.name
+
+        headers = schedule_entry.options.pop("headers", {})
+        headers.update(
+            {
+                "sentry-monitor-slug": monitor_name,
+                "sentry-monitor-config": monitor_config,
+            }
+        )
 
-    app.Beat(schedule=cloned_schedule.name).run()
+        check_in_id = capture_checkin(
+            monitor_slug=monitor_name,
+            monitor_config=monitor_config,
+            status=MonitorStatus.IN_PROGRESS,
+        )
+        headers.update({"sentry-monitor-check-in-id": check_in_id})
 
+        schedule_entry.options.update(headers)
+        return original_apply_entry(*args, **kwargs)
 
-# Nested functions do not work as Celery hook receiver,
-# so defining it here explicitly
-celery_beat_init = None
+    Scheduler.apply_entry = sentry_apply_entry
 
 
-def _patch_celery_beat_tasks():
+def _setup_celery_beat_signals():
     # type: () -> None
-
-    global celery_beat_init
-
-    def celery_beat_init(sender, **kwargs):
-        # type: (Service, Dict[Any, Any]) -> None
-
-        # Because we restart Celery Beat,
-        # make sure that this will not be called infinitely
-        beat_init.disconnect(celery_beat_init)
-
-        app = sender.app
-
-        add_updated_periodic_tasks = []
-
-        for name in sender.scheduler.schedule.keys():
-            # Ignore Celery's internal tasks
-            if name.startswith("celery."):
-                continue
-
-            monitor_name = name
-
-            schedule_entry = sender.scheduler.schedule[name]
-            celery_schedule = schedule_entry.schedule
-            monitor_config = _get_monitor_config(celery_schedule, app)
-
-            if monitor_config is None:
-                continue
-
-            headers = schedule_entry.options.pop("headers", {})
-            headers.update(
-                {
-                    "headers": {
-                        "sentry-monitor-slug": monitor_name,
-                        "sentry-monitor-config": monitor_config,
-                    },
-                }
-            )
-
-            task_signature = app.tasks.get(schedule_entry.task).s()
-            task_signature.set(headers=headers)
-
-            logger.debug(
-                "Set up Sentry Celery Beat monitoring for %s (%s)",
-                task_signature,
-                monitor_name,
-            )
-
-            add_updated_periodic_tasks.append(
-                functools.partial(
-                    app.add_periodic_task,
-                    celery_schedule,
-                    task_signature,
-                    args=schedule_entry.args,
-                    kwargs=schedule_entry.kwargs,
-                    name=schedule_entry.name,
-                    **(schedule_entry.options or {})
-                )
-            )
-
-        _reinstall_patched_tasks(app, sender, add_updated_periodic_tasks)
-
-    beat_init.connect(celery_beat_init)
-    task_prerun.connect(crons_task_before_run)
     task_success.connect(crons_task_success)
     task_failure.connect(crons_task_failure)
     task_retry.connect(crons_task_retry)
 
 
-def crons_task_before_run(sender, **kwargs):
-    # type: (Task, Dict[Any, Any]) -> None
-    logger.debug("celery_task_before_run %s", sender)
-    headers = _get_headers(sender)
-
-    if "sentry-monitor-slug" not in headers:
-        return
-
-    monitor_config = headers.get("sentry-monitor-config", {})
-
-    start_timestamp_s = now()
-
-    check_in_id = capture_checkin(
-        monitor_slug=headers["sentry-monitor-slug"],
-        monitor_config=monitor_config,
-        status=MonitorStatus.IN_PROGRESS,
-    )
-
-    headers.update({"sentry-monitor-check-in-id": check_in_id})
-    headers.update({"sentry-monitor-start-timestamp-s": start_timestamp_s})
-
-    sender.s().set(headers=headers)
-
-
 def crons_task_success(sender, **kwargs):
     # type: (Task, Dict[Any, Any]) -> None
     logger.debug("celery_task_success %s", sender)
@@ -514,7 +443,7 @@ def crons_task_success(sender, **kwargs):
 
     monitor_config = headers.get("sentry-monitor-config", {})
 
-    start_timestamp_s = headers["sentry-monitor-start-timestamp-s"]
+    start_timestamp_s = float(headers["sentry-monitor-start-timestamp-s"])
 
     capture_checkin(
         monitor_slug=headers["sentry-monitor-slug"],
@@ -535,7 +464,7 @@ def crons_task_failure(sender, **kwargs):
 
     monitor_config = headers.get("sentry-monitor-config", {})
 
-    start_timestamp_s = headers["sentry-monitor-start-timestamp-s"]
+    start_timestamp_s = float(headers["sentry-monitor-start-timestamp-s"])
 
     capture_checkin(
         monitor_slug=headers["sentry-monitor-slug"],
@@ -556,7 +485,7 @@ def crons_task_retry(sender, **kwargs):
 
     monitor_config = headers.get("sentry-monitor-config", {})
 
-    start_timestamp_s = headers["sentry-monitor-start-timestamp-s"]
+    start_timestamp_s = float(headers["sentry-monitor-start-timestamp-s"])
 
     capture_checkin(
         monitor_slug=headers["sentry-monitor-slug"],

tests/integrations/celery/test_celery_beat_crons.py

@@ -1,4 +1,3 @@
-import tempfile
 import mock
 
 import pytest
@@ -9,8 +8,6 @@
     _get_headers,
     _get_humanized_interval,
     _get_monitor_config,
-    _reinstall_patched_tasks,
-    crons_task_before_run,
     crons_task_success,
     crons_task_failure,
     crons_task_retry,
@@ -70,42 +67,6 @@ def test_get_humanized_interval(seconds, expected_tuple):
     assert _get_humanized_interval(seconds) == expected_tuple
 
 
-def test_crons_task_before_run():
-    fake_task = mock.MagicMock()
-    fake_task.request = {
-        "headers": {
-            "sentry-monitor-slug": "test123",
-            "sentry-monitor-config": {
-                "schedule": {
-                    "type": "interval",
-                    "value": 3,
-                    "unit": "day",
-                },
-                "timezone": "Europe/Vienna",
-            },
-            "sentry-monitor-some-future-key": "some-future-value",
-        },
-    }
-
-    with mock.patch(
-        "sentry_sdk.integrations.celery.capture_checkin"
-    ) as mock_capture_checkin:
-        crons_task_before_run(fake_task)
-
-        mock_capture_checkin.assert_called_once_with(
-            monitor_slug="test123",
-            monitor_config={
-                "schedule": {
-                    "type": "interval",
-                    "value": 3,
-                    "unit": "day",
-                },
-                "timezone": "Europe/Vienna",
-            },
-            status=MonitorStatus.IN_PROGRESS,
-        )
-
-
 def test_crons_task_success():
     fake_task = mock.MagicMock()
     fake_task.request = {
@@ -273,35 +234,3 @@ def test_get_monitor_config_default_timezone():
     monitor_config = _get_monitor_config(celery_schedule, app)
 
     assert monitor_config["timezone"] == "UTC"
-
-
-def test_reinstall_patched_tasks():
-    fake_beat = mock.MagicMock()
-    fake_beat.run = mock.MagicMock()
-
-    app = mock.MagicMock()
-    app.Beat = mock.MagicMock(return_value=fake_beat)
-
-    sender = mock.MagicMock()
-    sender.schedule_filename = "test_schedule_filename"
-    sender.stop = mock.MagicMock()
-
-    add_updated_periodic_tasks = [mock.MagicMock(), mock.MagicMock(), mock.MagicMock()]
-
-    mock_open = mock.Mock(return_value=tempfile.NamedTemporaryFile())
-
-    with mock.patch("sentry_sdk.integrations.celery.open", mock_open):
-        with mock.patch(
-            "sentry_sdk.integrations.celery.shutil.copyfileobj"
-        ) as mock_copyfileobj:
-            _reinstall_patched_tasks(app, sender, add_updated_periodic_tasks)
-
-            sender.stop.assert_called_once_with()
-
-            add_updated_periodic_tasks[0].assert_called_once_with()
-            add_updated_periodic_tasks[1].assert_called_once_with()
-            add_updated_periodic_tasks[2].assert_called_once_with()
-
-            mock_copyfileobj.assert_called_once()
-
-            fake_beat.run.assert_called_once_with()