Repository: celery/celery

Issue:
Handle Redis connection errors in result consumer
<!--
Please fill this template entirely and do not erase parts of it.
We reserve the right to close without a response
enhancement requests which are incomplete.
-->
# Checklist
<!--
To check an item on the list replace [ ] with [x].
-->

- [x] I have checked the [issues list](https://github.com/celery/celery/issues?q=is%3Aissue+label%3A%22Issue+Type%3A+Enhancement%22+-label%3A%22Category%3A+Documentation%22)
  for similar or identical enhancement to an existing feature.
- [x] I have checked the [pull requests list](https://github.com/celery/celery/pulls?q=is%3Apr+label%3A%22Issue+Type%3A+Enhancement%22+-label%3A%22Category%3A+Documentation%22)
  for existing proposed enhancements.
- [x] I have checked the [commit log](https://github.com/celery/celery/commits/master)
  to find out if the if the same enhancement was already implemented in the
  master branch.
- [x] I have included all related issues and possible duplicate issues in this issue
      (If there are none, check this box anyway).

## Related Issues and Possible Duplicates
<!--
Please make sure to search and mention any related issues
or possible duplicates to this issue as requested by the checklist above.

This may or may not include issues in other repositories that the Celery project
maintains or other repositories that are dependencies of Celery.

If you don't know how to mention issues, please refer to Github's documentation
on the subject: https://help.github.com/en/articles/autolinked-references-and-urls#issues-and-pull-requests
-->

#### Related Issues

- None

#### Possible Duplicates

- None

# Brief Summary
<!--
Please include a brief summary of what the enhancement is
and why it is needed.
-->
When there is a connection error with Redis while executing a command, in most cases, the redis client will [discard the connection](https://github.com/andymccurdy/redis-py/blob/272d3139e1c82c2d89551f87d12df7c18d938ea2/redis/client.py#L885-L886), causing the next command sent to Redis to open a new connection. This allows applications to recover from connection errors by simply retrying, a property that is used in Celery, for example when setting keys in the Redis result backend: https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L324-L325

This is not the case however when the [connection to Redis is in a pubsub state](https://github.com/andymccurdy/redis-py/blob/272d3139e1c82c2d89551f87d12df7c18d938ea2/redis/client.py#L3397-L3414). The reason for that is that some state associated with the connection (namely the list of keys subscibed to). The Redis client doesn't keep track of this state, so it can't possibly restore it when creating a new connection and leaves the connection handling to the application code.

The Celery Redis result consumer uses pubsub in order to be notified when results are available, but doesn't handle connection errors at all, causing a result consumer to end up in a state where it can't connect to the result backend any more after a single connection error, as any further attempt will reuse the same faulty connection.

The solution would be to add error handling logic to the result consumer, so it will recreate the connection on connection errors and initialize it to the proper state.

# Design

## Architectural Considerations
<!--
If more components other than Celery are involved,
describe them here and the effect it would have on Celery.
-->
None

## Proposed Behavior
<!--
Please describe in detail how this enhancement is going to change the behavior
of an existing feature.
Describe what happens in case of failures as well if applicable.
-->
Add error handling in all places the Redis result consumer sends a Redis command in a pubsub context:
* https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L127
* https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L142
* https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L148

We should catch all Redis connection errors, and call a new method that will reinitialize a pubsub connection in the proper state (discard the current connection from the pool, start the pubsub context, subscribe to all keys in `ResultConsumer.subscribed_to`) using the retry policy. If in `drain_events`, we should try to get new messages again. 

This will take care of most issues with connection errors. I see two remaining issues:
1.Some message might have been lost (sent between losing the connection and reconnecting). We could read all keys subscribed to right after reconnecting and before starting the pubsub context and call `on_state_change` for each existing key, but this might cause some messages to be delivered twice and I don't know how Celery will react to that.
2. If the connection can't be re-established despite the retries and reaches max-retries, the result consumer will end up with a faulty connection that can't be recovered from. This should be communicated somehow to the user (documentation, logging an explicit error message, custom exception).

## Proposed UI/UX
<!--
Please provide your ideas for the API, CLI options,
configuration key names etc. that will be adjusted for this enhancement.
-->
None

## Diagrams
<!--
Please include any diagrams that might be relevant
to the implementation of this enhancement such as:
* Class Diagrams
* Sequence Diagrams
* Activity Diagrams
You can drag and drop images into the text box to attach them to this issue.
-->
N/A

## Alternatives
<!--
If you have considered any alternative implementations
describe them in detail below.
-->
None

Handle Redis connection errors in result consumer
<!--
Please fill this template entirely and do not erase parts of it.
We reserve the right to close without a response
enhancement requests which are incomplete.
-->
# Checklist
<!--
To check an item on the list replace [ ] with [x].
-->

- [x] I have checked the [issues list](https://github.com/celery/celery/issues?q=is%3Aissue+label%3A%22Issue+Type%3A+Enhancement%22+-label%3A%22Category%3A+Documentation%22)
  for similar or identical enhancement to an existing feature.
- [x] I have checked the [pull requests list](https://github.com/celery/celery/pulls?q=is%3Apr+label%3A%22Issue+Type%3A+Enhancement%22+-label%3A%22Category%3A+Documentation%22)
  for existing proposed enhancements.
- [x] I have checked the [commit log](https://github.com/celery/celery/commits/master)
  to find out if the if the same enhancement was already implemented in the
  master branch.
- [x] I have included all related issues and possible duplicate issues in this issue
      (If there are none, check this box anyway).

## Related Issues and Possible Duplicates
<!--
Please make sure to search and mention any related issues
or possible duplicates to this issue as requested by the checklist above.

This may or may not include issues in other repositories that the Celery project
maintains or other repositories that are dependencies of Celery.

If you don't know how to mention issues, please refer to Github's documentation
on the subject: https://help.github.com/en/articles/autolinked-references-and-urls#issues-and-pull-requests
-->

#### Related Issues

- None

#### Possible Duplicates

- None

# Brief Summary
<!--
Please include a brief summary of what the enhancement is
and why it is needed.
-->
When there is a connection error with Redis while executing a command, in most cases, the redis client will [discard the connection](https://github.com/andymccurdy/redis-py/blob/272d3139e1c82c2d89551f87d12df7c18d938ea2/redis/client.py#L885-L886), causing the next command sent to Redis to open a new connection. This allows applications to recover from connection errors by simply retrying, a property that is used in Celery, for example when setting keys in the Redis result backend: https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L324-L325

This is not the case however when the [connection to Redis is in a pubsub state](https://github.com/andymccurdy/redis-py/blob/272d3139e1c82c2d89551f87d12df7c18d938ea2/redis/client.py#L3397-L3414). The reason for that is that some state associated with the connection (namely the list of keys subscibed to). The Redis client doesn't keep track of this state, so it can't possibly restore it when creating a new connection and leaves the connection handling to the application code.

The Celery Redis result consumer uses pubsub in order to be notified when results are available, but doesn't handle connection errors at all, causing a result consumer to end up in a state where it can't connect to the result backend any more after a single connection error, as any further attempt will reuse the same faulty connection.

The solution would be to add error handling logic to the result consumer, so it will recreate the connection on connection errors and initialize it to the proper state.

# Design

## Architectural Considerations
<!--
If more components other than Celery are involved,
describe them here and the effect it would have on Celery.
-->
None

## Proposed Behavior
<!--
Please describe in detail how this enhancement is going to change the behavior
of an existing feature.
Describe what happens in case of failures as well if applicable.
-->
Add error handling in all places the Redis result consumer sends a Redis command in a pubsub context:
* https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L127
* https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L142
* https://github.com/celery/celery/blob/d0563058f8f47f347ac1b56c44f833f569764482/celery/backends/redis.py#L148

We should catch all Redis connection errors, and call a new method that will reinitialize a pubsub connection in the proper state (discard the current connection from the pool, start the pubsub context, subscribe to all keys in `ResultConsumer.subscribed_to`) using the retry policy. If in `drain_events`, we should try to get new messages again. 

This will take care of most issues with connection errors. I see two remaining issues:
1.Some message might have been lost (sent between losing the connection and reconnecting). We could read all keys subscribed to right after reconnecting and before starting the pubsub context and call `on_state_change` for each existing key, but this might cause some messages to be delivered twice and I don't know how Celery will react to that.
2. If the connection can't be re-established despite the retries and reaches max-retries, the result consumer will end up with a faulty connection that can't be recovered from. This should be communicated somehow to the user (documentation, logging an explicit error message, custom exception).

## Proposed UI/UX
<!--
Please provide your ideas for the API, CLI options,
configuration key names etc. that will be adjusted for this enhancement.
-->
None

## Diagrams
<!--
Please include any diagrams that might be relevant
to the implementation of this enhancement such as:
* Class Diagrams
* Sequence Diagrams
* Activity Diagrams
You can drag and drop images into the text box to attach them to this issue.
-->
N/A

## Alternatives
<!--
If you have considered any alternative implementations
describe them in detail below.
-->
None



Discussion:

Resolving patch:
diff --git a/celery/backends/redis.py b/celery/backends/redis.py
--- a/celery/backends/redis.py
+++ b/celery/backends/redis.py
@@ -3,6 +3,7 @@
 from __future__ import absolute_import, unicode_literals
 
 import time
+from contextlib import contextmanager
 from functools import partial
 from ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED
 
@@ -78,6 +79,11 @@
 
 E_LOST = 'Connection to Redis lost: Retry (%s/%s) %s.'
 
+E_RETRY_LIMIT_EXCEEDED = """
+Retry limit exceeded while trying to reconnect to the Celery redis result \
+store backend. The Celery application must be restarted.
+"""
+
 logger = get_logger(__name__)
 
 
@@ -88,6 +94,8 @@ def __init__(self, *args, **kwargs):
         super(ResultConsumer, self).__init__(*args, **kwargs)
         self._get_key_for_task = self.backend.get_key_for_task
         self._decode_result = self.backend.decode_result
+        self._ensure = self.backend.ensure
+        self._connection_errors = self.backend.connection_errors
         self.subscribed_to = set()
 
     def on_after_fork(self):
@@ -99,6 +107,31 @@ def on_after_fork(self):
             logger.warning(text_t(e))
         super(ResultConsumer, self).on_after_fork()
 
+    def _reconnect_pubsub(self):
+        self._pubsub = None
+        self.backend.client.connection_pool.reset()
+        # task state might have changed when the connection was down so we
+        # retrieve meta for all subscribed tasks before going into pubsub mode
+        metas = self.backend.client.mget(self.subscribed_to)
+        metas = [meta for meta in metas if meta]
+        for meta in metas:
+            self.on_state_change(self._decode_result(meta), None)
+        self._pubsub = self.backend.client.pubsub(
+            ignore_subscribe_messages=True,
+        )
+        self._pubsub.subscribe(*self.subscribed_to)
+
+    @contextmanager
+    def reconnect_on_error(self):
+        try:
+            yield
+        except self._connection_errors:
+            try:
+                self._ensure(self._reconnect_pubsub, ())
+            except self._connection_errors:
+                logger.critical(E_RETRY_LIMIT_EXCEEDED)
+                raise
+
     def _maybe_cancel_ready_task(self, meta):
         if meta['status'] in states.READY_STATES:
             self.cancel_for(meta['task_id'])
@@ -124,9 +157,10 @@ def stop(self):
 
     def drain_events(self, timeout=None):
         if self._pubsub:
-            message = self._pubsub.get_message(timeout=timeout)
-            if message and message['type'] == 'message':
-                self.on_state_change(self._decode_result(message['data']), message)
+            with self.reconnect_on_error():
+                message = self._pubsub.get_message(timeout=timeout)
+                if message and message['type'] == 'message':
+                    self.on_state_change(self._decode_result(message['data']), message)
         elif timeout:
             time.sleep(timeout)
 
@@ -139,13 +173,15 @@ def _consume_from(self, task_id):
         key = self._get_key_for_task(task_id)
         if key not in self.subscribed_to:
             self.subscribed_to.add(key)
-            self._pubsub.subscribe(key)
+            with self.reconnect_on_error():
+                self._pubsub.subscribe(key)
 
     def cancel_for(self, task_id):
+        key = self._get_key_for_task(task_id)
+        self.subscribed_to.discard(key)
         if self._pubsub:
-            key = self._get_key_for_task(task_id)
-            self.subscribed_to.discard(key)
-            self._pubsub.unsubscribe(key)
+            with self.reconnect_on_error():
+                self._pubsub.unsubscribe(key)
 
 
 class RedisBackend(BaseKeyValueStoreBackend, AsyncBackendMixin):


Test patch:
diff --git a/t/unit/backends/test_redis.py b/t/unit/backends/test_redis.py
--- a/t/unit/backends/test_redis.py
+++ b/t/unit/backends/test_redis.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, unicode_literals
 
+import json
 import random
 import ssl
 from contextlib import contextmanager
@@ -26,6 +27,10 @@ def on_first_call(*args, **kwargs):
         mock.return_value, = retval
 
 
+class ConnectionError(Exception):
+    pass
+
+
 class Connection(object):
     connected = True
 
@@ -55,9 +60,27 @@ def execute(self):
         return [step(*a, **kw) for step, a, kw in self.steps]
 
 
+class PubSub(mock.MockCallbacks):
+    def __init__(self, ignore_subscribe_messages=False):
+        self._subscribed_to = set()
+
+    def close(self):
+        self._subscribed_to = set()
+
+    def subscribe(self, *args):
+        self._subscribed_to.update(args)
+
+    def unsubscribe(self, *args):
+        self._subscribed_to.difference_update(args)
+
+    def get_message(self, timeout=None):
+        pass
+
+
 class Redis(mock.MockCallbacks):
     Connection = Connection
     Pipeline = Pipeline
+    pubsub = PubSub
 
     def __init__(self, host=None, port=None, db=None, password=None, **kw):
         self.host = host
@@ -71,6 +94,9 @@ def __init__(self, host=None, port=None, db=None, password=None, **kw):
     def get(self, key):
         return self.keyspace.get(key)
 
+    def mget(self, keys):
+        return [self.get(key) for key in keys]
+
     def setex(self, key, expires, value):
         self.set(key, value)
         self.expire(key, expires)
@@ -144,7 +170,9 @@ class _RedisBackend(RedisBackend):
         return _RedisBackend(app=self.app)
 
     def get_consumer(self):
-        return self.get_backend().result_consumer
+        consumer = self.get_backend().result_consumer
+        consumer._connection_errors = (ConnectionError,)
+        return consumer
 
     @patch('celery.backends.asynchronous.BaseResultConsumer.on_after_fork')
     def test_on_after_fork(self, parent_method):
@@ -194,6 +222,33 @@ def test_drain_events_before_start(self):
         # drain_events shouldn't crash when called before start
         consumer.drain_events(0.001)
 
+    def test_consume_from_connection_error(self):
+        consumer = self.get_consumer()
+        consumer.start('initial')
+        consumer._pubsub.subscribe.side_effect = (ConnectionError(), None)
+        consumer.consume_from('some-task')
+        assert consumer._pubsub._subscribed_to == {b'celery-task-meta-initial', b'celery-task-meta-some-task'}
+
+    def test_cancel_for_connection_error(self):
+        consumer = self.get_consumer()
+        consumer.start('initial')
+        consumer._pubsub.unsubscribe.side_effect = ConnectionError()
+        consumer.consume_from('some-task')
+        consumer.cancel_for('some-task')
+        assert consumer._pubsub._subscribed_to == {b'celery-task-meta-initial'}
+
+    @patch('celery.backends.redis.ResultConsumer.cancel_for')
+    @patch('celery.backends.asynchronous.BaseResultConsumer.on_state_change')
+    def test_drain_events_connection_error(self, parent_on_state_change, cancel_for):
+        meta = {'task_id': 'initial', 'status': states.SUCCESS}
+        consumer = self.get_consumer()
+        consumer.start('initial')
+        consumer.backend.set(b'celery-task-meta-initial', json.dumps(meta))
+        consumer._pubsub.get_message.side_effect = ConnectionError()
+        consumer.drain_events()
+        parent_on_state_change.assert_called_with(meta, None)
+        assert consumer._pubsub._subscribed_to == {b'celery-task-meta-initial'}
+
 
 class test_RedisBackend:
     def get_backend(self):