Add retry decorator allowing to retry DB operations on request

Improve retrying decorator such that wrapped code could request
retry attempt in the arbitrary conditions by raising RetryRequest
exception. Wrapped code should provide inner exception that is
raised when amount of retries are exceeded.

Change-Id: I596fea8dc798a11fcdfdd69208af3313f41d755b

4 files changed, 75 insertions, 4 deletions

diff --git a/oslo_db/api.py b/oslo_db/api.py
index dc71d96..98e4d75 100644
--- a/oslo_db/api.py
+++ b/oslo_db/api.py
@@ -24,6 +24,7 @@ API methods.
 """
 
 import logging
+import sys
 import threading
 import time
 
@@ -69,6 +70,16 @@ def retry_on_deadlock(f):
     return f
 
 
+def retry_on_request(f):
+    """Retry a DB API call if RetryRequest exception was received.
+
+    wrap_db_entry will be applied to all db.api functions marked with this
+    decorator.
+    """
+    f.enable_retry_on_request = True
+    return f
+
+
 class wrap_db_retry(object):
     """Retry db.api methods, if db_error raised
 
@@ -91,8 +102,9 @@ class wrap_db_retry(object):
     :type max_retry_interval: int
     """
 
-    def __init__(self, retry_interval, max_retries, inc_retry_interval,
-                 max_retry_interval, retry_on_disconnect, retry_on_deadlock):
+    def __init__(self, retry_interval=0, max_retries=0, inc_retry_interval=0,
+                 max_retry_interval=0, retry_on_disconnect=False,
+                 retry_on_deadlock=False, retry_on_request=False):
         super(wrap_db_retry, self).__init__()
 
         self.db_error = ()
@@ -100,6 +112,8 @@ class wrap_db_retry(object):
             self.db_error += (exception.DBConnectionError, )
         if retry_on_deadlock:
             self.db_error += (exception.DBDeadlock, )
+        if retry_on_request:
+            self.db_error += (exception.RetryRequest, )
         self.retry_interval = retry_interval
         self.max_retries = max_retries
         self.inc_retry_interval = inc_retry_interval
@@ -118,6 +132,10 @@ class wrap_db_retry(object):
                 except db_error as e:
                     if remaining == 0:
                         LOG.exception(_LE('DB exceeded retry limit.'))
+                        if isinstance(e, exception.RetryRequest):
+                            six.reraise(type(e.inner_exc),
+                                        e.inner_exc,
+                                        sys.exc_info()[2])
                         raise e
                     if remaining != -1:
                         remaining -= 1
@@ -210,15 +228,17 @@ class DBAPI(object):
         retry_on_disconnect = self.use_db_reconnect and getattr(
             attr, 'enable_retry_on_disconnect', False)
         retry_on_deadlock = getattr(attr, 'enable_retry_on_deadlock', False)
+        retry_on_request = getattr(attr, 'enable_retry_on_request', False)
 
-        if retry_on_disconnect or retry_on_deadlock:
+        if retry_on_disconnect or retry_on_deadlock or retry_on_request:
             attr = wrap_db_retry(
                 retry_interval=self.retry_interval,
                 max_retries=self.max_retries,
                 inc_retry_interval=self.inc_retry_interval,
                 max_retry_interval=self.max_retry_interval,
                 retry_on_disconnect=retry_on_disconnect,
-                retry_on_deadlock=retry_on_deadlock)(attr)
+                retry_on_deadlock=retry_on_deadlock,
+                retry_on_request=retry_on_request)(attr)
 
         return attr
 
diff --git a/oslo_db/exception.py b/oslo_db/exception.py
index 5de7f1e..f950f6a 100644
--- a/oslo_db/exception.py
+++ b/oslo_db/exception.py
@@ -171,3 +171,12 @@ class BackendNotAvailable(Exception):
     within a test suite.
 
     """
+
+
+class RetryRequest(Exception):
+    """Error raised when DB operation needs to be retried.
+
+    That could be intentionally raised by the code without any real DB errors.
+    """
+    def __init__(self, inner_exc):
+        self.inner_exc = inner_exc
diff --git a/oslo_db/tests/old_import_api/test_api.py b/oslo_db/tests/old_import_api/test_api.py
index 3664e24..1fe3bf3 100644
--- a/oslo_db/tests/old_import_api/test_api.py
+++ b/oslo_db/tests/old_import_api/test_api.py
@@ -279,3 +279,24 @@ class DBDeadlockTestCase(DBAPITestCase):
         self.assertEqual(
             0, self.test_db_api.error_counter,
             'Counter not decremented, retry logic probably failed.')
+
+
+class DBRetryRequestCase(DBAPITestCase):
+    def test_retry_wrapper_succeeds(self):
+        @api.wrap_db_retry(max_retries=10, retry_on_request=True)
+        def some_method():
+            pass
+
+        some_method()
+
+    def test_retry_wrapper_reaches_limit(self):
+        max_retries = 10
+
+        @api.wrap_db_retry(max_retries=10, retry_on_request=True)
+        def some_method(res):
+            res['result'] += 1
+            raise exception.RetryRequest(ValueError())
+
+        res = {'result': 0}
+        self.assertRaises(ValueError, some_method, res)
+        self.assertEqual(max_retries + 1, res['result'])
diff --git a/oslo_db/tests/test_api.py b/oslo_db/tests/test_api.py
index 171e360..18dc586 100644
--- a/oslo_db/tests/test_api.py
+++ b/oslo_db/tests/test_api.py
@@ -175,3 +175,24 @@ class DBReconnectTestCase(DBAPITestCase):
         self.assertNotEqual(
             0, self.test_db_api.error_counter,
             'Retry did not stop after sql_max_retries iterations.')
+
+
+class DBRetryRequestCase(DBAPITestCase):
+    def test_retry_wrapper_succeeds(self):
+        @api.wrap_db_retry(max_retries=10, retry_on_request=True)
+        def some_method():
+            pass
+
+        some_method()
+
+    def test_retry_wrapper_reaches_limit(self):
+        max_retries = 10
+
+        @api.wrap_db_retry(max_retries=10, retry_on_request=True)
+        def some_method(res):
+            res['result'] += 1
+            raise exception.RetryRequest(ValueError())
+
+        res = {'result': 0}
+        self.assertRaises(ValueError, some_method, res)
+        self.assertEqual(max_retries + 1, res['result'])