Adding ability to retry signal receiver after raised exception (celery#4192)

Author: daviddavis

CONTRIBUTORS.txt

@@ -243,5 +243,6 @@ Samuel Dion-Girardeau, 2017/05/29
 Aydin Sen, 2017/06/14
 Preston Moore, 2017/06/18
 Nicolas Mota, 2017/08/10
+David Davis, 2017/08/11
 Martial Pageau, 2017/08/16
 Sammie S. Taunton, 2017/08/17

celery/utils/dispatch/signal.py

@@ -5,10 +5,12 @@
 import threading
 import weakref
 import warnings
+from kombu.utils.functional import retry_over_time
 from celery.exceptions import CDeprecationWarning
 from celery.five import python_2_unicode_compatible, range, text_t
 from celery.local import PromiseProxy, Proxy
 from celery.utils.functional import fun_accepts_kwargs
+from celery.utils.time import humanize_seconds
 from celery.utils.log import get_logger
 try:
     from weakref import WeakMethod
@@ -36,6 +38,10 @@ def _make_id(target):  # pragma: no cover
 
 NO_RECEIVERS = object()
 
+RECEIVER_RETRY_ERROR = """\
+Could not process signal receiver %(receiver)s. Retrying %(when)s...\
+"""
+
 
 @python_2_unicode_compatible
 class Signal(object):  # pragma: no cover
@@ -103,12 +109,49 @@ def connect(self, *args, **kwargs):
             dispatch_uid (Hashable): An identifier used to uniquely identify a
                 particular instance of a receiver.  This will usually be a
                 string, though it may be anything hashable.
+
+            retry (bool): If the signal receiver raises an exception
+                (e.g. ConnectionError), the receiver will be retried until it
+                runs successfully. A strong ref to the receiver will be stored
+                and the `weak` option will be ignored.
         """
-        def _handle_options(sender=None, weak=True, dispatch_uid=None):
+        def _handle_options(sender=None, weak=True, dispatch_uid=None,
+                            retry=False):
 
             def _connect_signal(fun):
-                self._connect_signal(fun, sender, weak, dispatch_uid)
+
+                options = {'dispatch_uid': dispatch_uid,
+                           'weak': weak}
+
+                def _retry_receiver(retry_fun):
+
+                    def _try_receiver_over_time(*args, **kwargs):
+                        def on_error(exc, intervals, retries):
+                            interval = next(intervals)
+                            err_msg = RECEIVER_RETRY_ERROR % \
+                                {'receiver': retry_fun,
+                                 'when': humanize_seconds(interval, 'in', ' ')}
+                            logger.error(err_msg)
+                            return interval
+
+                        return retry_over_time(retry_fun, Exception, args,
+                                               kwargs, on_error)
+
+                    return _try_receiver_over_time
+
+                if retry:
+                    options['weak'] = False
+                    if not dispatch_uid:
+                        # if there's no dispatch_uid then we need to set the
+                        # dispatch uid to the original func id so we can look
+                        # it up later with the original func id
+                        options['dispatch_uid'] = _make_id(fun)
+                    fun = _retry_receiver(fun)
+
+                self._connect_signal(fun, sender, options['weak'],
+                                     options['dispatch_uid'])
                 return fun
+
             return _connect_signal
 
         if args and callable(args[0]):
@@ -158,6 +201,7 @@ def _connect_signal(self, receiver, sender, weak, dispatch_uid):
             else:
                 self.receivers.append((lookup_key, receiver))
             self.sender_receivers_cache.clear()
+
         return receiver
 
     def disconnect(self, receiver=None, sender=None, weak=None,

t/unit/utils/test_dispatcher.py

@@ -143,3 +143,32 @@ def test_disconnection(self):
         finally:
             a_signal.disconnect(receiver_3)
         self._testIsClean(a_signal)
+
+    def test_retry(self):
+
+        class non_local:
+            counter = 1
+
+        def succeeds_eventually(val, **kwargs):
+            non_local.counter += 1
+            if non_local.counter < 3:
+                raise ValueError('this')
+
+            return val
+
+        a_signal.connect(succeeds_eventually, sender=self, retry=True)
+        try:
+            result = a_signal.send(sender=self, val='test')
+            assert non_local.counter == 3
+            assert result[0][1] == 'test'
+        finally:
+            a_signal.disconnect(succeeds_eventually, sender=self)
+        self._testIsClean(a_signal)
+
+    def test_retry_with_dispatch_uid(self):
+        uid = 'abc123'
+        a_signal.connect(receiver_1_arg, sender=self, retry=True,
+                         dispatch_uid=uid)
+        assert a_signal.receivers[0][0][0] == uid
+        a_signal.disconnect(receiver_1_arg, sender=self, dispatch_uid=uid)
+        self._testIsClean(a_signal)