Add in safety checks and better connection cleanup in certain situations when executing commands

Author: Grokzen

rediscluster/client.py

@@ -668,6 +668,7 @@ def _execute_command(self, *args, **kwargs):
                     )
             except TimeoutError:
                 log.exception("TimeoutError")
+                connection.disconnect()
 
                 if ttl < self.RedisClusterRequestTTL / 2:
                     time.sleep(0.05)
@@ -679,6 +680,7 @@ def _execute_command(self, *args, **kwargs):
                 self.connection_pool.disconnect()
                 self.connection_pool.reset()
                 self.refresh_table_asap = True
+                connection = None
 
                 raise e
             except MovedError as e:
@@ -702,6 +704,10 @@ def _execute_command(self, *args, **kwargs):
                 log.exception("AskError")
 
                 redirect_addr, asking = "{0}:{1}".format(e.host, e.port), True
+            except BaseException as e:
+                log.exception("BaseException")
+                connection.disconnect()
+                raise e
             finally:
                 if connection is not None:
                     self.connection_pool.release(connection)

rediscluster/connection.py

@@ -234,6 +234,27 @@ def get_connection(self, command_name, *keys, **options):
 
         self._in_use_connections[node['name']].add(connection)
 
+        try:
+            # ensure this connection is connected to Redis
+            connection.connect()
+            # connections that the pool provides should be ready to send
+            # a command. if not, the connection was either returned to the
+            # pool before all data has been read or the socket has been
+            # closed. either way, reconnect and verify everything is good.
+            try:
+                if connection.can_read():
+                    raise ConnectionError('Connection has data')
+            except ConnectionError:
+                connection.disconnect()
+                connection.connect()
+                if connection.can_read():
+                    raise ConnectionError('Connection not ready')
+        except BaseException:
+            # release the connection back to the pool so that we don't
+            # leak it
+            self.release(connection)
+            raise
+
         return connection
 
     def make_connection(self, node):