refactor: share task cache instance between celery and procrastinate

Author: snopoke

taskbadger/_integrations.py

@@ -1,10 +1,10 @@
 """Shared internals for taskbadger's optional system integrations
 (Celery, Procrastinate). Not part of the public API.
 
-Each integration creates its own module-level ``TaskCache`` instance and
-defines a thin ``safe_get_task`` wrapper around the shared one defined here.
-``BaseSystemIntegration`` provides the common ctor/include-exclude shape;
-subclasses override ``track_task`` if they need to filter additional
+A single module-level ``TaskCache`` (``task_cache``) is shared across all
+integrations; task ids are UUIDs so cross-integration key collisions are not
+a concern. ``BaseSystemIntegration`` provides the common ctor/include-exclude
+shape; subclasses override ``track_task`` if they need to filter additional
 task names (e.g. Procrastinate built-ins).
 """
 
@@ -51,20 +51,23 @@ def unset(self, key) -> None:
         self.cache.pop(key, None)
 
 
-def safe_get_task(cache: TaskCache, task_id: str):
+task_cache = TaskCache()
+
+
+def safe_get_task(task_id: str):
     """Cache-aware ``get_task``: returns the cached entry if present, otherwise
-    fetches via the SDK ``get_task`` and caches the result. Errors are logged and
-    swallowed (returns ``None``). ``None`` results are not cached.
+    fetches from the API and caches the result. Errors are logged and swallowed
+    (returns ``None``). ``None`` results are not cached.
     """
-    cached = cache.get(task_id)
+    cached = task_cache.get(task_id)
     if cached is not None:
         return cached
     try:
         task = sdk.get_task(task_id)
     except Exception as e:
         log.warning("Error fetching task '%s': %s", task_id, e)
         return None
-    cache.set(task_id, task)
+    task_cache.set(task_id, task)
     return task
 
 

taskbadger/celery.py

@@ -13,8 +13,7 @@
 from kombu import serialization
 
 from . import sdk
-from ._integrations import TERMINAL_STATES, TaskCache
-from ._integrations import safe_get_task as _shared_safe_get_task
+from ._integrations import TERMINAL_STATES, safe_get_task, task_cache
 from .internal.models import StatusEnum
 from .mug import Badger
 from .safe_sdk import create_task_safe, update_task_safe
@@ -27,8 +26,6 @@
 
 log = logging.getLogger("taskbadger")
 
-_task_cache = TaskCache()
-
 
 class Task(celery.Task):
     """A Celery Task that tracks itself with TaskBadger.
@@ -249,7 +246,7 @@ def _maybe_create_task(signal_sender):
     if task:
         # Store the task ID in the request so _update_task can find it
         signal_sender.request.update({TB_TASK_ID: task.id})
-        _task_cache.set(task.id, task)
+        task_cache.set(task.id, task)
 
 
 @task_prerun.connect
@@ -301,7 +298,7 @@ def _update_task(signal_sender, status, einfo=None):
         data = DefaultMergeStrategy().merge(task.data, {"exception": str(einfo)})
     task = update_task_safe(task.id, status=status, data=data)
     if task:
-        _task_cache.set(task_id, task)
+        task_cache.set(task_id, task)
 
 
 def enter_session():
@@ -321,17 +318,13 @@ def exit_session(signal_sender):
     if not task_id or not Badger.is_configured():
         return
 
-    _task_cache.unset(task_id)
+    task_cache.unset(task_id)
 
     session = Badger.current.session()
     if session.client:
         session.__exit__()
 
 
-def safe_get_task(task_id: str):
-    return _shared_safe_get_task(_task_cache, task_id)
-
-
 def _get_taskbadger_task_id(request):
     if not request:
         return

taskbadger/procrastinate.py

@@ -17,8 +17,7 @@
 import logging
 from contextvars import ContextVar
 
-from ._integrations import TERMINAL_STATES, TaskCache
-from ._integrations import safe_get_task as _shared_safe_get_task
+from ._integrations import TERMINAL_STATES, safe_get_task, task_cache
 from .internal.models import StatusEnum
 from .mug import Badger
 from .safe_sdk import create_task_safe, update_task_safe
@@ -118,8 +117,8 @@ def _update_status(tb_id, status, exception=None):
         # Bypass the cache for the terminal-state check: the user may have
         # updated the task to a terminal state via the regular SDK during
         # the body, which wouldn't be reflected in our local cache.
-        _task_cache.unset(tb_id)
-        current = _safe_get_task(tb_id)
+        task_cache.unset(tb_id)
+        current = safe_get_task(tb_id)
         if current is not None and current.status in TERMINAL_STATES:
             return
         data = None
@@ -134,14 +133,7 @@ def _update_status(tb_id, status, exception=None):
         updated = update_task_safe(tb_id, status=status)
 
     if updated is not None:
-        _task_cache.set(tb_id, updated)
-
-
-_task_cache = TaskCache()
-
-
-def _safe_get_task(task_id):
-    return _shared_safe_get_task(_task_cache, task_id)
+        task_cache.set(tb_id, updated)
 
 
 def _wrap_defer(task):
@@ -267,7 +259,7 @@ def current_task():
     tb_id = _current_tb_task_id.get()
     if tb_id is None:
         return None
-    return _safe_get_task(tb_id)
+    return safe_get_task(tb_id)
 
 
 def _patch_app_task(app, system):

tests/conftest.py

@@ -1,8 +1,18 @@
 import pytest
 
+from taskbadger._integrations import task_cache
 from taskbadger.mug import Badger, Settings
 
 
+@pytest.fixture(autouse=True)
+def _clear_task_cache():
+    """Clear the shared integrations task cache around every test so cached
+    entries from earlier tests can't leak into later ones."""
+    task_cache.cache.clear()
+    yield
+    task_cache.cache.clear()
+
+
 @pytest.fixture
 def _bind_settings():
     Badger.current.bind(Settings("https://taskbadger.net", "token", "org", "proj"))

tests/test_procrastinate.py

@@ -7,7 +7,7 @@
 from procrastinate import testing
 
 from taskbadger import StatusEnum
-from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task, _task_cache, current_task, track
+from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task, current_task, track
 from tests.utils import task_for_test
 
 
@@ -19,13 +19,6 @@ def _check_log_errors(caplog):
         pytest.fail(f"log errors during tests: {errors}")
 
 
-@pytest.fixture(autouse=True)
-def _clear_task_cache():
-    _task_cache.cache.clear()
-    yield
-    _task_cache.cache.clear()
-
-
 @pytest.fixture
 def app():
     in_memory = testing.InMemoryConnector()

tests/test_procrastinate_system_integration.py

@@ -4,18 +4,11 @@
 import pytest
 from procrastinate import testing
 
-from taskbadger.procrastinate import _INSTRUMENTED_ATTR, TB_TASK_ID_KWARG, _task_cache, track
+from taskbadger.procrastinate import _INSTRUMENTED_ATTR, TB_TASK_ID_KWARG, track
 from taskbadger.systems.procrastinate import ProcrastinateSystemIntegration
 from tests.utils import task_for_test
 
 
-@pytest.fixture(autouse=True)
-def _clear_task_cache():
-    _task_cache.cache.clear()
-    yield
-    _task_cache.cache.clear()
-
-
 @pytest.fixture
 def app():
     in_memory = testing.InMemoryConnector()