feat(insights): event_context (#214)

* refactor threadlocals

This now uses the ContextVar and a helper ContextStore class.

* Basic event_context

* More tests

* Fix other Notice usage

Author: rabidpraxis

dev-requirements.txt

@@ -15,6 +15,7 @@ psutil
 pylint
 pytest
 pytest-cov
+pytest-asyncio
 six
 sqlalchemy
 testfixtures

honeybadger/connection.py

@@ -3,7 +3,6 @@
 import threading
 
 from urllib.error import HTTPError, URLError
-from typing import Protocol
 from six.moves.urllib import request
 from six import b
 

honeybadger/context_store.py

@@ -0,0 +1,47 @@
+from contextvars import ContextVar
+from contextlib import contextmanager
+from typing import Any, Dict, Optional
+
+
+class ContextStore:
+    def __init__(self, name: str):
+        self._ctx: ContextVar[Optional[Dict[str, Any]]] = ContextVar(name, default=None)
+
+    def get(self) -> Dict[str, Any]:
+        data = self._ctx.get()
+        return {} if data is None else data.copy()
+
+    def clear(self) -> None:
+        self._ctx.set({})
+
+    def update(self, ctx: Optional[Dict[str, Any]] = None, **kwargs: Any) -> None:
+        """
+        Merge into the current context. Accepts either:
+          - update({'foo': 'bar'})
+          - update(foo='bar', baz=123)
+          - or both: update({'foo': 'bar'}, baz=123)
+        """
+        to_merge: Dict[str, Any] = {}
+        if ctx:
+            to_merge.update(ctx)
+        to_merge.update(kwargs)
+
+        new = self.get()
+        new.update(to_merge)
+        self._ctx.set(new)
+
+    @contextmanager
+    def override(self, ctx: Optional[Dict[str, Any]] = None, **kwargs: Any):
+        """
+        Temporarily merge these into context for the duration of the with-block.
+        """
+        to_merge: Dict[str, Any] = {}
+        if ctx:
+            to_merge.update(ctx)
+        to_merge.update(kwargs)
+
+        token = self._ctx.set({**self.get(), **to_merge})
+        try:
+            yield
+        finally:
+            self._ctx.reset(token)

honeybadger/core.py

@@ -2,29 +2,33 @@
 from contextlib import contextmanager
 import sys
 import logging
-import copy
-import time
 import datetime
 import atexit
+from typing import Optional, Dict, Any, List
 
 from honeybadger.plugins import default_plugin_manager
 import honeybadger.connection as connection
 import honeybadger.fake_connection as fake_connection
 from .events_worker import EventsWorker
 from .config import Configuration
 from .notice import Notice
+from .context_store import ContextStore
 
 logger = logging.getLogger("honeybadger")
 logger.addHandler(logging.NullHandler())
 
+error_context = ContextStore("honeybadger_error_context")
+event_context = ContextStore("honeybadger_event_context")
+
 
 class Honeybadger(object):
     TS_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 
     def __init__(self):
+        error_context.clear()
+        event_context.clear()
+
         self.config = Configuration()
-        self.thread_local = threading.local()
-        self.thread_local.context = {}
         self.events_worker = EventsWorker(
             self._connection(), self.config, logger=logging.getLogger("honeybadger")
         )
@@ -47,21 +51,16 @@ def _send_notice(self, notice):
 
         self._connection().send_notice(self.config, notice)
 
-    def _get_context(self):
-        return getattr(self.thread_local, "context", {})
-
-    def begin_request(self, request):
-        self.thread_local.context = self._get_context()
+    def begin_request(self, _):
+        error_context.clear()
+        event_context.clear()
 
     def wrap_excepthook(self, func):
         self.existing_except_hook = func
         sys.excepthook = self.exception_hook
 
     def exception_hook(self, type, exception, exc_traceback):
-        notice = Notice(
-            exception=exception, thread_local=self.thread_local, config=self.config
-        )
-        self._send_notice(notice)
+        self.notify(exception=exception)
         self.existing_except_hook(type, exception, exc_traceback)
 
     def shutdown(self):
@@ -72,18 +71,22 @@ def notify(
         exception=None,
         error_class=None,
         error_message=None,
-        context={},
+        context: Optional[Dict[str, Any]] = None,
         fingerprint=None,
-        tags=[],
+        tags: Optional[List[str]] = None,
     ):
+        base = error_context.get()
+        tag_ctx = base.pop("_tags", [])
+        merged_ctx = {**base, **(context or {})}
+        merged_tags = list({*tag_ctx, *(tags or [])})
+
         notice = Notice(
             exception=exception,
             error_class=error_class,
             error_message=error_message,
-            context=context,
+            context=merged_ctx,
             fingerprint=fingerprint,
-            tags=tags,
-            thread_local=self.thread_local,
+            tags=merged_tags,
             config=self.config,
         )
         return self._send_notice(notice)
@@ -124,7 +127,9 @@ def event(self, event_type=None, data=None, **kwargs):
         if isinstance(payload["ts"], datetime.datetime):
             payload["ts"] = payload["ts"].strftime(self.TS_FORMAT)
 
-        return self.events_worker.push(payload)
+        final_payload = {**self._get_event_context(), **payload}
+
+        return self.events_worker.push(final_payload)
 
     def configure(self, **kwargs):
         self.config.set_config_from_dict(kwargs)
@@ -141,28 +146,37 @@ def auto_discover_plugins(self):
         if self.config.is_aws_lambda_environment:
             default_plugin_manager.register(contrib.AWSLambdaPlugin())
 
-    def set_context(self, ctx=None, **kwargs):
-        # This operation is an update, not a set!
-        if not ctx:
-            ctx = kwargs
-        else:
-            ctx.update(kwargs)
-        self.thread_local.context = self._get_context()
-        self.thread_local.context.update(ctx)
+    # Error context
+    #
+    def _get_context(self):
+        return error_context.get()
+
+    def set_context(self, ctx: Optional[Dict[str, Any]] = None, **kwargs):
+        error_context.update(ctx, **kwargs)
 
     def reset_context(self):
-        self.thread_local.context = {}
+        error_context.clear()
 
     @contextmanager
-    def context(self, **kwargs):
-        original_context = copy.copy(self._get_context())
-        self.set_context(**kwargs)
-        try:
+    def context(self, ctx: Optional[Dict[str, Any]] = None, **kwargs):
+        with error_context.override(ctx, **kwargs):
+            yield
+
+    # Event context
+    #
+    def _get_event_context(self):
+        return event_context.get()
+
+    def set_event_context(self, ctx: Optional[Dict[str, Any]] = None, **kwargs):
+        event_context.update(ctx, **kwargs)
+
+    def reset_event_context(self):
+        event_context.clear()
+
+    @contextmanager
+    def event_context(self, ctx: Optional[Dict[str, Any]] = None, **kwargs):
+        with event_context.override(ctx, **kwargs):
             yield
-        except:
-            raise
-        else:
-            self.thread_local.context = original_context
 
     def _connection(self):
         if self.config.is_dev() and not self.config.force_report_data:

honeybadger/notice.py

@@ -9,14 +9,11 @@ def __init__(self, *args, **kwargs):
         self.error_message = kwargs.get("error_message", None)
         self.exc_traceback = kwargs.get("exc_traceback", None)
         self.fingerprint = kwargs.get("fingerprint", None)
-        self.thread_local = kwargs.get("thread_local", None)
         self.config = kwargs.get("config", None)
         self.context = kwargs.get("context", {})
         self.tags = self._construct_tags(kwargs.get("tags", []))
 
         self._process_exception()
-        self._process_context()
-        self._process_tags()
 
     def _process_exception(self):
         if self.exception is None and self.error_class:
@@ -28,15 +25,6 @@ def _process_exception(self):
         elif self.exception and self.error_message:
             self.context["error_message"] = self.error_message
 
-    def _process_context(self):
-        self.context = dict(**self._get_thread_context(), **self.context)
-
-    def _process_tags(self):
-        tags_from_context = self._construct_tags(
-            self._get_thread_context().get("_tags", [])
-        )
-        self.tags = list(set(tags_from_context + self.tags))
-
     @cached_property
     def payload(self):
         return create_payload(
@@ -61,15 +49,24 @@ def excluded_exception(self):
                 return True
         return False
 
-    def _get_thread_context(self):
-        if self.thread_local is None:
-            return {}
-        return getattr(self.thread_local, "context", {})
-
     def _construct_tags(self, tags):
-        constructed_tags = []
+        """
+        Accepts either:
+          - a single string (possibly comma-separated)
+          - a list of strings (each possibly comma-separated)
+        and returns a flat list of stripped tags.
+        """
+        raw = []
         if isinstance(tags, str):
-            constructed_tags = [tag.strip() for tag in tags.split(",")]
-        elif isinstance(tags, list):
-            constructed_tags = tags
-        return constructed_tags
+            raw = [tags]
+        elif isinstance(tags, (list, tuple)):
+            raw = tags
+        out = []
+        for item in raw:
+            if not isinstance(item, str):
+                continue
+            for part in item.split(","):
+                t = part.strip()
+                if t:
+                    out.append(t)
+        return out

honeybadger/protocols.py

@@ -1,9 +1,9 @@
-from typing import Protocol, Dict, Any, Optional, List
+from typing import Protocol, Any, Optional, List
 from .types import EventsSendResult, Notice, Event
 
 
 class Connection(Protocol):
-    def send_notice(self, config: Any, payload: Notice) -> Optional[str]:
+    def send_notice(self, config: Any, notice: Notice) -> Optional[str]:
         """
         Send an error notice to Honeybadger.