Repository: redis/redis-py

Issue:
Changed the PubSub's health check command to be performed only on the…
… first command execution.

### Pull Request check-list

_Please make sure to review and check all of these items:_

- [V ] Does `$ tox` pass with this change (including linting)?
- [V ] Do the CI tests pass with this change (enable it first in your forked repo and wait for the github action build to finish)?
- [ V] Is the new or changed code fully tested?
- [ V] Is a documentation update included (if this change modifies existing APIs, or introduces new ones)?

_NOTE: these things are not required to open a PR and can be done
afterwards / while the PR is open._

### Description of change

Closing #1720:

I reproduced the PubSub bug described in #1720 and was able to find the RC.
When the PubSub's 'execute_command' method is being called it passes a 'health_check' bool to determine if it needs to run a health check. The 'health_check' value is set to **not self.subscribed**, which checks if the pubsub instance has any items in the channels/patterns lists. That means, that we perform a health check within the execute_command function only if we are not yet subscribed. All subsequent commands, after the first subscription, should be executed without performing a health check, since the channels/patterns list is no longer empty.
The pubsub's 'get_messages()' method can be used to poll published messages after a pubsub instance has been created. If a poller thread is created (thread that waits on get_message()) it will listen on the same socket as the pubsub execute_command is listening to when it performs a health check. Hence, we should not send a healthcheck using the pubsub execute_command function after the poller thread is initiated, since then it will be racing the poller thread to read the response from the socket.

In the example in #1720 we see the following flow:

1. PubSub instance is being created
2. Channel 'foo' is being subscribed - health check is performed since self.channels is empty
3. 'PONG' is received 
4. A poller thread is being started, looping on 'get_messages()'
5. The poller thread polls the 'subscribe' response
6. Channel 'foo' is being unsubscribed, health check is not being performed since self.channels still contains 'foo'. 'foo' is removed from self.channels
7. The poller thread polls the 'unsubscribe' response
8. Channel 'baz' is being subscribed - health check is performed since self.channels is empty again
9. The poller thread tries to poll the 'subscribe' response and gets the 'PONG' response instead
10. The health check waits for a PONG response that has already been received by the poller thread, and is therefore timed out

Therefore, we shouldn't use self.channels and self.patterns to determine whether a health check needs to be executed, but we should have another variable to indicate whether this is the first command execution, and if so, to run a health check. 

However, a poller thread may be started before subscribing to a channel, e.g. :

1. ps = r.pubsub()
2. poller = threading.Thread(target=poll, args=(ps,))
3. ps.subscribe('foo')

In this case, the health check will be performed and we will still get a race reading from the socket with the poller thread. 
So, my suggestion is to add a new 'cmd_execution_health_check' variable initiated with 'True' to the pubsub class and to set it to False on: 
1. The end of execute_command method, so the health check will be performed only on the first execution), or
2. get_message() function, so the health check will not be performed from the execute_command function at all.
health checks are being done by the get_message() method, so no need to execute it also from the main command execution.  

This change fixes the reported bug.


Discussion:

Resolving patch:
diff --git a/redis/client.py b/redis/client.py
--- a/redis/client.py
+++ b/redis/client.py
@@ -1276,18 +1276,17 @@ def __init__(
         self.shard_hint = shard_hint
         self.ignore_subscribe_messages = ignore_subscribe_messages
         self.connection = None
+        self.subscribed_event = threading.Event()
         # we need to know the encoding options for this connection in order
         # to lookup channel and pattern names for callback handlers.
         self.encoder = encoder
         if self.encoder is None:
             self.encoder = self.connection_pool.get_encoder()
+        self.health_check_response_b = self.encoder.encode(self.HEALTH_CHECK_MESSAGE)
         if self.encoder.decode_responses:
             self.health_check_response = ["pong", self.HEALTH_CHECK_MESSAGE]
         else:
-            self.health_check_response = [
-                b"pong",
-                self.encoder.encode(self.HEALTH_CHECK_MESSAGE),
-            ]
+            self.health_check_response = [b"pong", self.health_check_response_b]
         self.reset()
 
     def __enter__(self):
@@ -1312,9 +1311,11 @@ def reset(self):
             self.connection_pool.release(self.connection)
             self.connection = None
         self.channels = {}
+        self.health_check_response_counter = 0
         self.pending_unsubscribe_channels = set()
         self.patterns = {}
         self.pending_unsubscribe_patterns = set()
+        self.subscribed_event.clear()
 
     def close(self):
         self.reset()
@@ -1340,7 +1341,7 @@ def on_connect(self, connection):
     @property
     def subscribed(self):
         "Indicates if there are subscriptions to any channels or patterns"
-        return bool(self.channels or self.patterns)
+        return self.subscribed_event.is_set()
 
     def execute_command(self, *args):
         "Execute a publish/subscribe command"
@@ -1358,8 +1359,28 @@ def execute_command(self, *args):
             self.connection.register_connect_callback(self.on_connect)
         connection = self.connection
         kwargs = {"check_health": not self.subscribed}
+        if not self.subscribed:
+            self.clean_health_check_responses()
         self._execute(connection, connection.send_command, *args, **kwargs)
 
+    def clean_health_check_responses(self):
+        """
+        If any health check responses are present, clean them
+        """
+        ttl = 10
+        conn = self.connection
+        while self.health_check_response_counter > 0 and ttl > 0:
+            if self._execute(conn, conn.can_read, timeout=conn.socket_timeout):
+                response = self._execute(conn, conn.read_response)
+                if self.is_health_check_response(response):
+                    self.health_check_response_counter -= 1
+                else:
+                    raise PubSubError(
+                        "A non health check response was cleaned by "
+                        "execute_command: {0}".format(response)
+                    )
+            ttl -= 1
+
     def _disconnect_raise_connect(self, conn, error):
         """
         Close the connection and raise an exception
@@ -1399,11 +1420,23 @@ def parse_response(self, block=True, timeout=0):
             return None
         response = self._execute(conn, conn.read_response)
 
-        if conn.health_check_interval and response == self.health_check_response:
+        if self.is_health_check_response(response):
             # ignore the health check message as user might not expect it
+            self.health_check_response_counter -= 1
             return None
         return response
 
+    def is_health_check_response(self, response):
+        """
+        Check if the response is a health check response.
+        If there are no subscriptions redis responds to PING command with a
+        bulk response, instead of a multi-bulk with "pong" and the response.
+        """
+        return response in [
+            self.health_check_response,  # If there was a subscription
+            self.health_check_response_b,  # If there wasn't
+        ]
+
     def check_health(self):
         conn = self.connection
         if conn is None:
@@ -1414,6 +1447,7 @@ def check_health(self):
 
         if conn.health_check_interval and time.time() > conn.next_health_check:
             conn.send_command("PING", self.HEALTH_CHECK_MESSAGE, check_health=False)
+            self.health_check_response_counter += 1
 
     def _normalize_keys(self, data):
         """
@@ -1443,6 +1477,11 @@ def psubscribe(self, *args, **kwargs):
         # for the reconnection.
         new_patterns = self._normalize_keys(new_patterns)
         self.patterns.update(new_patterns)
+        if not self.subscribed:
+            # Set the subscribed_event flag to True
+            self.subscribed_event.set()
+            # Clear the health check counter
+            self.health_check_response_counter = 0
         self.pending_unsubscribe_patterns.difference_update(new_patterns)
         return ret_val
 
@@ -1477,6 +1516,11 @@ def subscribe(self, *args, **kwargs):
         # for the reconnection.
         new_channels = self._normalize_keys(new_channels)
         self.channels.update(new_channels)
+        if not self.subscribed:
+            # Set the subscribed_event flag to True
+            self.subscribed_event.set()
+            # Clear the health check counter
+            self.health_check_response_counter = 0
         self.pending_unsubscribe_channels.difference_update(new_channels)
         return ret_val
 
@@ -1508,6 +1552,20 @@ def get_message(self, ignore_subscribe_messages=False, timeout=0):
         before returning. Timeout should be specified as a floating point
         number.
         """
+        if not self.subscribed:
+            # Wait for subscription
+            start_time = time.time()
+            if self.subscribed_event.wait(timeout) is True:
+                # The connection was subscribed during the timeout time frame.
+                # The timeout should be adjusted based on the time spent
+                # waiting for the subscription
+                time_spent = time.time() - start_time
+                timeout = max(0.0, timeout - time_spent)
+            else:
+                # The connection isn't subscribed to any channels or patterns,
+                # so no messages are available
+                return None
+
         response = self.parse_response(block=False, timeout=timeout)
         if response:
             return self.handle_message(response, ignore_subscribe_messages)
@@ -1561,6 +1619,10 @@ def handle_message(self, response, ignore_subscribe_messages=False):
                 if channel in self.pending_unsubscribe_channels:
                     self.pending_unsubscribe_channels.remove(channel)
                     self.channels.pop(channel, None)
+            if not self.channels and not self.patterns:
+                # There are no subscriptions anymore, set subscribed_event flag
+                # to false
+                self.subscribed_event.clear()
 
         if message_type in self.PUBLISH_MESSAGE_TYPES:
             # if there's a message handler, invoke it


Test patch:
diff --git a/tests/test_pubsub.py b/tests/test_pubsub.py
--- a/tests/test_pubsub.py
+++ b/tests/test_pubsub.py
@@ -2,6 +2,7 @@
 import threading
 import time
 from unittest import mock
+from unittest.mock import patch
 
 import pytest
 
@@ -348,15 +349,6 @@ def test_unicode_pattern_message_handler(self, r):
             "pmessage", channel, "test message", pattern=pattern
         )
 
-    def test_get_message_without_subscribe(self, r):
-        p = r.pubsub()
-        with pytest.raises(RuntimeError) as info:
-            p.get_message()
-        expect = (
-            "connection not set: " "did you forget to call subscribe() or psubscribe()?"
-        )
-        assert expect in info.exconly()
-
 
 class TestPubSubAutoDecoding:
     "These tests only validate that we get unicode values back"
@@ -549,6 +541,39 @@ def test_get_message_with_timeout_returns_none(self, r):
         assert wait_for_message(p) == make_message("subscribe", "foo", 1)
         assert p.get_message(timeout=0.01) is None
 
+    def test_get_message_not_subscribed_return_none(self, r):
+        p = r.pubsub()
+        assert p.subscribed is False
+        assert p.get_message() is None
+        assert p.get_message(timeout=0.1) is None
+        with patch.object(threading.Event, "wait") as mock:
+            mock.return_value = False
+            assert p.get_message(timeout=0.01) is None
+            assert mock.called
+
+    def test_get_message_subscribe_during_waiting(self, r):
+        p = r.pubsub()
+
+        def poll(ps, expected_res):
+            assert ps.get_message() is None
+            message = ps.get_message(timeout=1)
+            assert message == expected_res
+
+        subscribe_response = make_message("subscribe", "foo", 1)
+        poller = threading.Thread(target=poll, args=(p, subscribe_response))
+        poller.start()
+        time.sleep(0.2)
+        p.subscribe("foo")
+        poller.join()
+
+    def test_get_message_wait_for_subscription_not_being_called(self, r):
+        p = r.pubsub()
+        p.subscribe("foo")
+        with patch.object(threading.Event, "wait") as mock:
+            assert p.subscribed is True
+            assert wait_for_message(p) == make_message("subscribe", "foo", 1)
+            assert mock.called is False
+
 
 class TestPubSubWorkerThread:
     @pytest.mark.skipif(