undo prevision fix attempts in async client and cluster

Author: kristjanvalur

redis/asyncio/client.py

@@ -500,23 +500,6 @@ async def _disconnect_raise(self, conn: Connection, error: Exception):
         ):
             raise error
 
-    async def _try_send_command_parse_response(self, conn, *args, **options):
-        try:
-            return await conn.retry.call_with_retry(
-                lambda: self._send_command_parse_response(
-                    conn, args[0], *args, **options
-                ),
-                lambda error: self._disconnect_raise(conn, error),
-            )
-        except asyncio.CancelledError:
-            await conn.disconnect(nowait=True)
-            raise
-        finally:
-            if self.single_connection_client:
-                self._single_conn_lock.release()
-            if not self.connection:
-                await self.connection_pool.release(conn)
-
     # COMMAND EXECUTION AND PROTOCOL PARSING
     async def execute_command(self, *args, **options):
         """Execute a command and return a parsed response"""
@@ -527,10 +510,18 @@ async def execute_command(self, *args, **options):
 
         if self.single_connection_client:
             await self._single_conn_lock.acquire()
-
-        return await asyncio.shield(
-            self._try_send_command_parse_response(conn, *args, **options)
-        )
+        try:
+            return await conn.retry.call_with_retry(
+                lambda: self._send_command_parse_response(
+                    conn, command_name, *args, **options
+                ),
+                lambda error: self._disconnect_raise(conn, error),
+            )
+        finally:
+            if self.single_connection_client:
+                self._single_conn_lock.release()
+            if not self.connection:
+                await pool.release(conn)
 
     async def parse_response(
         self, connection: Connection, command_name: Union[str, bytes], **options
@@ -774,18 +765,10 @@ async def _disconnect_raise_connect(self, conn, error):
         is not a TimeoutError. Otherwise, try to reconnect
         """
         await conn.disconnect()
-
         if not (conn.retry_on_timeout and isinstance(error, TimeoutError)):
             raise error
         await conn.connect()
 
-    async def _try_execute(self, conn, command, *arg, **kwargs):
-        try:
-            return await command(*arg, **kwargs)
-        except asyncio.CancelledError:
-            await conn.disconnect()
-            raise
-
     async def _execute(self, conn, command, *args, **kwargs):
         """
         Connect manually upon disconnection. If the Redis server is down,
@@ -794,11 +777,9 @@ async def _execute(self, conn, command, *args, **kwargs):
         called by the # connection to resubscribe us to any channels and
         patterns we were previously listening to
         """
-        return await asyncio.shield(
-            conn.retry.call_with_retry(
-                lambda: self._try_execute(conn, command, *args, **kwargs),
-                lambda error: self._disconnect_raise_connect(conn, error),
-            )
+        return await conn.retry.call_with_retry(
+            lambda: command(*args, **kwargs),
+            lambda error: self._disconnect_raise_connect(conn, error),
         )
 
     async def parse_response(self, block: bool = True, timeout: float = 0):
@@ -1202,18 +1183,6 @@ async def _disconnect_reset_raise(self, conn, error):
             await self.reset()
             raise
 
-    async def _try_send_command_parse_response(self, conn, *args, **options):
-        try:
-            return await conn.retry.call_with_retry(
-                lambda: self._send_command_parse_response(
-                    conn, args[0], *args, **options
-                ),
-                lambda error: self._disconnect_reset_raise(conn, error),
-            )
-        except asyncio.CancelledError:
-            await conn.disconnect()
-            raise
-
     async def immediate_execute_command(self, *args, **options):
         """
         Execute a command immediately, but don't auto-retry on a
@@ -1229,8 +1198,12 @@ async def immediate_execute_command(self, *args, **options):
                 command_name, self.shard_hint
             )
             self.connection = conn
-        return await asyncio.shield(
-            self._try_send_command_parse_response(conn, *args, **options)
+
+        return await conn.retry.call_with_retry(
+            lambda: self._send_command_parse_response(
+                conn, command_name, *args, **options
+            ),
+            lambda error: self._disconnect_reset_raise(conn, error),
         )
 
     def pipeline_execute_command(self, *args, **options):
@@ -1398,19 +1371,6 @@ async def _disconnect_raise_reset(self, conn: Connection, error: Exception):
             await self.reset()
             raise
 
-    async def _try_execute(self, conn, execute, stack, raise_on_error):
-        try:
-            return await conn.retry.call_with_retry(
-                lambda: execute(conn, stack, raise_on_error),
-                lambda error: self._disconnect_raise_reset(conn, error),
-            )
-        except asyncio.CancelledError:
-            # not supposed to be possible, yet here we are
-            await conn.disconnect(nowait=True)
-            raise
-        finally:
-            await self.reset()
-
     async def execute(self, raise_on_error: bool = True):
         """Execute all the commands in the current pipeline"""
         stack = self.command_stack
@@ -1432,11 +1392,10 @@ async def execute(self, raise_on_error: bool = True):
         conn = cast(Connection, conn)
 
         try:
-            return await asyncio.shield(
-                self._try_execute(conn, execute, stack, raise_on_error)
+            return await conn.retry.call_with_retry(
+                lambda: execute(conn, stack, raise_on_error),
+                lambda error: self._disconnect_raise_reset(conn, error),
             )
-        except RuntimeError:
-            await self.reset()
         finally:
             await self.reset()
 

redis/asyncio/cluster.py

@@ -1016,33 +1016,12 @@ async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
         await connection.send_packed_command(connection.pack_command(*args), False)
 
         # Read response
-        return await asyncio.shield(
-            self._parse_and_release(connection, args[0], **kwargs)
-        )
-
-    async def _parse_and_release(self, connection, *args, **kwargs):
         try:
-            return await self.parse_response(connection, *args, **kwargs)
-        except asyncio.CancelledError:
-            # should not be possible
-            await connection.disconnect(nowait=True)
-            raise
+            return await self.parse_response(connection, args[0], **kwargs)
         finally:
+            # Release connection
             self._free.append(connection)
 
-    async def _try_parse_response(self, cmd, connection, ret):
-        try:
-            cmd.result = await asyncio.shield(
-                self.parse_response(connection, cmd.args[0], **cmd.kwargs)
-            )
-        except asyncio.CancelledError:
-            await connection.disconnect(nowait=True)
-            raise
-        except Exception as e:
-            cmd.result = e
-            ret = True
-        return ret
-
     async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
         # Acquire connection
         connection = self.acquire_connection()
@@ -1055,7 +1034,13 @@ async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
         # Read responses
         ret = False
         for cmd in commands:
-            ret = await asyncio.shield(self._try_parse_response(cmd, connection, ret))
+            try:
+                cmd.result = await self.parse_response(
+                    connection, cmd.args[0], **cmd.kwargs
+                )
+            except Exception as e:
+                cmd.result = e
+                ret = True
 
         # Release connection
         self._free.append(connection)