Merge pull request #324 from martinnj/feature/retrying-client

Implement RetryingClient

3 files changed, 502 insertions, 0 deletions

diff --git a/pymemcache/client/__init__.py b/pymemcache/client/__init__.py
index 158cf15..250f4e1 100644
--- a/pymemcache/client/__init__.py
+++ b/pymemcache/client/__init__.py
@@ -3,6 +3,7 @@
 from pymemcache.client.base import Client  # noqa
 from pymemcache.client.base import PooledClient  # noqa
 from pymemcache.client.hash import HashClient  # noqa
+from pymemcache.client.retrying import RetryingClient  # noqa
 
 from pymemcache.exceptions import MemcacheError  # noqa
 from pymemcache.exceptions import MemcacheClientError  # noqa
diff --git a/pymemcache/client/retrying.py b/pymemcache/client/retrying.py
new file mode 100644
index 0000000..8c8f3a9
--- /dev/null
+++ b/pymemcache/client/retrying.py
@@ -0,0 +1,185 @@
+""" Module containing the RetryingClient wrapper class. """
+
+from time import sleep
+
+
+def _ensure_tuple_argument(argument_name, argument_value):
+    """
+    Helper function to ensure the given arguments are tuples of Exceptions (or
+    subclasses), or can at least be converted to such.
+
+    Args:
+      argument_name: str, name of the argument we're checking, only used for
+        raising meaningful exceptions.
+      argument: any, the argument itself.
+
+    Returns:
+      tuple[Exception]: A tuple with the elements from the argument if they are
+        valid.
+
+    Exceptions:
+      ValueError: If the argument was not None, tuple or Iterable.
+      ValueError: If any of the elements of the argument is not a subclass of
+        Exception.
+    """
+
+    # Ensure the argument is a tuple, set or list.
+    if argument_value is None:
+        return tuple()
+    elif not isinstance(argument_value, (tuple, set, list)):
+        raise ValueError(
+            "%s must be either a tuple, a set or a list." % argument_name
+        )
+
+    # Convert the argument before checking contents.
+    argument_tuple = tuple(argument_value)
+
+    # Check that all the elements are actually inherited from Exception.
+    # (Catchable)
+    if not all([issubclass(arg, Exception) for arg in argument_tuple]):
+        raise ValueError(
+            "%s is only allowed to contain elements that are subclasses of "
+            "Exception." % argument_name
+        )
+
+    return argument_tuple
+
+
+class RetryingClient(object):
+    """
+    Client that allows retrying calls for the other clients.
+    """
+
+    def __init__(
+        self,
+        client,
+        attempts=2,
+        retry_delay=0,
+        retry_for=None,
+        do_not_retry_for=None
+    ):
+        """
+        Constructor for RetryingClient.
+
+        Args:
+          client: Client|PooledClient|HashClient, inner client to use for
+            performing actual work.
+          attempts: optional int, how many times to attempt an action before
+            failing. Must be 1 or above. Defaults to 2.
+          retry_delay: optional int|float, how many seconds to sleep between
+            each attempt.
+            Defaults to 0.
+
+          retry_for: optional None|tuple|set|list, what exceptions to
+            allow retries for. Will allow retries for all exceptions if None.
+            Example:
+                `(MemcacheClientError, MemcacheUnexpectedCloseError)`
+            Accepts any class that is a subclass of Exception.
+            Defaults to None.
+
+          do_not_retry_for: optional None|tuple|set|list, what
+            exceptions should be retried. Will not block retries for any
+            Exception if None.
+            Example:
+                `(IOError, MemcacheIllegalInputError)`
+            Accepts any class that is a subclass of Exception.
+            Defaults to None.
+
+        Exceptions:
+          ValueError: If `attempts` is not 1 or above.
+          ValueError: If `retry_for` or `do_not_retry_for` is not None, tuple or
+            Iterable.
+          ValueError: If any of the elements of `retry_for` or
+            `do_not_retry_for` is not a subclass of Exception.
+          ValueError: If there is any overlap between `retry_for` and
+            `do_not_retry_for`.
+        """
+
+        if attempts < 1:
+            raise ValueError(
+                "`attempts` argument must be at least 1. "
+                "Otherwise no attempts are made."
+            )
+
+        self._client = client
+        self._attempts = attempts
+        self._retry_delay = retry_delay
+        self._retry_for = _ensure_tuple_argument("retry_for", retry_for)
+        self._do_not_retry_for = _ensure_tuple_argument(
+            "do_not_retry_for", do_not_retry_for
+        )
+
+        # Verify no overlap in the go/no-go exception collections.
+        for exc_class in self._retry_for:
+            if exc_class in self._do_not_retry_for:
+                raise ValueError(
+                    "Exception class \"%s\" was present in both `retry_for` "
+                    "and `do_not_retry_for`. Any exception class is only "
+                    "allowed in a single argument." % repr(exc_class)
+                )
+
+        # Take dir from the client to speed up future checks.
+        self._client_dir = dir(self._client)
+
+    def _retry(self, name, func, *args, **kwargs):
+        """
+        Workhorse function, handles retry logic.
+
+        Args:
+          name: str, Name of the function called.
+          func: callable, the function to retry.
+          *args: args, array arguments to pass to the function.
+          **kwargs: kwargs, keyword arguments to pass to the function.
+        """
+        for attempt in range(self._attempts):
+            try:
+                result = func(*args, **kwargs)
+                return result
+
+            except Exception as exc:
+                # Raise the exception to caller if either is met:
+                # - We've used the last attempt.
+                # - self._retry_for is set, and we do not match.
+                # - self._do_not_retry_for is set, and we do match.
+                # - name is not actually a member of the client class.
+                if attempt >= self._attempts - 1 \
+                   or (self._retry_for
+                       and not isinstance(exc, self._retry_for)) \
+                   or (self._do_not_retry_for
+                       and isinstance(exc, self._do_not_retry_for)) \
+                   or name not in self._client_dir:
+                    raise exc
+
+                # Sleep and try again.
+                sleep(self._retry_delay)
+
+    # This is the real magic soup of the class, we catch anything that isn't
+    # strictly defined for ourselves and pass it on to whatever client we've
+    # been given.
+    def __getattr__(self, name):
+
+        return lambda *args, **kwargs: self._retry(
+            name,
+            self._client.__getattribute__(name),
+            *args,
+            **kwargs
+        )
+
+    # We implement these explicitly because they're "magic" functions and won't
+    # get passed on by __getattr__.
+
+    def __dir__(self):
+        return self._client_dir
+
+    # These magics are copied from the base client.
+    def __setitem__(self, key, value):
+        self.set(key, value, noreply=True)
+
+    def __getitem__(self, key):
+        value = self.get(key)
+        if value is None:
+            raise KeyError
+        return value
+
+    def __delitem__(self, key):
+        self.delete(key, noreply=True)
diff --git a/pymemcache/test/test_client_retry.py b/pymemcache/test/test_client_retry.py
new file mode 100644
index 0000000..230a941
--- /dev/null
+++ b/pymemcache/test/test_client_retry.py
@@ -0,0 +1,316 @@
+""" Test collection for the RetryingClient. """
+
+import functools
+import unittest
+
+import mock
+import pytest
+
+from .test_client import ClientTestMixin, MockSocket
+from pymemcache.client.retrying import RetryingClient
+from pymemcache.client.base import Client
+from pymemcache.exceptions import MemcacheUnknownError, MemcacheClientError
+
+
+# Test pure passthroughs with no retry action.
+class TestRetryingClientPassthrough(ClientTestMixin, unittest.TestCase):
+
+    def make_base_client(self, mock_socket_values, **kwargs):
+        base_client = Client(None, **kwargs)
+        # mock out client._connect() rather than hard-settting client.sock to
+        # ensure methods are checking whether self.sock is None before
+        # attempting to use it
+        sock = MockSocket(list(mock_socket_values))
+        base_client._connect = mock.Mock(side_effect=functools.partial(
+            setattr, base_client, "sock", sock))
+        return base_client
+
+    def make_client(self, mock_socket_values, **kwargs):
+        # Create a base client to wrap.
+        base_client = self.make_base_client(
+            mock_socket_values=mock_socket_values, **kwargs
+        )
+
+        # Wrap the client in the retrying class, disable retries.
+        client = RetryingClient(base_client, attempts=1)
+        return client
+
+
+# Retry specific tests.
+@pytest.mark.unit()
+class TestRetryingClient(object):
+
+    def make_base_client(self, mock_socket_values, **kwargs):
+        """ Creates a regular mock client to wrap in the RetryClient. """
+        base_client = Client(None, **kwargs)
+        # mock out client._connect() rather than hard-settting client.sock to
+        # ensure methods are checking whether self.sock is None before
+        # attempting to use it
+        sock = MockSocket(list(mock_socket_values))
+        base_client._connect = mock.Mock(side_effect=functools.partial(
+            setattr, base_client, "sock", sock))
+        return base_client
+
+    def make_client(self, mock_socket_values, **kwargs):
+        """
+        Creates a RetryingClient that will respond with the given values,
+        configured using kwargs.
+        """
+        # Create a base client to wrap.
+        base_client = self.make_base_client(
+            mock_socket_values=mock_socket_values
+        )
+
+        # Wrap the client in the retrying class, and pass kwargs on.
+        client = RetryingClient(base_client, **kwargs)
+        return client
+
+    # Start testing.
+    def test_constructor_default(self):
+        base_client = self.make_base_client([])
+        RetryingClient(base_client)
+
+        with pytest.raises(TypeError):
+            RetryingClient()
+
+    def test_constructor_attempts(self):
+        base_client = self.make_base_client([])
+        rc = RetryingClient(base_client, attempts=1)
+        assert rc._attempts == 1
+
+        with pytest.raises(ValueError):
+            RetryingClient(base_client, attempts=0)
+
+    def test_constructor_retry_for(self):
+        base_client = self.make_base_client([])
+
+        # Try none/default.
+        rc = RetryingClient(base_client, retry_for=None)
+        assert rc._retry_for == tuple()
+
+        # Try with tuple.
+        rc = RetryingClient(base_client, retry_for=tuple([Exception]))
+        assert rc._retry_for == tuple([Exception])
+
+        # Try with list.
+        rc = RetryingClient(base_client, retry_for=[Exception])
+        assert rc._retry_for == tuple([Exception])
+
+        # Try with multi element list.
+        rc = RetryingClient(base_client, retry_for=[Exception, IOError])
+        assert rc._retry_for == (Exception, IOError)
+
+        # With string?
+        with pytest.raises(ValueError):
+            RetryingClient(base_client, retry_for="haha!")
+
+        # With collectino of string and exceptions?
+        with pytest.raises(ValueError):
+            RetryingClient(base_client, retry_for=[Exception, str])
+
+    def test_constructor_do_no_retry_for(self):
+        base_client = self.make_base_client([])
+
+        # Try none/default.
+        rc = RetryingClient(base_client, do_not_retry_for=None)
+        assert rc._do_not_retry_for == tuple()
+
+        # Try with tuple.
+        rc = RetryingClient(base_client, do_not_retry_for=tuple([Exception]))
+        assert rc._do_not_retry_for == tuple([Exception])
+
+        # Try with list.
+        rc = RetryingClient(base_client, do_not_retry_for=[Exception])
+        assert rc._do_not_retry_for == tuple([Exception])
+
+        # Try with multi element list.
+        rc = RetryingClient(base_client, do_not_retry_for=[Exception, IOError])
+        assert rc._do_not_retry_for == (Exception, IOError)
+
+        # With string?
+        with pytest.raises(ValueError):
+            RetryingClient(base_client, do_not_retry_for="haha!")
+
+        # With collectino of string and exceptions?
+        with pytest.raises(ValueError):
+            RetryingClient(base_client, do_not_retry_for=[Exception, str])
+
+    def test_constructor_both_filters(self):
+        base_client = self.make_base_client([])
+
+        # Try none/default.
+        rc = RetryingClient(base_client, retry_for=None, do_not_retry_for=None)
+        assert rc._retry_for == tuple()
+        assert rc._do_not_retry_for == tuple()
+
+        # Try a valid config.
+        rc = RetryingClient(
+            base_client,
+            retry_for=[Exception, IOError],
+            do_not_retry_for=[ValueError, MemcacheUnknownError]
+        )
+        assert rc._retry_for == (Exception, IOError)
+        assert rc._do_not_retry_for == (ValueError, MemcacheUnknownError)
+
+        # Try with overlapping filters
+        with pytest.raises(ValueError):
+            rc = RetryingClient(
+                base_client,
+                retry_for=[Exception, IOError, MemcacheUnknownError],
+                do_not_retry_for=[ValueError, MemcacheUnknownError]
+            )
+
+    def test_dir_passthrough(self):
+        base = self.make_base_client([])
+        client = RetryingClient(base)
+
+        assert dir(base) == dir(client)
+
+    def test_retry_dict_set_is_supported(self):
+        client = self.make_client([b'UNKNOWN\r\n', b'STORED\r\n'])
+        client[b'key'] = b'value'
+
+    def test_retry_dict_get_is_supported(self):
+        client = self.make_client(
+            [
+                b'UNKNOWN\r\n',
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+            ]
+        )
+        assert client[b'key'] == b'value'
+
+    def test_retry_dict_get_not_found_is_supported(self):
+        client = self.make_client([b'UNKNOWN\r\n', b'END\r\n'])
+
+        with pytest.raises(KeyError):
+            client[b'key']
+
+    def test_retry_dict_del_is_supported(self):
+        client = self.make_client([b'UNKNOWN\r\n', b'DELETED\r\n'])
+        del client[b'key']
+
+    def test_retry_get_found(self):
+        client = self.make_client([
+            b'UNKNOWN\r\n',
+            b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+        ], attempts=2)
+        result = client.get("key")
+        assert result == b'value'
+
+    def test_retry_get_not_found(self):
+        client = self.make_client([
+            b'UNKNOWN\r\n',
+            b'END\r\n'
+        ], attempts=2)
+        result = client.get("key")
+        assert result is None
+
+    def test_retry_get_exception(self):
+        client = self.make_client([
+            b'UNKNOWN\r\n',
+            b'UNKNOWN\r\n'
+        ], attempts=2)
+        with pytest.raises(MemcacheUnknownError):
+            client.get("key")
+
+    def test_retry_set_success(self):
+        client = self.make_client([
+            b'UNKNOWN\r\n',
+            b'STORED\r\n'
+        ], attempts=2)
+        result = client.set("key", "value", noreply=False)
+        assert result is True
+
+    def test_retry_set_fail(self):
+        client = self.make_client([
+            b'UNKNOWN\r\n',
+            b'UNKNOWN\r\n',
+            b'STORED\r\n'
+        ], attempts=2)
+        with pytest.raises(MemcacheUnknownError):
+            client.set("key", "value", noreply=False)
+
+    def test_no_retry(self):
+        client = self.make_client([
+            b'UNKNOWN\r\n',
+            b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+        ], attempts=1)
+
+        with pytest.raises(MemcacheUnknownError):
+            client.get("key")
+
+    def test_retry_for_exception_success(self):
+        # Test that we retry for the exception specified.
+        client = self.make_client(
+            [
+                MemcacheClientError("Whoops."),
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+            ],
+            attempts=2,
+            retry_for=tuple([MemcacheClientError])
+        )
+        result = client.get("key")
+        assert result == b'value'
+
+    def test_retry_for_exception_fail(self):
+        # Test that we do not retry for unapproved exception.
+        client = self.make_client(
+            [
+                MemcacheUnknownError("Whoops."),
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+            ],
+            attempts=2,
+            retry_for=tuple([MemcacheClientError])
+        )
+
+        with pytest.raises(MemcacheUnknownError):
+            client.get("key")
+
+    def test_do_not_retry_for_exception_success(self):
+        # Test that we retry for exceptions not specified.
+        client = self.make_client(
+            [
+                MemcacheClientError("Whoops."),
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+            ],
+            attempts=2,
+            do_not_retry_for=tuple([MemcacheUnknownError])
+        )
+        result = client.get("key")
+        assert result == b'value'
+
+    def test_do_not_retry_for_exception_fail(self):
+        # Test that we do not retry for the exception specified.
+        client = self.make_client(
+            [
+                MemcacheClientError("Whoops."),
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n'
+            ],
+            attempts=2,
+            do_not_retry_for=tuple([MemcacheClientError])
+        )
+
+        with pytest.raises(MemcacheClientError):
+            client.get("key")
+
+    def test_both_exception_filters(self):
+        # Test interacction between both exception filters.
+        client = self.make_client(
+            [
+                MemcacheClientError("Whoops."),
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n',
+                MemcacheUnknownError("Whoops."),
+                b'VALUE key 0 5\r\nvalue\r\nEND\r\n',
+            ],
+            attempts=2,
+            retry_for=tuple([MemcacheClientError]),
+            do_not_retry_for=tuple([MemcacheUnknownError])
+        )
+
+        # Check that we succeed where allowed.
+        result = client.get("key")
+        assert result == b'value'
+
+        # Check that no retries are attempted for the banned exception.
+        with pytest.raises(MemcacheUnknownError):
+            client.get("key")