Bug in retry-handling in the PubSub

[ManelCoutinhoSensei]

Hey y'all!

It seems like the retry-functionality in the PubSub-client is not working correctly.

Reproduce Error

To reproduce the error I am closing valkey while my system is trying to get new messages in a parallel thread. Even though our connection is set up with an infinite amount of retry, the PubSub-client always crashes and never recovers.
Here is a minimal pseudo code to help reproduce the issue:

def event_handler(msg):
    print(msg)

vk_pub_sub.psubscribe(**{"test": event_handler})

def exception_handler(ex, pubsub, thread):
    print(traceback.format_exc())

vk_pub_sub.run_in_thread(sleep_time=0.1, exception_handler=exception_handler)

# Close VK connection

Error Trace

Traceback (most recent call last):
  File "/path/to/env/lib/python3.10/site-packages/valkey/retry.py", line 62, in call_with_retry
    return do()
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 860, in <lambda>
    lambda: command(*args, **kwargs),
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 877, in try_read
    if not conn.can_read(timeout=timeout):
  File "/path/to/env/lib/python3.10/site-packages/valkey/connection.py", line 538, in can_read
    return self._parser.can_read(timeout)
  File "/path/to/env/lib/python3.10/site-packages/valkey/_parsers/libvalkey.py", line 80, in can_read
    return self.read_from_socket(timeout=timeout, raise_on_timeout=False)
  File "/path/to/env/lib/python3.10/site-packages/valkey/_parsers/libvalkey.py", line 91, in read_from_socket
    raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
valkey.exceptions.ConnectionError: Connection closed by server.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/path/to/env/lib/python3.10/site-packages/valkey/connection.py", line 560, in read_response
    response = self._parser.read_response(disable_decoding=disable_decoding)
  File "/path/to/env/lib/python3.10/site-packages/valkey/_parsers/libvalkey.py", line 129, in read_response
    self.read_from_socket()
  File "/path/to/env/lib/python3.10/site-packages/valkey/_parsers/libvalkey.py", line 89, in read_from_socket
    bufflen = self._sock.recv_into(self._buffer)
ConnectionResetError: [Errno 104] Connection reset by peer

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 1210, in run
    pubsub.get_message(ignore_subscribe_messages=True, timeout=sleep_time)
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 1072, in get_message
    response = self.parse_response(block=(timeout is None), timeout=timeout)
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 883, in parse_response
    response = self._execute(conn, try_read)
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 859, in _execute
    return conn.retry.call_with_retry(
  File "/path/to/env/lib/python3.10/site-packages/valkey/retry.py", line 65, in call_with_retry
    fail(error)
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 861, in <lambda>
    lambda error: self._disconnect_raise_connect(conn, error),
  File "/path/to/env/lib/python3.10/site-packages/valkey/client.py", line 849, in _disconnect_raise_connect
    conn.connect()
  File "/path/to/env/lib/python3.10/site-packages/valkey/connection.py", line 328, in connect
    self.on_connect()
  File "/path/to/env/lib/python3.10/site-packages/valkey/connection.py", line 430, in on_connect
    self.read_response()
  File "/path/to/env/lib/python3.10/site-packages/valkey/connection.py", line 568, in read_response
    raise ConnectionError(
valkey.exceptions.ConnectionError: Error while reading from localhost:6379 : (104, 'Connection reset by peer')

Possible Cause

From my debug and the traces, it seems that the issue is related to the connect method since the second part of the error trace occurs outside the retry scope.

Quick Patch

A quick workaround is to wrap everything inside a retry call, as shown below. However, this might introduce unexpected side effects:

def connect(self):
    """Connects to the Valkey server if not already connected."""
    if self._sock:
        return

    def _full_connect():
        self._sock = self._connect()
        if self.valkey_connect_func is None:
            # Use the default on_connect function
            self.on_connect()
        else:
            # Use the passed function valkey_connect_func
            self.valkey_connect_func(self)

        # run any user callbacks. right now the only internal callback
        # is for pubsub channel/pattern resubscription
        # first, remove any dead weakrefs
        self._connect_callbacks = [ref for ref in self._connect_callbacks if ref()]
        for ref in self._connect_callbacks:
            callback = ref()
            if callback:
                callback(self)

    try:
        sock = self.retry.call_with_retry(
            lambda: _full_connect(),
            lambda error: self.disconnect(error)
        )
    except socket.timeout:
        raise TimeoutError("Timeout connecting to server")
    except OSError as e:
        raise ConnectionError(self._error_message(e))

Thanks in advance for any help 🙂

[mkmkme]

Hey @ManelCoutinhoSensei,

Thanks for your feedback! Sorry it took me a while to respond, was quite busy lately.

Also I tried to create a small reproducer to understand your problem but failed. Here's a script I used:

#!/usr/bin/env python3

import time
import traceback
import valkey


def event_handler(msg):
    print(msg)


def exception_handler(ex, pubsub, thread):
    print(traceback.format_exc())


def main():
    client = valkey.ValkeyCluster.from_url("valkey://localhost:16379/0")
    vk_pubsub = client.pubsub()
    vk_pubsub.psubscribe(**{"test": event_handler})
    vk_pubsub.run_in_thread(
        sleep_time=0.1, daemon=True, exception_handler=exception_handler
    )

    while True:
        time.sleep(1)


if __name__ == "__main__":
    main()

Then, I did invoke devenv to start a Valkey cluster to connect to and started the script.
When I stopped the cluster via docker compose --profile all down, my terminal indeed started flooding with exceptions. However, when I started the cluster again, the exceptions stopped popping up and I managed to publish a new message and receive it in this client output:

......................
  File "/home/mikhail.koviazin/workspace/valkey-py/venv/lib64/python3.12/site-packages/valkey/connection.py", line 533, in can_read
    self.connect()
  File "/home/mikhail.koviazin/workspace/valkey-py/venv/lib64/python3.12/site-packages/valkey/connection.py", line 331, in connect
    self.valkey_connect_func(self)
  File "/home/mikhail.koviazin/workspace/valkey-py/venv/lib64/python3.12/site-packages/valkey/cluster.py", line 676, in on_connect
    connection.on_connect()
  File "/home/mikhail.koviazin/workspace/valkey-py/venv/lib64/python3.12/site-packages/valkey/connection.py", line 430, in on_connect
    self.read_response()
  File "/home/mikhail.koviazin/workspace/valkey-py/venv/lib64/python3.12/site-packages/valkey/connection.py", line 568, in read_response
    raise ConnectionError(
valkey.exceptions.ConnectionError: Error while reading from 127.0.0.1:16380 : (104, 'Connection reset by peer')





{'type': 'pmessage', 'pattern': b'test', 'channel': b'test', 'data': b'HELL'}

Please tell me:

1. Did I understand your issue correctly? If not, what did I do wrong?
2. Could you retry my steps and tell me whether the behaviour is same or different?

Thanks!

[ManelCoutinhoSensei]

Hey @mkmkme,

Sorry for the delayed response.

You're absolutely right in your assessment. My example was intended only to demonstrate the error, which I think, conceptually, shouldn't occur given where/how it happens in the code.

---

Let me try to explain it more clearly.

If you modify your code as follows:

def event_handler(msg):
    print(msg)


def exception_handler(ex, pubsub, thread):
    print(traceback.format_exc())


def main():
    client = valkey.Valkey.from_url("valkey://localhost:6379/0", Retry(ExponentialBackoff(10, 0.5), 5))
    vk_pubsub = client.pubsub()
    vk_pubsub.psubscribe(**{"test": event_handler})
    vk_pubsub.run_in_thread(
        sleep_time=0.1, daemon=True, exception_handler=exception_handler
    )

Instead of immediately triggering the retry mechanism, the process will first break (producing the mentioned error) before retrying automatically.

Furthermore, since the underlying vk_pubsub.get_message method fails the first time, if you remove the exception_handler from your example, the thread will break completely and no messages will be received (despite now having a Retry mechanism).

To better visualize this, consider running a separate thread that manually reads messages:

while True:
    try:
        msg = vk_pubsub.get_message(timeout=0.01, ignore_subscribe_messages=True)

        if msg is None:
            continue

        print(msg)
    except Exception:
        print(traceback.format_exc())

After stopping Valkey, its behavior will be:

1. get_message fails but thanks to the try-except guards, there is one manual retry
2. the automatic retry mechanism kicks in.

If you remove the try-except, the thread will break immediately, which is not the expected behavior.

I say this because, if you use other methods like psubscribe instead of get_message, you don't need a manual retry—everything is already handled internally.

Hope this clarifies things! Let me know if you have any questions.

[mkmkme]

Hey @ManelCoutinhoSensei!

Thanks, this does clarify it. Unfortunately I don't have to handle this right now as I'm going for a 2-weeks vacation. I'll have a look at it right after it.

[ManelCoutinhoSensei]

Hey @mkmkme
Did you have the time to take a look at this in the meantime?

[mkmkme]

Hey @ManelCoutinhoSensei,

Unfortunately I've been quite busy with the main job. Worry not, though, as I haven't forgotten about this! I'll try to handle this in the upcoming days and this is now in my high priority list in valkey-py.

Sorry for the delay!