Repository: stereobutter/reddish (language: python)

Issue:
using `Redis` concurrently may result in `BusyResourceError`
TODO: add lock around the underlying `trio.abc.Stream` instance


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 53e8243..1c82693 100644
--- a/README.md
+++ b/README.md
@@ -17,13 +17,13 @@ barring regular commands from being issued over the connection.
 
 ## Installation
 ```
-pip install reddish
+pip install reddish[trio]  # install with support for trio
 ```
 
 ## Minimal Example
 ```python
 import trio
-from reddish import Redis, Command
+from reddish.trio import Redis, Command
 
 redis = Redis(await trio.open_tcp_stream('localhost', 6379))
 
@@ -63,7 +63,7 @@ assert response == json.loads(data)
 
 ### Command with variadic arguments
 ```python
-from reddish import Args
+from reddish.trio import Args
 
 # inlining arguments
 Command('DEL {keys}', keys=Args(['foo', 'bar']))  # DEL foo bar
@@ -80,7 +80,7 @@ foo, bar = await redis.execute_many(Command('GET', 'foo'), Command('GET', 'bar')
 
 ### Transactions
 ```python
-from reddish import MultiExec
+from reddish.trio import MultiExec
 
 tx = MultiExec(
     Command('ECHO {}', 'foo'),
diff --git a/pyproject.toml b/pyproject.toml
index 8c8e588..f3c981f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,10 +10,13 @@ repository = "https://github.com/stereobutter/reddish"
 
 [tool.poetry.dependencies]
 python = "^3.7"
-trio = "^0.19.0"
+trio = {version="^0.19.0", optional=true}
 pydantic = "^1.8.2"
 hiredis = "^2.0.0"
 
+[tool.poetry.extras]
+trio = ["trio"]
+
 [tool.poetry.dev-dependencies]
 pytest = "^6.2.4"
 hypothesis = "^6.17.1"
diff --git a/reddish/__init__.py b/reddish/__init__.py
index 75fc2c9..e69de29 100644
--- a/reddish/__init__.py
+++ b/reddish/__init__.py
@@ -1,4 +0,0 @@
-import reddish.models  # noqa: F401
-
-from ._command import Args, Command, MultiExec  # noqa: F401
-from ._redis import Redis  # noqa: F401
diff --git a/reddish/_errors.py b/reddish/_errors.py
index dec2441..8beae75 100644
--- a/reddish/_errors.py
+++ b/reddish/_errors.py
@@ -1,4 +1,4 @@
-class ConnectionClosedError(Exception):
+class BrokenConnectionError(Exception):
     def __init__(self, msg='Connection closed.', *args, **kwargs):
         super().__init__(msg, *args, **kwargs)
 
diff --git a/reddish/_redis.py b/reddish/_redis.py
deleted file mode 100644
index 3e6ae08..0000000
--- a/reddish/_redis.py
+++ /dev/null
@@ -1,100 +0,0 @@
-from trio.abc import Stream
-from hiredis import Reader
-
-from typing import Union
-
-from ._command import Command, MultiExec
-from ._utils import partition
-from ._errors import ConnectionClosedError
-
-
-UNSUPPORTED_COMMANDS = (
-    'SUBSCRIBE',
-    'UNSUBSCRIBE'
-    'PSUBSCRIBE',
-    'PUNSUBSCRIBE',
-    'MONITOR'
-)
-
-
-class UnsupportedCommandError(Exception):
-    pass
-
-
-class Redis:
-
-    def __init__(self, stream: Stream):
-        """Redis client for executing commands.
-
-        Args:
-            stream: a `trio.abc.Stream` connected to a redis server.
-        """
-        self._stream = stream
-        self._reader = Reader()
-
-    async def execute(self, command: Union[Command, MultiExec]):
-        """Execute a single redis command.
-
-        Args:
-            command: The command to be executed.
-
-        Returns:
-            Response from redis as received or parsed into the type
-            provided to the command.
-        """
-        [response] = await self.execute_many(command)
-        return response
-
-    async def execute_many(self, *commands: Union[Command, MultiExec]):
-        """Execute multiple redis commands at once.
-
-        Args:
-            *commands: The commands to be executed.
-
-        Returns:
-            Responses from redis as received or parsed into the types
-            provided to the commands.
-        """
-        def guard(command):
-            if cmd._command_name.upper() in UNSUPPORTED_COMMANDS:
-                raise UnsupportedCommandError(f"The '{cmd._command_name}' command is not supported.")
-
-        for cmd in commands:
-            if isinstance(cmd, MultiExec):
-                tx = cmd
-                for cmd in tx:
-                    guard(cmd)
-            else:
-                guard(cmd)
-
-        await self._stream.send_all(b''.join(bytes(cmd) for cmd in commands))
-
-        expected_num_replies = [len(cmd) for cmd in commands]
-        replies = await self._read_reply(expected=sum(expected_num_replies))
-
-        output = tuple(
-            cmd._parse_response(*data)
-            for cmd, data in zip(commands, partition(replies, expected_num_replies))
-        )
-
-        return output
-
-    async def _read_reply(self, expected=1):
-        # Read and parse replies from connection.
-        # ``expected`` is the amount of expected replies. In case of
-        # pipelining this number is set to the amount of commands sent.
-
-        replies = []
-
-        while True:
-            data = await self._stream.receive_some()
-            if data == b'':
-                raise ConnectionClosedError('connection unexpectedly closed')
-            self._reader.feed(data)
-            while True:
-                reply = self._reader.gets()
-                if reply is False:
-                    break  # Read more data, go back to the outer loop.
-                replies.append(reply)
-                if len(replies) == expected:
-                    return replies
diff --git a/reddish/_sansio.py b/reddish/_sansio.py
index f64e5b2..42334dc 100644
--- a/reddish/_sansio.py
+++ b/reddish/_sansio.py
@@ -1,7 +1,7 @@
 import hiredis
 from ._utils import partition
 from ._command import MultiExec
-from ._errors import UnsupportedCommandError
+from ._errors import UnsupportedCommandError, BrokenConnectionError
 
 
 class ReplyBuffer:
@@ -40,12 +40,14 @@ class RedisSansIO:
     def __init__(self, reader=None):
         self._reader = reader or hiredis.Reader()
         self._reply_buffer = None
-        self._closed = False
+        self._broken = False
 
-    def close(self):
-        self._closed = True
+    def mark_broken(self):
+        self._broken = True
 
     def send(self, commands):
+        if self._broken:
+            raise BrokenConnectionError()
         if self._reply_buffer is not None:
             raise ProtocolError('Cannot send more commands')
         for cmd in commands:
@@ -56,6 +58,8 @@ def send(self, commands):
     def receive(self, data):
         reply_buffer = self._reply_buffer
 
+        if self._broken:
+            raise BrokenConnectionError()
         if reply_buffer is None:
             raise ProtocolError('Cannot receive replies because no commands where queued')
         reader = self._reader
diff --git a/reddish/_typing.py b/reddish/_typing.py
new file mode 100644
index 0000000..05cdd77
--- /dev/null
+++ b/reddish/_typing.py
@@ -0,0 +1,4 @@
+from ._command import Command, MultiExec
+from typing import Union
+
+CommandType = Union[Command, MultiExec]
diff --git a/reddish/trio/__init__.py b/reddish/trio/__init__.py
new file mode 100644
index 0000000..1fd4d47
--- /dev/null
+++ b/reddish/trio/__init__.py
@@ -0,0 +1,2 @@
+from ._client import Redis  # noqa
+from .._command import Args, Command, MultiExec  # noqa
\ No newline at end of file
diff --git a/reddish/trio/_client.py b/reddish/trio/_client.py
new file mode 100644
index 0000000..3b618ea
--- /dev/null
+++ b/reddish/trio/_client.py
@@ -0,0 +1,70 @@
+try:
+    import trio
+except ImportError:
+    raise ImportError("Execute 'pip install reddish[trio]' to enable trio support")
+from .._sansio import RedisSansIO
+from .._errors import BrokenConnectionError
+
+from .._typing import CommandType
+
+
+class Redis:
+
+    def __init__(self, stream: trio.abc.Stream) -> None:
+        """Redis client for executing commands.
+
+        Args:
+            stream: a `trio.abc.Stream` connected to a redis server.
+        """
+
+        if not isinstance(stream, trio.abc.Stream):
+            raise TypeError(f"'{repr(stream)}' is not an instance of 'trio.abc.Stream'")
+        self._stream = stream
+        self._lock = trio.Lock()
+        self._redis = RedisSansIO()
+
+    async def execute_many(self, *commands: CommandType):
+        """Execute multiple redis commands at once.
+
+        Args:
+            *commands: The commands to be executed.
+
+        Returns:
+            Responses from redis as received or parsed into the types
+            provided to the commands.
+        """
+
+        redis = self._redis
+        stream = self._stream
+
+        async with self._lock:
+            try:
+                request = redis.send(commands)
+                await stream.send_all(request)
+
+                while True:
+                    data = await stream.receive_some()
+                    if data == b'':
+                        raise BrokenConnectionError()
+                    replies = redis.receive(data)
+                    if replies:
+                        return replies
+            except (trio.BrokenResourceError, trio.ClosedResourceError):
+                redis.mark_broken()
+                raise BrokenConnectionError()
+            except trio.Cancelled:
+                redis.mark_broken()
+                raise
+
+    async def execute(self, command: CommandType):
+        """Execute a single redis command.
+
+        Args:
+            command: The command to be executed.
+
+        Returns:
+            Response from redis as received or parsed into the type
+            provided to the command.
+        """
+
+        return (await self.execute_many(command))[0]


Test patch:
diff --git a/tests/test_RedisSansIO.py b/tests/test_RedisSansIO.py
index 8246411..0d52ccd 100644
--- a/tests/test_RedisSansIO.py
+++ b/tests/test_RedisSansIO.py
@@ -1,7 +1,7 @@
 import pytest
 from reddish._command import Command, MultiExec
 from reddish._sansio import RedisSansIO, ProtocolError
-from reddish._errors import UnsupportedCommandError
+from reddish._errors import UnsupportedCommandError, BrokenConnectionError
 
 
 @pytest.fixture
@@ -9,31 +9,34 @@ def redis():
     return RedisSansIO()
 
 
-def test_sending_command(redis):
-    cmd = Command('PING')
-    assert bytes(cmd) == redis.send([Command('PING')])
+@pytest.fixture
+def ping():
+    return Command('PING').into(str)
+
 
+def test_sending_command(redis, ping):
+    assert bytes(ping) == redis.send([ping])
 
-def test_sending_multiple_commands(redis):
-    cmd = Command('PING')
-    assert 2 * bytes(cmd) == redis.send([Command('PING'), Command('PING')])
 
+def test_sending_multiple_commands(redis, ping):
+    assert 2 * bytes(ping) == redis.send([ping, ping])
 
-def test_receiving(redis):
-    redis.send([Command('PING').into(str)])
+
+def test_receiving(redis, ping):
+    redis.send([ping])
     assert ['PONG'] == redis.receive(b'+PONG\r\n')
 
 
-def test_receiving_in_chunks(redis):
-    redis.send([Command('PING').into(str)])
+def test_receiving_in_chunks(redis, ping):
+    redis.send([ping])
     assert [] == redis.receive(b'+PONG')
     assert ['PONG'] == redis.receive(b'\r\n')
 
 
-def test_sending_repeatedly_should_raise(redis):
+def test_sending_repeatedly_should_raise(redis, ping):
     with pytest.raises(ProtocolError):
-        redis.send([Command('PING')])
-        redis.send([Command('PING')])
+        redis.send([ping])
+        redis.send([ping])
 
 
 def test_receiving_without_send_should_raise(redis):
@@ -41,6 +44,18 @@ def test_receiving_without_send_should_raise(redis):
         redis.receive(b'+PONG')
 
 
+def test_broken_connection_send(redis):
+    redis.mark_broken()
+    with pytest.raises(BrokenConnectionError):
+        redis.send(ping)
+
+
+def test_broken_connection_receive(redis):
+    redis.mark_broken()
+    with pytest.raises(BrokenConnectionError):
+        redis.receive(b'')
+
+
 def test_unsupported_commands(redis):
     with pytest.raises(UnsupportedCommandError):
         redis.send([Command('SUBSCRIBE foo')])
diff --git a/tests/conftest.py b/tests/trio/conftest.py
similarity index 100%
rename from tests/conftest.py
rename to tests/trio/conftest.py
diff --git a/tests/trio/test_trio.py b/tests/trio/test_trio.py
new file mode 100644
index 0000000..e12ddb7
--- /dev/null
+++ b/tests/trio/test_trio.py
@@ -0,0 +1,55 @@
+import trio
+import pytest_trio
+import pytest
+from reddish.trio import Redis, Command
+from reddish._errors import BrokenConnectionError
+
+
+@pytest_trio.trio_fixture
+async def connection():
+    return trio.open_tcp_stream('localhost', 6379)
+
+
+@pytest_trio.trio_fixture
+async def redis(connection):
+    return Redis(await connection)
+
+
+@pytest.fixture
+def ping():
+    return Command('PING').into(str)
+
+
+async def test_execute(redis, ping):
+    'PONG' == await redis.execute(ping)
+
+
+async def test_execute_many(redis, ping):
+    ['PONG', 'PONG'] == await redis.execute_many(ping, ping)
+
+
+async def test_stream(connection):
+    with pytest.raises(TypeError):
+        Redis(connection)
+
+
+async def test_closed_connection(connection, ping):
+    async with await connection as stream:
+        redis = Redis(stream)
+
+    with pytest.raises(BrokenConnectionError):
+        await redis.execute(ping)
+
+
+async def test_concurrent_requests(redis, ping):
+    async with trio.open_nursery() as nursery:
+        for i in range(10):
+            nursery.start_soon(redis.execute, ping)
+
+
+async def test_cancellation(redis, ping):
+    with trio.move_on_after(0.0001):  # break the connection
+        await redis.execute(ping)
+
+    with pytest.raises(BrokenConnectionError):
+        await redis.execute(ping)