Merge pull request #504 from basho/features/lrb/retry-on-closed-gh-503

Smart re-try for operations

Author: lukebakken

Makefile

@@ -88,19 +88,19 @@ endif
 
 .PHONY: unit-test
 unit-test:
-	@$(PROJDIR)/runner unit-test
+	@$(PROJDIR)/.runner unit-test
 
 .PHONY: integration-test
 integration-test:
-	@$(PROJDIR)/runner integration-test
+	@$(PROJDIR)/.runner integration-test
 
 .PHONY: security-test
 security-test:
-	@$(PROJDIR)/runner security-test
+	@$(PROJDIR)/.runner security-test
 
 .PHONY: timeseries-test
 timeseries-test:
-	@$(PROJDIR)/runner timeseries-test
+	@$(PROJDIR)/.runner timeseries-test
 
 .PHONY: test
 test: integration-test

docs/advanced.rst

@@ -12,14 +12,17 @@ Connection pool
 
 .. currentmodule:: riak.transports.pool
 
-.. autoexception:: BadResource
 .. autoclass:: Resource
    :members:
+
 .. autoclass:: Pool
    :members:
 
 .. autoclass:: PoolIterator
 
+.. autoexception:: BadResource
+.. autoexception:: ConnectionClosed
+
 -----------
 Retry logic
 -----------

riak/client/operations.py

@@ -87,24 +87,21 @@ def stream_buckets(self, bucket_type=None, timeout=None):
 
         """
         _validate_timeout(timeout)
+
         if bucket_type:
             bucketfn = self._bucket_type_bucket_builder
         else:
             bucketfn = self._default_type_bucket_builder
 
-        resource = self._acquire()
-        transport = resource.object
-        stream = transport.stream_buckets(bucket_type=bucket_type,
-                                          timeout=timeout)
-        stream.attach(resource)
-        try:
-            for bucket_list in stream:
-                bucket_list = [bucketfn(bytes_to_str(name), bucket_type)
-                               for name in bucket_list]
-                if len(bucket_list) > 0:
-                    yield bucket_list
-        finally:
-            stream.close()
+        def make_op(transport):
+            return transport.stream_buckets(
+                    bucket_type=bucket_type, timeout=timeout)
+
+        for bucket_list in self._stream_with_retry(make_op):
+            bucket_list = [bucketfn(bytes_to_str(name), bucket_type)
+                           for name in bucket_list]
+            if len(bucket_list) > 0:
+                yield bucket_list
 
     @retryable
     def ping(self, transport):
@@ -266,6 +263,8 @@ def stream_index(self, bucket, index, startkey, endkey=None,
         :rtype: :class:`~riak.client.index_page.IndexPage`
 
         """
+        # TODO FUTURE: implement "retry on connection closed"
+        # as in stream_mapred
         _validate_timeout(timeout, infinity_ok=True)
 
         page = IndexPage(self, bucket, index, startkey, endkey,
@@ -338,6 +337,8 @@ def paginate_stream_index(self, bucket, index, startkey, endkey=None,
           :class:`~riak.client.index_page.IndexPage`
 
         """
+        # TODO FUTURE: implement "retry on connection closed"
+        # as in stream_mapred
         page = self.stream_index(bucket, index, startkey,
                                  endkey=endkey,
                                  max_results=max_results,
@@ -489,19 +490,16 @@ def stream_keys(self, bucket, timeout=None):
         :rtype: iterator
         """
         _validate_timeout(timeout)
-        resource = self._acquire()
-        transport = resource.object
-        stream = transport.stream_keys(bucket, timeout=timeout)
-        stream.attach(resource)
-        try:
-            for keylist in stream:
-                if len(keylist) > 0:
-                    if six.PY2:
-                        yield keylist
-                    else:
-                        yield [bytes_to_str(item) for item in keylist]
-        finally:
-            stream.close()
+
+        def make_op(transport):
+            return transport.stream_keys(bucket, timeout=timeout)
+
+        for keylist in self._stream_with_retry(make_op):
+            if len(keylist) > 0:
+                if six.PY2:
+                    yield keylist
+                else:
+                    yield [bytes_to_str(item) for item in keylist]
 
     @retryable
     def put(self, transport, robj, w=None, dw=None, pw=None, return_body=None,
@@ -799,15 +797,12 @@ def stream_mapred(self, inputs, query, timeout):
         :rtype: iterator
         """
         _validate_timeout(timeout)
-        resource = self._acquire()
-        transport = resource.object
-        stream = transport.stream_mapred(inputs, query, timeout)
-        stream.attach(resource)
-        try:
-            for phase, data in stream:
-                yield phase, data
-        finally:
-            stream.close()
+
+        def make_op(transport):
+            return transport.stream_mapred(inputs, query, timeout)
+
+        for phase, data in self._stream_with_retry(make_op):
+            yield phase, data
 
     @retryable
     def create_search_index(self, transport, index, schema=None, n_val=None,

riak/client/transport.py

@@ -1,9 +1,11 @@
 from contextlib import contextmanager
-from riak.transports.pool import BadResource
+from riak.transports.pool import BadResource, ConnectionClosed
 from riak.transports.tcp import is_retryable as is_tcp_retryable
 from riak.transports.http import is_retryable as is_http_retryable
-import threading
 from six import PY2
+
+import threading
+
 if PY2:
     from httplib import HTTPException
 else:
@@ -84,7 +86,8 @@ def _transport(self):
         _transport()
 
         Yields a single transport to the caller from the default pool,
-        without retries.
+        without retries. NB: no need to re-try as this method is only
+        used by CRDT operations that should never be re-tried.
         """
         pool = self._choose_pool()
         with pool.transaction() as transport:
@@ -98,6 +101,29 @@ def _acquire(self):
         """
         return self._choose_pool().acquire()
 
+    def _stream_with_retry(self, make_op):
+        first_try = True
+        while True:
+            resource = self._acquire()
+            transport = resource.object
+            streaming_op = make_op(transport)
+            streaming_op.attach(resource)
+            try:
+                for item in streaming_op:
+                    yield item
+                break
+            except BadResource as e:
+                resource.errored = True
+                # NB: *only* re-try if connection closed happened
+                # at the start of the streaming op
+                if first_try and not e.mid_stream:
+                    continue
+                else:
+                    raise
+            finally:
+                first_try = False
+                streaming_op.close()
+
     def _with_retries(self, pool, fn):
         """
         Performs the passed function with retries against the given pool.
@@ -112,26 +138,38 @@ def _with_retries(self, pool, fn):
         def _skip_bad_nodes(transport):
             return transport._node not in skip_nodes
 
-        retry_count = self.retries
-
-        for retry in range(retry_count):
+        retry_count = self.retries - 1
+        first_try = True
+        current_try = 0
+        while True:
             try:
-                with pool.transaction(_filter=_skip_bad_nodes) as transport:
+                with pool.transaction(
+                        _filter=_skip_bad_nodes,
+                        yield_resource=True) as resource:
+                    transport = resource.object
                     try:
                         return fn(transport)
-                    except (IOError, HTTPException) as e:
+                    except (IOError, HTTPException, ConnectionClosed) as e:
+                        resource.errored = True
                         if _is_retryable(e):
                             transport._node.error_rate.incr(1)
                             skip_nodes.append(transport._node)
-                            raise BadResource(e)
+                            if first_try:
+                                continue
+                            else:
+                                raise BadResource(e)
                         else:
                             raise
             except BadResource as e:
-                if retry < (retry_count - 1):
+                if current_try < retry_count:
+                    resource.errored = True
+                    current_try += 1
                     continue
                 else:
                     # Re-raise the inner exception
                     raise e.args[0]
+            finally:
+                first_try = False
 
     def _choose_pool(self, protocol=None):
         """
@@ -168,6 +206,7 @@ def _is_retryable(error):
     return is_tcp_retryable(error) or is_http_retryable(error)
 
 
+# http://thecodeship.com/patterns/guide-to-python-function-decorators/
 def retryable(fn, protocol=None):
     """
     Wraps a client operation that can be retried according to the set

riak/transports/pool.py

@@ -5,14 +5,31 @@
 from contextlib import contextmanager
 
 
-# This file is a rough port of the Innertube Ruby library
 class BadResource(Exception):
     """
     Users of a :class:`Pool` should raise this error when the pool
     resource currently in-use is bad and should be removed from the
     pool.
+
+    :param mid_stream: did this exception happen mid-streaming op?
+    :type mid_stream: boolean
+    """
+    def __init__(self, ex, mid_stream=False):
+        super(BadResource, self).__init__(ex)
+        self.mid_stream = mid_stream
+
+
+class ConnectionClosed(BadResource):
     """
-    pass
+    Users of a :class:`Pool` should raise this error when the pool
+    resource currently in-use has been closed and should be removed
+    from the pool.
+
+    :param mid_stream: did this exception happen mid-streaming op?
+    :type mid_stream: boolean
+    """
+    def __init__(self, ex, mid_stream=False):
+        super(ConnectionClosed, self).__init__(ex, mid_stream)
 
 
 class Resource(object):
@@ -30,20 +47,26 @@ def __init__(self, obj, pool):
         :type obj: object
         """
 
-        self.object = obj
         """The wrapped pool resource."""
+        self.object = obj
 
-        self.claimed = False
         """Whether the resource is currently in use."""
+        self.claimed = False
 
-        self.pool = pool
         """The pool that this resource belongs to."""
+        self.pool = pool
+
+        """True if this Resource errored."""
+        self.errored = False
 
     def release(self):
         """
         Releases this resource back to the pool it came from.
         """
-        self.pool.release(self)
+        if self.errored:
+            self.pool.delete_resource(self)
+        else:
+            self.pool.release(self)
 
 
 class Pool(object):
@@ -60,7 +83,7 @@ class Pool(object):
 
     Example::
 
-        from riak.Pool import Pool, BadResource
+        from riak.transports.pool import Pool
         class ListPool(Pool):
             def create_resource(self):
                 return []
@@ -74,7 +97,6 @@ def destroy_resource(self):
             resource.append(1)
         with pool.transaction() as resource2:
             print(repr(resource2)) # should be [1]
-
     """
 
     def __init__(self):
@@ -138,7 +160,7 @@ def release(self, resource):
             self.releaser.notify_all()
 
     @contextmanager
-    def transaction(self, _filter=None, default=None):
+    def transaction(self, _filter=None, default=None, yield_resource=False):
         """
         transaction(_filter=None, default=None)
 
@@ -152,10 +174,18 @@ def transaction(self, _filter=None, default=None):
         :type _filter: callable
         :param default: a value that will be used instead of calling
             :meth:`create_resource` if a new resource needs to be created
+        :param yield_resource: set to True to yield the Resource object
+            itself
+        :type yield_resource: boolean
         """
         resource = self.acquire(_filter=_filter, default=default)
         try:
-            yield resource.object
+            if yield_resource:
+                yield resource
+            else:
+                yield resource.object
+            if resource.errored:
+                self.delete_resource(resource)
         except BadResource:
             self.delete_resource(resource)
             raise

riak/transports/tcp/__init__.py

@@ -1,7 +1,7 @@
 import errno
 import socket
 
-from riak.transports.pool import Pool
+from riak.transports.pool import Pool, ConnectionClosed
 from riak.transports.tcp.transport import TcpTransport
 
 
@@ -48,7 +48,13 @@ def is_retryable(err):
 
     :rtype: boolean
     """
-    if isinstance(err, socket.error):
+    if isinstance(err, ConnectionClosed):
+        # NB: only retryable if we're not mid-streaming
+        if err.mid_stream:
+            return False
+        else:
+            return True
+    elif isinstance(err, socket.error):
         code = err.args[0]
         return code in CONN_CLOSED_ERRORS
     else: