Merge branch 'insights-instrumentation' into event-sampling

Author: rabidpraxis

honeybadger/contrib/celery.py

@@ -1,11 +1,11 @@
-import threading
 import time
+import logging
 
 from honeybadger import honeybadger
 from honeybadger.plugins import Plugin, default_plugin_manager
 from honeybadger.utils import filter_dict, extract_honeybadger_config, get_duration
 
-_listener_started = False
+logger = logging.getLogger(__name__)
 
 
 class CeleryPlugin(Plugin):
@@ -67,7 +67,13 @@ def init_app(self):
         """
         Initialize honeybadger and listen for errors.
         """
-        from celery.signals import task_failure, task_postrun, task_prerun, worker_ready
+        from celery.signals import (
+            task_failure,
+            task_postrun,
+            task_prerun,
+            before_task_publish,
+            worker_process_init,
+        )
 
         self._task_starts = {}
         self._initialize_honeybadger(self.app.conf)
@@ -79,9 +85,9 @@ def init_app(self):
         if honeybadger.config.insights_enabled:
             # Enable task events, as we need to listen to
             # task-finished events
-            self.app.conf.worker_send_task_events = True
+            worker_process_init.connect(self._on_worker_process_init, weak=False)
             task_prerun.connect(self._on_task_prerun, weak=False)
-            worker_ready.connect(self._start_task_event_listener, weak=False)
+            before_task_publish.connect(self._on_before_task_publish, weak=False)
 
     def _initialize_honeybadger(self, config):
         """
@@ -96,35 +102,34 @@ def _initialize_honeybadger(self, config):
         honeybadger.configure(**config_kwargs)
         honeybadger.config.set_12factor_config()  # environment should override celery settings
 
-    def _start_task_event_listener(self, *args, **kwargs):
-        # only start the listener once
-        global _listener_started
-        if _listener_started:
-            return
-        _listener_started = True
-
-        from celery.events import EventReceiver  # type: ignore[import]
-
-        def run():
-            with self.app.connection() as conn:
-                recv = EventReceiver(
-                    conn, handlers={"task-finished": self._on_task_finished}
-                )
-                recv.capture(limit=None, timeout=None, wakeup=False)
-
-        self._listen_thread = threading.Thread(target=run, daemon=True)
-        self._listen_thread.start()
-
-    def _on_task_finished(self, payload, **kwargs):
-        honeybadger.event("celery.task_finished", payload["payload"])
+    def _on_worker_process_init(self, *args, **kwargs):
+        # Restart the events worker to ensure it is running in the new worker
+        # process.
+        try:
+            honeybadger.events_worker.restart()
+        except Exception as e:
+            logger.warning(f"Warning: Failed to restart Honeybadger events worker: {e}")
+
+    def _on_before_task_publish(self, sender=None, body=None, headers=None, **kwargs):
+        # Inject Honeybadger event context into task headers
+        if headers is not None:
+            current_context = honeybadger._get_event_context()
+            if current_context:
+                headers["honeybadger_event_context"] = current_context
 
     def _on_task_prerun(self, task_id=None, task=None, *args, **kwargs):
         self._task_starts[task_id] = time.time()
 
+        if task:
+            context = getattr(task.request, "honeybadger_event_context", None)
+            if context:
+                honeybadger.set_event_context(context)
+
     def _on_task_postrun(self, task_id=None, task=None, *args, **kwargs):
         """
         Callback executed after a task is finished.
         """
+
         insights_config = honeybadger.config.insights_config
 
         exclude = insights_config.celery.exclude_tasks
@@ -159,7 +164,7 @@ def _on_task_postrun(self, task_id=None, task=None, *args, **kwargs):
                     remove_keys=True,
                 )
 
-            task.send_event("task-finished", payload=payload)
+            honeybadger.event("celery.task_finished", payload=payload)
 
         honeybadger.reset_context()
 
@@ -180,13 +185,15 @@ def tearDown(self):
             task_failure.disconnect(self._on_task_failure)
 
         if honeybadger.config.insights_enabled:
-            from celery.signals import worker_ready, task_prerun
+            from celery.signals import (
+                task_prerun,
+                worker_process_init,
+                before_task_publish,
+            )
 
             task_prerun.disconnect(self._on_task_prerun)
-            worker_ready.disconnect(self._start_task_event_listener)
-
-        if hasattr(self, "_listen_thread"):
-            self._listen_thread.join(timeout=1)
+            worker_process_init.disconnect(self._on_worker_process_init, weak=False)
+            before_task_publish.disconnect(self._on_before_task_publish, weak=False)
 
     # Keep the misspelled method for backward compatibility
     def tearDowm(self):

honeybadger/contrib/django.py

@@ -1,12 +1,18 @@
 from __future__ import absolute_import
 import re
 import time
+import uuid
 
 from six import iteritems
 
 from honeybadger import honeybadger
 from honeybadger.plugins import Plugin, default_plugin_manager
-from honeybadger.utils import filter_dict, filter_env_vars, get_duration
+from honeybadger.utils import (
+    filter_dict,
+    filter_env_vars,
+    get_duration,
+    sanitize_request_id,
+)
 from honeybadger.contrib.db import DBHoneybadger
 
 try:
@@ -138,6 +144,7 @@ def __call__(self, request):
         set_request(request)
         start_time = time.time()
         honeybadger.begin_request(request)
+        self._set_request_id(request)
         response = self.get_response(request)
 
         if (
@@ -151,6 +158,19 @@ def __call__(self, request):
 
         return response
 
+    def _set_request_id(self, request):
+        # Attempt to get request ID from various sources
+        request_id = (
+            getattr(request, "id", None)
+            or getattr(request, "request_id", None)
+            or request.headers.get("X-Request-ID", None)
+        )
+        request_id = sanitize_request_id(request_id)
+        if not request_id:
+            request_id = str(uuid.uuid4())
+
+        honeybadger.set_event_context(request_id=request_id)
+
     def _patch_cursor(self):
         from django.db.backends.utils import CursorWrapper
 

honeybadger/contrib/flask.py

@@ -3,6 +3,7 @@
 
 import logging
 import time
+import uuid
 
 from honeybadger import honeybadger
 from honeybadger.plugins import Plugin, default_plugin_manager
@@ -11,6 +12,7 @@
     filter_env_vars,
     get_duration,
     extract_honeybadger_config,
+    sanitize_request_id,
 )
 from honeybadger.contrib.db import DBHoneybadger
 from six import iteritems
@@ -201,6 +203,12 @@ def _initialize_honeybadger(self, config):
     def _handle_request_started(self, sender, *args, **kwargs):
         from flask import request
 
+        request_id = sanitize_request_id(request.headers.get("X-Request-ID"))
+        if not request_id:
+            request_id = str(uuid.uuid4())
+
+        honeybadger.set_event_context(request_id=request_id)
+
         _request_info.set(
             {
                 "start_time": time.time(),

honeybadger/core.py

@@ -83,6 +83,8 @@ def notify(
         merged_ctx = {**base, **(context or {})}
         merged_tags = list({*tag_ctx, *(tags or [])})
 
+        request_id = self._get_event_context().get("request_id", None)
+
         notice = Notice(
             exception=exception,
             error_class=error_class,
@@ -91,6 +93,7 @@ def notify(
             fingerprint=fingerprint,
             tags=merged_tags,
             config=self.config,
+            request_id=request_id,
         )
         return self._send_notice(notice)
 

honeybadger/events_worker.py

@@ -1,3 +1,4 @@
+import os
 import time
 import threading
 import logging
@@ -41,12 +42,29 @@ def __init__(
 
         self._thread = threading.Thread(
             target=self._run,
-            name="honeybadger-events-worker",
+            name=f"honeybadger-events-worker-{os.getpid()}",
             daemon=True,
         )
         self._thread.start()
         self.log.debug("Events worker started")
 
+    def restart(self):
+        """Restart the batch worker thread (useful after process forking)"""
+        if hasattr(self, "_thread") and self._thread and self._thread.is_alive():
+            self.shutdown()
+
+        # Reset state
+        self._stop = False
+
+        self._thread = threading.Thread(
+            target=self._run,
+            name=f"honeybadger-events-worker-{os.getpid()}",
+            daemon=True,
+        )
+        self._thread.start()
+
+        return self._thread.is_alive()
+
     def push(self, event: Event) -> bool:
         with self._cond:
             if self._all_events_queued_len() >= self.config.events_max_queue_size:

honeybadger/notice.py

@@ -11,6 +11,7 @@ def __init__(self, *args, **kwargs):
         self.fingerprint = kwargs.get("fingerprint", None)
         self.config = kwargs.get("config", None)
         self.context = kwargs.get("context", {})
+        self.request_id = kwargs.get("request_id", None)
         self.tags = self._construct_tags(kwargs.get("tags", []))
 
         self._process_exception()
@@ -34,6 +35,7 @@ def payload(self):
             context=self.context,
             tags=self.tags,
             config=self.config,
+            correlation_context=self._correlation_context(),
         )
 
     def excluded_exception(self):
@@ -49,6 +51,11 @@ def excluded_exception(self):
                 return True
         return False
 
+    def _correlation_context(self):
+        if self.request_id:
+            return {"request_id": self.request_id}
+        return None
+
     def _construct_tags(self, tags):
         """
         Accepts either:

honeybadger/payload.py

@@ -158,6 +158,7 @@ def create_payload(
     config=None,
     context=None,
     fingerprint=None,
+    correlation_context=None,
     tags=None,
 ):
     # if using local_variables get them
@@ -188,4 +189,7 @@ def create_payload(
         "request": {"context": context, "local_variables": local_variables},
     }
 
+    if correlation_context:
+        payload["correlation_context"] = correlation_context
+
     return default_plugin_manager.generate_payload(payload, config, context)