fix(watson): fix + batch full-text index updates on Celery worker tasks

dojo/middleware.py

@@ -1,7 +1,7 @@
 import logging
 import re
 import time
-from contextlib import suppress
+from contextlib import contextmanager, suppress
 from threading import local
 from urllib.parse import quote
 
@@ -291,35 +291,75 @@ def _drain_search_context_to_async(objects, source):
         objects.discard(entry)
 
 
+@contextmanager
+def watson_search_context_for_task():
+    """
+    Batch watson index updates for saves that happen inside a Celery task.
+
+    Celery workers serve no HTTP request, so ``AsyncSearchContextMiddleware`` never runs
+    and no watson ``search_context`` is active while the task executes. Without an active
+    context, django-watson's ``post_save`` receiver indexes every saved object
+    synchronously — one DELETE + INSERT into ``watson_searchentry`` per object. A bulk
+    import running in a worker (Pro's ``AsyncImporter``) then re-indexes findings one at a
+    time, flooding the DB with thousands of index writes.
+
+    Opening a search_context around the task makes those saves accumulate and drain in
+    ``WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE``-sized async batches on exit, exactly like the
+    request path does via ``AsyncSearchContextMiddleware``. Wire it in through
+    ``CELERY_TASK_CONTEXT_MANAGERS`` so ``PluggableContextTask`` enters it around every
+    task. An empty context (a task that saved no indexed models) drains to nothing, so this
+    is a cheap no-op for tasks that do not touch registered models.
+    """
+    search_context_manager.start()
+    try:
+        yield
+    finally:
+        # Drain accumulated objects to async batched index-update tasks, then end the
+        # (now-empty) context so watson's bulk-save short-circuits — mirrors
+        # AsyncSearchContextMiddleware._close_search_context for the request path.
+        if search_context_manager.is_active():
+            objects, _is_invalid = search_context_manager._stack[-1]
+            _drain_search_context_to_async(objects, source="watson_search_context_for_task")
+        search_context_manager.end()
+
+
 def install_intermediate_flush_hook():
     """
-    Wrap `watson.search.search_context_manager.add_to_context` with a
-    size-based flush. Once the per-request set reaches
+    Wrap `add_to_context` on the module-global `watson.search.search_context_manager`
+    with a size-based flush. Once the shared accumulation set reaches
     `WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE`, drain it into async tasks
     and clear it in place. Bounds memory on long-running requests
     (large imports) and starts celery batches earlier instead of
     dispatching all at end-of-request.
 
+    The wrapper is bound to the singleton INSTANCE, not the class: only the shared
+    request/task accumulation context batches. A throwaway local SearchContextManager
+    — e.g. the one update_watson_search_index_for_model builds to index one
+    already-bounded batch — keeps the stock method and indexes its own batch on
+    end(). If local contexts flushed too, that index task would re-dispatch itself
+    for the same batch (infinite recursion under eager celery / re-dispatch loop on
+    a worker). watson's post_save always adds to the module-global instance, so the
+    singleton is the only place the flush is needed.
+
     Idempotent — safe to call multiple times.
     Setting WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE to 0 or below disables
     the hook at runtime.
     """
-    cls = search_context_manager.__class__
-    if getattr(cls, "_dd_intermediate_flush_installed", False):
+    if getattr(search_context_manager, "_dd_intermediate_flush_installed", False):
         return
 
-    original_add = cls.add_to_context
+    original_add = search_context_manager.add_to_context  # bound method
 
-    def add_to_context_with_flush(self, engine, obj):
-        original_add(self, engine, obj)
+    def add_to_context_with_flush(engine, obj):
+        original_add(engine, obj)
         threshold = getattr(settings, "WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE", 1000)
-        if threshold <= 0 or not self._stack:
+        if threshold <= 0 or not search_context_manager._stack:
             return
-        objects, is_invalid = self._stack[-1]
+        objects, is_invalid = search_context_manager._stack[-1]
         if is_invalid or len(objects) < threshold:
             return
         _drain_search_context_to_async(objects, source="AsyncSearchContextMiddleware[intermediate]")
 
-    cls.add_to_context = add_to_context_with_flush
-    cls._dd_intermediate_flush_installed = True
-    logger.debug("AsyncSearchContextMiddleware: intermediate flush hook installed on %s", cls.__name__)
+    search_context_manager.add_to_context = add_to_context_with_flush
+    search_context_manager._dd_intermediate_flush_installed = True
+    logger.debug("AsyncSearchContextMiddleware: intermediate flush hook installed on the global search_context_manager")

dojo/settings/settings.dist.py

@@ -902,6 +902,16 @@ def generate_url(scheme, double_slashes, user, password, host, port, path, param
 WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE = env("DD_WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE")
 WATSON_INDEX_PREFETCH_ENABLED = env("DD_WATSON_INDEX_PREFETCH_ENABLED")
 
+# Context managers wrapped around every Celery task by PluggableContextTask (see
+# dojo/celery.py). watson_search_context_for_task opens a watson search_context so bulk
+# finding saves that happen inside a worker task (which serves no HTTP request, so
+# AsyncSearchContextMiddleware never runs) accumulate and drain in batched async index
+# updates instead of indexing one finding at a time. Extend this list downstream (e.g. Pro)
+# rather than replacing it, so this batching stays wired.
+CELERY_TASK_CONTEXT_MANAGERS = [
+    "dojo.middleware.watson_search_context_for_task",
+]
+
 # Celery beat scheduled tasks
 CELERY_BEAT_SCHEDULE = {
     "add-alerts": {

dojo/tasks.py

@@ -204,6 +204,11 @@ def update_watson_search_index_for_model(model_name, pk_list, *args, **kwargs):
         instances_added = 0
         instances_skipped = 0
 
+        # This task IS the terminal drain: it accumulates one already-bounded batch (<= 1000
+        # PKs) into its own local SearchContextManager and bulk-saves it once via end(). The
+        # intermediate size-flush hook is bound to the global singleton instance only, so
+        # this private context keeps the stock add_to_context, never re-triggers the flush,
+        # and cannot re-dispatch itself.
         for instance in instances:
             try:
                 # Add to watson context (this will trigger indexing on end())

unittests/test_importers_performance.py

@@ -392,14 +392,14 @@ def test_import_reimport_reimport_performance_pghistory_no_async_with_product_gr
         self.system_settings(enable_product_grade=True)
 
         self._import_reimport_performance(
-            expected_num_queries1=183,
-            expected_num_async_tasks1=4,
-            expected_num_queries2=140,
-            expected_num_async_tasks2=3,
+            expected_num_queries1=181,
+            expected_num_async_tasks1=5,
+            expected_num_queries2=138,
+            expected_num_async_tasks2=4,
             expected_num_queries3=44,
             expected_num_async_tasks3=3,
-            expected_num_queries4=109,
-            expected_num_async_tasks4=2,
+            expected_num_queries4=107,
+            expected_num_async_tasks4=3,
         )
 
     # Deduplication is enabled in the tests above, but to properly test it we must run the same import twice and capture the results.
@@ -719,14 +719,14 @@ def test_import_reimport_reimport_performance_pghistory_no_async_with_product_gr
         self.system_settings(enable_product_grade=True)
 
         self._import_reimport_performance(
-            expected_num_queries1=195,
-            expected_num_async_tasks1=4,
-            expected_num_queries2=154,
-            expected_num_async_tasks2=3,
+            expected_num_queries1=193,
+            expected_num_async_tasks1=5,
+            expected_num_queries2=152,
+            expected_num_async_tasks2=4,
             expected_num_queries3=54,
             expected_num_async_tasks3=3,
-            expected_num_queries4=113,
-            expected_num_async_tasks4=2,
+            expected_num_queries4=111,
+            expected_num_async_tasks4=3,
         )
 
     def _deduplication_performance(self, expected_num_queries1, expected_num_async_tasks1, expected_num_queries2, expected_num_async_tasks2, *, check_duplicates=True):

unittests/test_watson_intermediate_flush.py

@@ -10,7 +10,7 @@
 from unittest.mock import patch
 
 from django.test import override_settings
-from watson.search import search_context_manager
+from watson.search import SearchContextManager, search_context_manager
 
 from dojo.middleware import (
     _drain_search_context_to_async,  # noqa: PLC2701 -- internal helper under test
@@ -119,3 +119,22 @@ def test_invalidated_context_skips_drain(self):
                 # hook should detect the invalid flag and bail out.
                 search_context_manager.add_to_context(engine_marker, p)
             drain.assert_not_called()
+
+    @override_settings(WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE=2)
+    def test_ad_hoc_context_manager_does_not_drain(self):
+        """
+        The flush wrapper is bound to the global singleton instance only. An ad-hoc
+        SearchContextManager -- e.g. the one update_watson_search_index_for_model builds to
+        index its own batch -- keeps the stock add_to_context and must NOT drain, or it would
+        dispatch a clone of itself and loop forever (queue ~0, worker pegged, nothing indexed).
+        """
+        adhoc = SearchContextManager()
+        adhoc.start()
+        try:
+            with patch("dojo.middleware._drain_search_context_to_async") as drain:
+                for p in self.products[:3]:  # past the threshold of 2
+                    adhoc.add_to_context(object(), p)
+                drain.assert_not_called()
+        finally:
+            adhoc.invalidate()
+            adhoc.end()