Repository: redis/redis-py

Issue:
Sub-unsub-resub causes redispy to forget the channel
We recently came across some strange behaviour in our web servers, and tracked the problem down to a race condition in redispy, which is triggered when you unsubscribe and quickly resubscribe to a channel.

I can trivially reproduce it with this:

```
pubsub.subscribe("channel1")
# pubsub.channels is {"channel1": None}
time.sleep(1)
pubsub.unsubscribe("channel1")
pubsub.subscribe("channel1")
time.sleep(1)
# pubsub.channels is {} and pubsub.subscribed is False!!
```

pubsub.subscribe() synchronously adds the channel to pubsub.channels, however pubsub.unsubscribe() doesn't: it waits for the redis server to asynchronously send an unsubscribe message.

So, what happens is, the unsubscribe request is sent, subscribe() readds the channel to the dict, but then the unsubscribe message is returned from the server and the newly-added entry is removed from the dict. Both our application and the redis server end up thinking the channel is subscribed, however redispy has forgotten about it (meaning our callback function becomes unregistered and we stop receiving messages). In theory, the same can happen with the pattern calls.

A possible fix is to just not wait for the server to respond, e.g:

```
class RedisPyPubSubFix(PubSub):
    _sentinel = object()

    def __init__(self, redis_client, **kwargs):
        super(RedisPyPubSubFix, self).__init__(redis_client.connection_pool, **kwargs)

    def unsubscribe(self, *args):
        resp = super(RedisPyPubSubFix, self).unsubscribe(*args)
        if not args:
            # Unsub(empty) is unsub all
            self.channels.clear()
        else:
            for arg in args:
                self.channels.pop(self.encode(arg), None)
        return resp

    def punsubscribe(self, *args):
        resp = super(RedisPyPubSubFix, self).punsubscribe(*args)
        if not args:
            self.patterns.clear()
        else:
            for arg in args:
                self.patterns.pop(self.encode(arg), None)
        return resp

    def handle_message(self, response, ignore_subscribe_messages=False):
        """
        Parses a pub/sub message. If the channel or pattern was subscribed to
        with a message handler, the handler is invoked instead of a parsed
        message being returned.
        """
        message_type = nativestr(response[0])
        if message_type == 'pmessage':
            message = {
                'type': message_type,
                'pattern': response[1],
                'channel': response[2],
                'data': response[3]
            }
        else:
            message = {
                'type': message_type,
                'pattern': None,
                'channel': response[1],
                'data': response[2]
            }

        if message_type in self.PUBLISH_MESSAGE_TYPES:
            # if there's a message handler, invoke it
            if message_type == 'pmessage':
                handler = self.patterns.get(message['pattern'], self._sentinel)
            else:
                handler = self.channels.get(message['channel'], self._sentinel)
            if handler is self._sentinel:
                # We're no longer subbed to the channel, so we
                # don't care about the message.
                return None
            if handler:
                handler(message)
                return None
        else:
            # this is a subscribe/unsubscribe message. ignore if we don't
            # want them
            if ignore_subscribe_messages or self.ignore_subscribe_messages:
                return None

        return message
```



Discussion:
Until the client receives confirmation of the unsubscribe, messages can still arrive on that channel. That's why `handle_message` removes the channel/pattern rather than the `unsubscribe` method itself.

Perhaps we have another dict on the pubsub object that is a temporary store for channels/patterns  being subscribed to. The `subscribe` methods write to this new dict, and then when we receive confirmation of the subscription in `handle_message`, we move it from the temporary store to the `channels` or `patterns` dicts.

Hi Andy, thanks very much for your response (and the excellent redis client)!

I suspected that was the intention, though I'm not sure it's the correct approach. I think that messages received after the application has unsubscribed from a channel can be safely ignored, since obviously the user code doesn't expect or want to receive them.

When I was writing the code I posted earlier, I had already realised that since the channels/patterns dicts can be updated synchronously, they can also be used to filter the messages against. Specifically, the line after the "don't care about the message" comment would filter out messages from channels we're no longer listening to.

I wrote a simple test function which shows a problem.
Problem is since 2.10.0 version.

```python
def test_sub_unsub_sub():
    # redis since 2.10.0 to newest (3.0.1)
    import redis

    cli = redis.Redis()
    pubsub = cli.pubsub()
    pubsub.subscribe('test-chan')
    pubsub.unsubscribe('test-chan')
    pubsub.subscribe('test-chan')
    assert pubsub.subscribed

    message = pubsub.get_message()
    assert message
    assert message['channel'] == 'test-chan', message
    assert message['type'] == 'subscribe', message

    message = pubsub.get_message()
    assert message
    assert message['channel'] == 'test-chan', message
    assert message['type'] == 'unsubscribe', message

    message = pubsub.get_message()
    assert message
    assert message['channel'] == 'test-chan', message
    assert message['type'] == 'subscribe', message

    assert pubsub.subscribed  #  <-- this fails!

if __name__ == '__main__':
    test_sub_unsub_sub()
```
The ugly consequence of this bug is that in case of reconnect, our channel will be not re-subscribed
(because `self.channels` does not have it)!

```python
    def on_connect(self, connection):
        "Re-subscribe to any channels and patterns previously subscribed to"
        if self.channels:
            channels = {}
            for k, v in iteritems(self.channels):
                if not self.decode_responses:
                    k = k.decode(self.encoding, self.encoding_errors)
                channels[k] = v
            self.subscribe(**channels)
```
> A possible fix is to just not wait for the server to respond, e.g:
> 
> ```
>     def unsubscribe(self, *args):
>         resp = super(RedisPyPubSubFix, self).unsubscribe(*args)
>         if not args:
>             # Unsub(empty) is unsub all
>             self.channels.clear()
>         else:
>             for arg in args:
>                 self.channels.pop(self.encode(arg), None)
>         return resp

@costasgambit The proposed fix was similar to code before 2.10.0 which was:

```python
    def unsubscribe(self, channels=[]):
        if isinstance(channels, basestring):
            channels = [channels]
        for channel in channels:
            try:
                self.channels.remove(channel)
            except KeyError:
                pass
        return self.execute_command('UNSUBSCRIBE', *channels)
```

The changing commit was https://github.com/andymccurdy/redis-py/commit/169cc0241a4f8824f35add9a872c067b70b1e95c commented as "fixes PubSub.subscribe once and for all. "  :)
Possible fix (non breaking current semantic of `unsubscribe`) is to add `pending_unsubscribe` set which keeps channels to be unsubscribed:

```python
    def __init__(...):
       ...
       self._pending_unsubscribe = set()

    def on_connect(self, connection):
      ...
      self._pending_unsubscribe.clear()

    def subscribe(self, *args, **kwargs):
        ...
        self._pending_unsubscribe.difference_update(new_channels)
        return ret_val

    def unsubscribe(self, *args):
        ...
        ret_val = self.execute_command('UNSUBSCRIBE', *args)
        self._pending_unsubscribe.update(args)
        return ret_val

    def handle_message(self, response, ignore_subscribe_messages=False):
       ...
        if message_type in self.UNSUBSCRIBE_MESSAGE_TYPES:
            channel = message['channel']
            if channel in self._pending_unsubscribe:
                ...
                self._pending_unsubscribe.remove(channel)
                self.channels.pop(channel, None)
```
Is there any contributor, here? I could prepare a pull request, but I see a lot of pull requests (78) nobody cares of - it is very sad.
A PR would be great. The `pending_unsubscribe` solution seems good. 

Resolving patch:
diff --git a/redis/client.py b/redis/client.py
--- a/redis/client.py
+++ b/redis/client.py
@@ -2903,7 +2903,9 @@ def reset(self):
             self.connection_pool.release(self.connection)
             self.connection = None
         self.channels = {}
+        self.pending_unsubscribe_channels = set()
         self.patterns = {}
+        self.pending_unsubscribe_patterns = set()
 
     def close(self):
         self.reset()
@@ -2913,6 +2915,8 @@ def on_connect(self, connection):
         # NOTE: for python3, we can't pass bytestrings as keyword arguments
         # so we need to decode channel/pattern names back to unicode strings
         # before passing them to [p]subscribe.
+        self.pending_unsubscribe_channels.clear()
+        self.pending_unsubscribe_patterns.clear()
         if self.channels:
             channels = {}
             for k, v in iteritems(self.channels):
@@ -2999,17 +3003,25 @@ def psubscribe(self, *args, **kwargs):
         # update the patterns dict AFTER we send the command. we don't want to
         # subscribe twice to these patterns, once for the command and again
         # for the reconnection.
-        self.patterns.update(self._normalize_keys(new_patterns))
+        new_patterns = self._normalize_keys(new_patterns)
+        self.patterns.update(new_patterns)
+        self.pending_unsubscribe_patterns.difference_update(new_patterns)
         return ret_val
 
     def punsubscribe(self, *args):
         """
-        Unsubscribe from the supplied patterns. If empy, unsubscribe from
+        Unsubscribe from the supplied patterns. If empty, unsubscribe from
         all patterns.
         """
         if args:
             args = list_or_args(args[0], args[1:])
-        return self.execute_command('PUNSUBSCRIBE', *args)
+        retval = self.execute_command('PUNSUBSCRIBE', *args)
+        if args:
+            patterns = self._normalize_keys(dict.fromkeys(args))
+        else:
+            patterns = self.patterns
+        self.pending_unsubscribe_patterns.update(patterns)
+        return retval
 
     def subscribe(self, *args, **kwargs):
         """
@@ -3027,7 +3039,9 @@ def subscribe(self, *args, **kwargs):
         # update the channels dict AFTER we send the command. we don't want to
         # subscribe twice to these channels, once for the command and again
         # for the reconnection.
-        self.channels.update(self._normalize_keys(new_channels))
+        new_channels = self._normalize_keys(new_channels)
+        self.channels.update(new_channels)
+        self.pending_unsubscribe_channels.difference_update(new_channels)
         return ret_val
 
     def unsubscribe(self, *args):
@@ -3037,7 +3051,13 @@ def unsubscribe(self, *args):
         """
         if args:
             args = list_or_args(args[0], args[1:])
-        return self.execute_command('UNSUBSCRIBE', *args)
+        retval = self.execute_command('UNSUBSCRIBE', *args)
+        if args:
+            channels = self._normalize_keys(dict.fromkeys(args))
+        else:
+            channels = self.channels
+        self.pending_unsubscribe_channels.update(channels)
+        return retval
 
     def listen(self):
         "Listen for messages on channels this client has been subscribed to"
@@ -3094,22 +3114,21 @@ def handle_message(self, response, ignore_subscribe_messages=False):
                 'channel': response[1],
                 'data': response[2]
             }
-
         # if this is an unsubscribe message, remove it from memory
         if message_type in self.UNSUBSCRIBE_MESSAGE_TYPES:
-            subscribed_dict = None
             if message_type == 'punsubscribe':
-                subscribed_dict = self.patterns
+                pattern = response[1]
+                if pattern in self.pending_unsubscribe_patterns:
+                    self.pending_unsubscribe_patterns.remove(pattern)
+                    self.patterns.pop(pattern, None)
             else:
-                subscribed_dict = self.channels
-            try:
-                del subscribed_dict[message['channel']]
-            except KeyError:
-                pass
+                channel = response[1]
+                if channel in self.pending_unsubscribe_channels:
+                    self.pending_unsubscribe_channels.remove(channel)
+                    self.channels.pop(channel, None)
 
         if message_type in self.PUBLISH_MESSAGE_TYPES:
             # if there's a message handler, invoke it
-            handler = None
             if message_type == 'pmessage':
                 handler = self.patterns.get(message['pattern'], None)
             else:


Test patch:
diff --git a/tests/test_pubsub.py b/tests/test_pubsub.py
--- a/tests/test_pubsub.py
+++ b/tests/test_pubsub.py
@@ -212,6 +212,48 @@ def test_ignore_individual_subscribe_messages(self, r):
             assert message is None
         assert p.subscribed is False
 
+    def test_sub_unsub_resub_channels(self, r):
+        kwargs = make_subscribe_test_data(r.pubsub(), 'channel')
+        self._test_sub_unsub_resub(**kwargs)
+
+    def test_sub_unsub_resub_patterns(self, r):
+        kwargs = make_subscribe_test_data(r.pubsub(), 'pattern')
+        self._test_sub_unsub_resub(**kwargs)
+
+    def _test_sub_unsub_resub(self, p, sub_type, unsub_type, sub_func,
+                              unsub_func, keys):
+        # https://github.com/andymccurdy/redis-py/issues/764
+        key = keys[0]
+        sub_func(key)
+        unsub_func(key)
+        sub_func(key)
+        assert p.subscribed is True
+        assert wait_for_message(p) == make_message(sub_type, key, 1)
+        assert wait_for_message(p) == make_message(unsub_type, key, 0)
+        assert wait_for_message(p) == make_message(sub_type, key, 1)
+        assert p.subscribed is True
+
+    def test_sub_unsub_all_resub_channels(self, r):
+        kwargs = make_subscribe_test_data(r.pubsub(), 'channel')
+        self._test_sub_unsub_all_resub(**kwargs)
+
+    def test_sub_unsub_all_resub_patterns(self, r):
+        kwargs = make_subscribe_test_data(r.pubsub(), 'pattern')
+        self._test_sub_unsub_all_resub(**kwargs)
+
+    def _test_sub_unsub_all_resub(self, p, sub_type, unsub_type, sub_func,
+                                  unsub_func, keys):
+        # https://github.com/andymccurdy/redis-py/issues/764
+        key = keys[0]
+        sub_func(key)
+        unsub_func()
+        sub_func(key)
+        assert p.subscribed is True
+        assert wait_for_message(p) == make_message(sub_type, key, 1)
+        assert wait_for_message(p) == make_message(unsub_type, key, 0)
+        assert wait_for_message(p) == make_message(sub_type, key, 1)
+        assert p.subscribed is True
+
 
 class TestPubSubMessages(object):
     def setup_method(self, method):