propagate event drainer errors to prevent infinite loop and require m…

…anual restart Co-authored-by: Linus Phan <13613724+linusphan@users.noreply.github.com> Co-authored-by: Jack <57678801+mothershipper@users.noreply.github.com>
celery · linusphan · Oct 21, 2024 · Oct 21, 2024 · Oct 21, 2024 · Oct 21, 2024
commit 43d8395be712e5f4f929259a4a59f108117f1624
diff --git a/celery/backends/asynchronous.py b/celery/backends/asynchronous.py
@@ -3,8 +3,10 @@
 import threading
 import time
 from collections import deque
+from collections.abc import Callable
 from queue import Empty
 from time import sleep
+from typing import Optional
 from weakref import WeakKeyDictionary
 
 from kombu.utils.compat import detect_environment
@@ -13,6 +15,10 @@
 from celery.exceptions import TimeoutError
 from celery.utils.threads import THREAD_TIMEOUT_MAX
 
+E_CELERY_RESTART_REQUIRED = """
+Celery must be restarted due to an error encountered while draining events.
+"""
+
 __all__ = (
     'AsyncBackendMixin', 'BaseResultConsumer', 'Drainer',
     'register_drainer',
@@ -64,7 +70,8 @@ def wait_for(self, p, wait, timeout=None):
 
 
 class greenletDrainer(Drainer):
-    spawn = None
+    spawn: Callable[[Callable[[], None]], object]
+    _exc: Optional[Exception] = None
     _g = None
     _drain_complete_event = None    # event, sended (and recreated) after every drain_events iteration
 
@@ -78,23 +85,37 @@ def _send_drain_complete_event(self):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self._mutex = threading.Lock()
         self._started = threading.Event()
         self._stopped = threading.Event()
         self._shutdown = threading.Event()
         self._create_drain_complete_event()
 
     def run(self):
         self._started.set()
-        while not self._stopped.is_set():
-            try:
-                self.result_consumer.drain_events(timeout=1)
-                self._send_drain_complete_event()
-                self._create_drain_complete_event()
-            except socket.timeout:
-                pass
-        self._shutdown.set()
+
+        try:
+            while not self._stopped.is_set():
+                try:
+                    self.result_consumer.drain_events(timeout=1)
+                    self._send_drain_complete_event()
+                    self._create_drain_complete_event()
+                except socket.timeout:
+                    pass
+        except Exception as e:
+            with self._mutex:
+                self._exc = e
+        finally:
+            self._shutdown.set()
 
     def start(self):
+        if self._shutdown.is_set():
+            with self._mutex:
+                if self._exc is not None:
+                    raise self._exc
+                else:
+                    raise Exception(E_CELERY_RESTART_REQUIRED)
+
         if not self._started.is_set():
             self._g = self.spawn(self.run)
             self._started.wait()
@@ -109,6 +130,10 @@ def wait_for(self, p, wait, timeout=None):
         if not p.ready:
             self._drain_complete_event.wait(timeout=timeout)
 
+            with self._mutex:
+                if self._exc is not None:
+                    raise self._exc
+
 
 @register_drainer('eventlet')
 class eventletDrainer(greenletDrainer):

diff --git a/celery/backends/redis.py b/celery/backends/redis.py
@@ -128,9 +128,9 @@ def reconnect_on_error(self):
         except self._connection_errors:
             try:
                 self._ensure(self._reconnect_pubsub, ())
-            except self._connection_errors:
+            except self._connection_errors as e:
                 logger.critical(E_RETRY_LIMIT_EXCEEDED)
-                raise
+                raise RuntimeError(E_RETRY_LIMIT_EXCEEDED) from e
 
     def _maybe_cancel_ready_task(self, meta):
         if meta['status'] in states.READY_STATES: