Configurable reduction in Celery chatter on Redis queues (#8077)

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
Co-authored-by: John Ewart <john@johnewart.net>

Author: 3 people

changelog/8077-individual-celery-worker-flags.yaml

@@ -0,0 +1,7 @@
+# Copy this file and rename it to <pr_number>-<short-description>.yaml (e.g., 1234-add-user-endpoint.yaml)
+# Fill in the required fields and delete this comment block
+
+type: Changed
+description: Replaced the single worker_disable_gossip_heartbeat_mingle config flag with three independent flags (worker_disable_heartbeat, worker_disable_gossip, worker_disable_mingle) for granular control over Celery worker flags
+pr: 8077
+labels: []

src/fides/api/worker/__init__.py

@@ -22,6 +22,7 @@
     PRIVACY_PREFERENCES_QUEUE_NAME,
     celery_app,
 )
+from fides.config import CONFIG
 
 
 class _PythonAndYamlFilter(DefaultFilter):
@@ -37,15 +38,29 @@ def __call__(self, change: Any, path: str) -> bool:
 
 def _run_celery_worker(worker_queues: str) -> None:
     """Run the Celery worker process. Extracted so it can be used as a watchfiles target."""
-    celery_app.worker_main(
-        argv=[
-            "--quiet",  # Disable Celery startup banner
-            "worker",
-            "--loglevel=info",
-            "--concurrency=2",
-            f"--queues={worker_queues}",
-        ]
-    )
+    argv = [
+        "--quiet",  # Disable Celery startup banner
+        "worker",
+        "--loglevel=info",
+        f"--concurrency={CONFIG.celery.worker_concurrency}",
+        f"--queues={worker_queues}",
+    ]
+    without_flags = []
+    if CONFIG.celery.worker_disable_heartbeat:
+        without_flags.append("--without-heartbeat")
+    if CONFIG.celery.worker_disable_gossip:
+        without_flags.append("--without-gossip")
+    if CONFIG.celery.worker_disable_mingle:
+        without_flags.append("--without-mingle")
+    if without_flags:
+        argv += without_flags
+        logger.info(
+            f"Worker started with {' '.join(without_flags)} "
+            f"(FIDES__CELERY__WORKER_DISABLE_HEARTBEAT={CONFIG.celery.worker_disable_heartbeat}, "
+            f"FIDES__CELERY__WORKER_DISABLE_GOSSIP={CONFIG.celery.worker_disable_gossip}, "
+            f"FIDES__CELERY__WORKER_DISABLE_MINGLE={CONFIG.celery.worker_disable_mingle})"
+        )
+    celery_app.worker_main(argv=argv)
 
 
 def start_worker(

src/fides/config/celery_settings.py

@@ -33,6 +33,26 @@ class CelerySettings(FidesSettings):
     healthcheck_ping_timeout: float = Field(
         default=2.0, description="The timeout in seconds for the health check ping"
     )
+    worker_disable_heartbeat: bool = Field(
+        default=False,
+        description="If true, starts the worker with --without-heartbeat. "
+        "Use this as a workaround for Celery BRPOP connection drop issues (celery/celery#7276).",
+    )
+    worker_disable_gossip: bool = Field(
+        default=False,
+        description="If true, starts the worker with --without-gossip. "
+        "Use this as a workaround for Celery BRPOP connection drop issues (celery/celery#7276).",
+    )
+    worker_disable_mingle: bool = Field(
+        default=False,
+        description="If true, starts the worker with --without-mingle. "
+        "Use this as a workaround for Celery BRPOP connection drop issues (celery/celery#7276).",
+    )
+    worker_concurrency: int = Field(
+        default=2,
+        ge=1,
+        description="Number of worker processes/threads passed to `celery worker --concurrency`.",
+    )
     broker_url: Optional[str] = Field(
         default=None,
         description="Celery broker URL. When set, overrides the default. With redis.cluster_enabled, "

tests/ctl/api/test_worker.py

@@ -132,3 +132,113 @@ def test_no_reload_calls_worker_main_directly(self, worker_main_mock: MagicMock)
         start_worker()
 
         worker_main_mock.assert_called_once()
+
+
+class TestWorkerDisableFlags:
+    """Tests for the individual worker_disable_heartbeat/gossip/mingle config flags."""
+
+    @pytest.fixture(autouse=True)
+    def mock_celery_config(self):
+        """Provide a mock CONFIG.celery with the three disable flags."""
+        mock_config = MagicMock()
+        mock_config.celery.worker_disable_heartbeat = False
+        mock_config.celery.worker_disable_gossip = False
+        mock_config.celery.worker_disable_mingle = False
+        mock_config.celery.worker_concurrency = 2
+        with patch("fides.api.worker.CONFIG", mock_config):
+            yield mock_config
+
+    @pytest.fixture
+    def worker_main_mock(self):
+        with patch("fides.api.worker.celery_app.worker_main") as mock:
+            yield mock
+
+    def test_no_without_flags_when_all_disabled_false(
+        self, mock_celery_config, worker_main_mock: MagicMock
+    ):
+        """Default: no --without-* flags when all flags are False."""
+        start_worker()
+
+        call_args = worker_main_mock.call_args
+        argv = call_args.kwargs["argv"]
+        assert "--without-heartbeat" not in argv
+        assert "--without-gossip" not in argv
+        assert "--without-mingle" not in argv
+
+    @pytest.mark.parametrize(
+        "flag_name",
+        [
+            "worker_disable_heartbeat",
+            "worker_disable_gossip",
+            "worker_disable_mingle",
+        ],
+    )
+    def test_single_flag_enabled(
+        self, mock_celery_config, worker_main_mock: MagicMock, flag_name: str
+    ):
+        """When a disable flag is True, the matching --without-* appears in argv."""
+        setattr(mock_celery_config.celery, flag_name, True)
+        start_worker()
+
+        call_args = worker_main_mock.call_args
+        argv = call_args.kwargs["argv"]
+        flag_to_arg = {
+            "worker_disable_heartbeat": "--without-heartbeat",
+            "worker_disable_gossip": "--without-gossip",
+            "worker_disable_mingle": "--without-mingle",
+        }
+        assert flag_to_arg[flag_name] in argv
+
+    def test_all_three_flags_enabled(
+        self, mock_celery_config, worker_main_mock: MagicMock
+    ):
+        """All three flags together produces all --without-* flags."""
+        mock_celery_config.celery.worker_disable_heartbeat = True
+        mock_celery_config.celery.worker_disable_gossip = True
+        mock_celery_config.celery.worker_disable_mingle = True
+
+        start_worker()
+
+        call_args = worker_main_mock.call_args
+        argv = call_args.kwargs["argv"]
+        assert "--without-heartbeat" in argv
+        assert "--without-gossip" in argv
+        assert "--without-mingle" in argv
+
+
+class TestWorkerConcurrency:
+    """Tests for CONFIG.celery.worker_concurrency."""
+
+    @pytest.fixture(autouse=True)
+    def mock_celery_config(self):
+        mock_config = MagicMock()
+        mock_config.celery.worker_disable_heartbeat = False
+        mock_config.celery.worker_disable_gossip = False
+        mock_config.celery.worker_disable_mingle = False
+        mock_config.celery.worker_concurrency = 2
+        with patch("fides.api.worker.CONFIG", mock_config):
+            yield mock_config
+
+    @pytest.fixture
+    def worker_main_mock(self):
+        with patch("fides.api.worker.celery_app.worker_main") as mock:
+            yield mock
+
+    def test_default_concurrency_in_argv(
+        self, mock_celery_config: MagicMock, worker_main_mock: MagicMock
+    ):
+        start_worker()
+        argv = worker_main_mock.call_args.kwargs["argv"]
+        assert "--concurrency=2" in argv
+
+    @pytest.mark.parametrize("concurrency", [1, 4, 16])
+    def test_custom_concurrency_in_argv(
+        self,
+        mock_celery_config: MagicMock,
+        worker_main_mock: MagicMock,
+        concurrency: int,
+    ):
+        mock_celery_config.celery.worker_concurrency = concurrency
+        start_worker()
+        argv = worker_main_mock.call_args.kwargs["argv"]
+        assert f"--concurrency={concurrency}" in argv

tests/ops/tasks/test_celery.py

@@ -40,6 +40,7 @@ def test_celery_default_config() -> None:
     assert config.celery.task_always_eager
     assert config.celery.event_queue_prefix == "fides_worker"
     assert config.celery.task_default_queue == "fides"
+    assert config.celery.worker_concurrency == 2
 
     celery_app = _create_celery(config)
     expected_broker = (