diff --git a/dojo/middleware.py b/dojo/middleware.py index 127ab462a09..8f4967fad33 100644 --- a/dojo/middleware.py +++ b/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") diff --git a/dojo/settings/settings.dist.py b/dojo/settings/settings.dist.py index d1b661e095e..9ac007db054 100644 --- a/dojo/settings/settings.dist.py +++ b/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": { diff --git a/dojo/tasks.py b/dojo/tasks.py index 54c62427b13..56187f4c574 100644 --- a/dojo/tasks.py +++ b/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()) diff --git a/unittests/test_importers_performance.py b/unittests/test_importers_performance.py index 7a226f089a8..bc0b5c645e5 100644 --- a/unittests/test_importers_performance.py +++ b/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): diff --git a/unittests/test_watson_intermediate_flush.py b/unittests/test_watson_intermediate_flush.py index 624b516bf77..a4012e2ac73 100644 --- a/unittests/test_watson_intermediate_flush.py +++ b/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()