Repository: celery/celery

Issue:
Celery worker OOM due to unbounded prefetch on eta tasks
# Summary
We are encountering Out Of Memory infinite loop issues on celery workers, when our redis queue is suddenly loaded with either lot of failed tasks, or lot of scheduled tasks.

The reason is that a celery worker keeps fetching scheduled tasks that have an `eta` or `countdown` until it satisfies the number of tasks ready-to-run reaches the `worker_prefetch_multiplier` configuration. When the queue contains lot of scheduled tasks (for us, it's thousands of tasks), the worker keeps fetching them, until it is full and goes out of memory. Then those tasks go back to the queue, and get fetched by other workers, which go in their turn out of memory, and so on.

I checked the code and saw that there is currently no limit on prefetching tasks that have `eta`. Here is the line in [strategy.py](https://github.com/celery/celery/blob/v5.5.3/celery/worker/strategy.py#L198)

Where in kombu, we would increment the prefetch count immediately. In the [QoS class](https://github.com/celery/kombu/blob/main/kombu/common.py#L404).

I was expecting to have a configuration property that can control the limit, but currently there is no such configuration. Could we please add it ?

<!--
Please fill this template entirely and do not erase parts of it.
We reserve the right to close without a response
bug reports which are incomplete.
-->
# Checklist
<!--
To check an item on the list replace [ ] with [x].
-->
- [x] I have verified that the issue exists against the `main` branch of Celery.
- [ ] This has already been asked to the [discussions forum](https://github.com/celery/celery/discussions) first.
- [x] I have read the relevant section in the
  [contribution guide](https://docs.celeryq.dev/en/main/contributing.html#other-bugs)
  on reporting bugs.
- [ ] I have checked the [issues list](https://github.com/celery/celery/issues?q=is%3Aissue+label%3A%22Issue+Type%3A+Bug+Report%22+-label%3A%22Category%3A+Documentation%22)
  for similar or identical bug reports.
- [ ] I have checked the [pull requests list](https://github.com/celery/celery/pulls?q=is%3Apr+label%3A%22PR+Type%3A+Bugfix%22+-label%3A%22Category%3A+Documentation%22)
  for existing proposed fixes.
- [x] I have checked the [commit log](https://github.com/celery/celery/commits/main)
  to find out if the bug was already fixed in the main branch.
- [ ] I have included all related issues and possible duplicate issues
  in this issue (If there are none, check this box anyway).
- [ ] I have tried to reproduce the issue with [pytest-celery](https://docs.celeryq.dev/projects/pytest-celery/en/latest/userguide/celery-bug-report.html) and added the reproduction script below.

## Mandatory Debugging Information

- [ ] I have included the output of ``celery -A proj report`` in the issue.
    (if you are not able to do this, then at least specify the Celery
     version affected).
- [ ] I have verified that the issue exists against the `main` branch of Celery.
- [ ] I have included the contents of ``pip freeze`` in the issue.
- [ ] I have included all the versions of all the external dependencies required
  to reproduce this bug.

## Optional Debugging Information
<!--
Try some of the below if you think they are relevant.
It will help us figure out the scope of the bug and how many users it affects.
-->
- [ ] I have tried reproducing the issue on more than one Python version
  and/or implementation.
- [ ] I have tried reproducing the issue on more than one message broker and/or
  result backend.
- [ ] I have tried reproducing the issue on more than one version of the message
  broker and/or result backend.
- [ ] I have tried reproducing the issue on more than one operating system.
- [ ] I have tried reproducing the issue on more than one workers pool.
- [x] I have tried reproducing the issue with autoscaling, retries,
  ETA/Countdown & rate limits disabled.
- [ ] I have tried reproducing the issue after downgrading
  and/or upgrading Celery and its dependencies.



Discussion:

Patch:
diff --git a/celery/app/defaults.py b/celery/app/defaults.py
index bd44d8bfbb..1e577faa82 100644
--- a/celery/app/defaults.py
+++ b/celery/app/defaults.py
@@ -342,6 +342,7 @@ def __repr__(self):
         pool_restarts=Option(False, type='bool'),
         proc_alive_timeout=Option(4.0, type='float'),
         prefetch_multiplier=Option(4, type='int'),
+        eta_task_limit=Option(None, type='int'),
         enable_prefetch_count_reduction=Option(True, type='bool'),
         redirect_stdouts=Option(
             True, type='bool', old={'celery_redirect_stdouts'},
diff --git a/celery/worker/consumer/tasks.py b/celery/worker/consumer/tasks.py
index 67cbfc1207..95b7ec0eb0 100644
--- a/celery/worker/consumer/tasks.py
+++ b/celery/worker/consumer/tasks.py
@@ -46,7 +46,10 @@ def set_prefetch_count(prefetch_count):
                 prefetch_count=prefetch_count,
                 apply_global=qos_global,
             )
-        c.qos = QoS(set_prefetch_count, c.initial_prefetch_count)
+        eta_task_limit = c.app.conf.worker_eta_task_limit
+        c.qos = QoS(
+            set_prefetch_count, c.initial_prefetch_count, max_prefetch=eta_task_limit
+        )
 
     def stop(self, c):
         """Stop task consumer."""
@@ -79,7 +82,9 @@ def qos_global(self, c) -> bool:
         qos_global = not c.connection.qos_semantics_matches_spec
 
         if c.app.conf.worker_detect_quorum_queues:
-            using_quorum_queues, qname = detect_quorum_queues(c.app, c.connection.transport.driver_type)
+            using_quorum_queues, _ = detect_quorum_queues(
+                c.app, c.connection.transport.driver_type
+            )
 
             if using_quorum_queues:
                 qos_global = False
diff --git a/celery/worker/strategy.py b/celery/worker/strategy.py
index 3fe5fa145c..6a1c6225b4 100644
--- a/celery/worker/strategy.py
+++ b/celery/worker/strategy.py
@@ -109,7 +109,6 @@ def default(task, app, consumer,
     hostname = consumer.hostname
     connection_errors = consumer.connection_errors
     _does_info = logger.isEnabledFor(logging.INFO)
-
     # task event related
     # (optimized to avoid calling request.send_event)
     eventer = consumer.event_dispatcher
@@ -125,7 +124,8 @@ def default(task, app, consumer,
     limit_task = consumer._limit_task
     limit_post_eta = consumer._limit_post_eta
     Request = symbol_by_name(task.Request)
-    Req = create_request_cls(Request, task, consumer.pool, hostname, eventer, app=app)
+    Req = create_request_cls(Request, task, consumer.pool, hostname, eventer,
+                             app=app)
 
     revoked_tasks = consumer.controller.state.revoked
 
@@ -194,6 +194,7 @@ def task_message_handler(message, body, ack, reject, callbacks,
             consumer.qos.increment_eventually()
             return call_at(eta, limit_post_eta, (req, bucket, 1),
                            priority=6)
+
         if eta:
             consumer.qos.increment_eventually()
             call_at(eta, apply_eta_task, (req,), priority=6)
diff --git a/docs/userguide/configuration.rst b/docs/userguide/configuration.rst
index 0de8eba8a5..46fd4599f3 100644
--- a/docs/userguide/configuration.rst
+++ b/docs/userguide/configuration.rst
@@ -174,6 +174,7 @@ have been moved into a new  ``task_`` prefix.
 ``CELERYD_POOL_PUTLOCKS``                  :setting:`worker_pool_putlocks`
 ``CELERYD_POOL_RESTARTS``                  :setting:`worker_pool_restarts`
 ``CELERYD_PREFETCH_MULTIPLIER``            :setting:`worker_prefetch_multiplier`
+``CELERYD_ETA_TASK_LIMIT``                 :setting:`worker_eta_task_limit`
 ``CELERYD_ENABLE_PREFETCH_COUNT_REDUCTION``:setting:`worker_enable_prefetch_count_reduction`
 ``CELERYD_REDIRECT_STDOUTS``               :setting:`worker_redirect_stdouts`
 ``CELERYD_REDIRECT_STDOUTS_LEVEL``         :setting:`worker_redirect_stdouts_level`
@@ -3186,6 +3187,24 @@ as many messages as it wants.
 
 For more on prefetching, read :ref:`optimizing-prefetch-limit`
 
+.. setting:: worker_eta_task_limit
+
+``worker_eta_task_limit``
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 5.6
+
+Default: No limit (None).
+
+The maximum number of ETA/countdown tasks that a worker can hold in memory at once.
+When this limit is reached, the worker will not receive new tasks from the broker
+until some of the existing ETA tasks are executed.
+
+This setting helps prevent memory exhaustion when a queue contains a large number
+of tasks with ETA/countdown values, as these tasks are held in memory until their
+execution time. Without this limit, workers may fetch thousands of ETA tasks into
+memory, potentially causing out-of-memory issues.
+
 .. note::
 
     Tasks with ETA/countdown aren't affected by prefetch limits.


Test patch:
diff --git a/t/unit/worker/test_consumer.py b/t/unit/worker/test_consumer.py
index 04d167e3d8..1805e58154 100644
--- a/t/unit/worker/test_consumer.py
+++ b/t/unit/worker/test_consumer.py
@@ -706,6 +706,72 @@ def test_log_when_qos_is_false(self, caplog):
         assert record.levelname == "INFO"
         assert record.msg == "Global QoS is disabled. Prefetch count in now static."
 
+    def test_qos_with_worker_eta_task_limit(self):
+        """Test QoS is instantiated with worker_eta_task_limit as max_prefetch."""
+        c = self.c
+        c.app.conf.worker_eta_task_limit = 100
+        c.initial_prefetch_count = 10
+        c.task_consumer = Mock()
+        c.app.amqp.TaskConsumer = Mock(return_value=c.task_consumer)
+        c.connection.default_channel.basic_qos = Mock()
+        c.update_strategies = Mock()
+        c.on_decode_error = Mock()
+
+        tasks = Tasks(c)
+
+        with patch('celery.worker.consumer.tasks.QoS') as mock_qos:
+            tasks.start(c)
+
+            # Verify QoS was called with max_prefetch set to worker_eta_task_limit
+            mock_qos.assert_called_once()
+            args, kwargs = mock_qos.call_args
+            assert len(args) == 2  # callback and initial_value
+            assert kwargs.get('max_prefetch') == 100
+
+    def test_qos_without_worker_eta_task_limit(self):
+        """Test QoS is instantiated with None max_prefetch when worker_eta_task_limit is None."""
+        c = self.c
+        c.app.conf.worker_eta_task_limit = None
+        c.initial_prefetch_count = 10
+        c.task_consumer = Mock()
+        c.app.amqp.TaskConsumer = Mock(return_value=c.task_consumer)
+        c.connection.default_channel.basic_qos = Mock()
+        c.update_strategies = Mock()
+        c.on_decode_error = Mock()
+
+        tasks = Tasks(c)
+
+        with patch('celery.worker.consumer.tasks.QoS') as mock_qos:
+            tasks.start(c)
+
+            # Verify QoS was called with max_prefetch set to None
+            mock_qos.assert_called_once()
+            args, kwargs = mock_qos.call_args
+            assert len(args) == 2  # callback and initial_value
+            assert kwargs.get('max_prefetch') is None
+
+    def test_qos_with_zero_worker_eta_task_limit(self):
+        """Test that QoS respects zero as a valid worker_eta_task_limit value."""
+        c = self.c
+        c.app.conf.worker_eta_task_limit = 0
+        c.initial_prefetch_count = 10
+        c.task_consumer = Mock()
+        c.app.amqp.TaskConsumer = Mock(return_value=c.task_consumer)
+        c.connection.default_channel.basic_qos = Mock()
+        c.update_strategies = Mock()
+        c.on_decode_error = Mock()
+
+        tasks = Tasks(c)
+
+        with patch('celery.worker.consumer.tasks.QoS') as mock_qos:
+            tasks.start(c)
+
+            # Verify QoS was called with max_prefetch set to 0
+            mock_qos.assert_called_once()
+            args, kwargs = mock_qos.call_args
+            assert len(args) == 2  # callback and initial_value
+            assert kwargs.get('max_prefetch') == 0
+
 
 class test_Agent:
 
diff --git a/t/unit/worker/test_strategy.py b/t/unit/worker/test_strategy.py
index 30c50b9845..b2b829c4f4 100644
--- a/t/unit/worker/test_strategy.py
+++ b/t/unit/worker/test_strategy.py
@@ -1,5 +1,4 @@
 import logging
-from collections import defaultdict
 from contextlib import contextmanager
 from unittest.mock import ANY, Mock, patch
 
@@ -99,8 +98,8 @@ def was_limited_with_eta(self):
             assert not self.was_reserved()
             called = self.consumer.timer.call_at.called
             if called:
-                assert self.consumer.timer.call_at.call_args[0][1] == \
-                    self.consumer._limit_post_eta
+                callback = self.consumer.timer.call_at.call_args[0][1]
+                assert callback == self.consumer._limit_post_eta
             return called
 
         def was_scheduled(self):
@@ -128,10 +127,15 @@ def _context(self, sig,
 
         reserved = Mock()
         consumer = Mock()
-        consumer.task_buckets = defaultdict(lambda: None)
+        # Create a proper mock for task_buckets that supports __getitem__
+        task_buckets_mock = Mock()
+        task_buckets_mock.__getitem__ = Mock(side_effect=lambda key: None)
+        consumer.task_buckets = task_buckets_mock
         if limit:
             bucket = TokenBucket(rate(limit), capacity=1)
-            consumer.task_buckets[sig.task] = bucket
+            task_buckets_mock.__getitem__.side_effect = (
+                lambda key: bucket if key == sig.task else None
+            )
         consumer.controller.state.revoked = set()
         consumer.disable_rate_limits = not rate_limits
         consumer.event_dispatcher.enabled = events