move retry logic to standalone module

Author: vgvoleg

tests/query/test_query_transaction.py

@@ -65,6 +65,6 @@ def test_context_manager_normal_flow(self, tx: BaseTxContext):
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
     def test_context_manager_does_not_hide_exceptions(self, tx: BaseTxContext):
-        with pytest.raises(RuntimeError):
+        with pytest.raises(Exception):
             with tx:
-                tx.commit()
+                raise Exception()

ydb/_retries.py

@@ -0,0 +1,136 @@
+import random
+import time
+
+from . import issues
+from ._errors import check_retriable_error
+
+
+class BackoffSettings(object):
+    def __init__(self, ceiling=6, slot_duration=0.001, uncertain_ratio=0.5):
+        self.ceiling = ceiling
+        self.slot_duration = slot_duration
+        self.uncertain_ratio = uncertain_ratio
+
+    def calc_timeout(self, retry_number):
+        slots_count = 1 << min(retry_number, self.ceiling)
+        max_duration_ms = slots_count * self.slot_duration * 1000.0
+        # duration_ms = random.random() * max_duration_ms * uncertain_ratio) + max_duration_ms * (1 - uncertain_ratio)
+        duration_ms = max_duration_ms * (random.random() * self.uncertain_ratio + 1.0 - self.uncertain_ratio)
+        return duration_ms / 1000.0
+
+
+class RetrySettings(object):
+    def __init__(
+        self,
+        max_retries=10,
+        max_session_acquire_timeout=None,
+        on_ydb_error_callback=None,
+        backoff_ceiling=6,
+        backoff_slot_duration=1,
+        get_session_client_timeout=5,
+        fast_backoff_settings=None,
+        slow_backoff_settings=None,
+        idempotent=False,
+    ):
+        self.max_retries = max_retries
+        self.max_session_acquire_timeout = max_session_acquire_timeout
+        self.on_ydb_error_callback = (lambda e: None) if on_ydb_error_callback is None else on_ydb_error_callback
+        self.fast_backoff = BackoffSettings(10, 0.005) if fast_backoff_settings is None else fast_backoff_settings
+        self.slow_backoff = (
+            BackoffSettings(backoff_ceiling, backoff_slot_duration)
+            if slow_backoff_settings is None
+            else slow_backoff_settings
+        )
+        self.retry_not_found = True
+        self.idempotent = idempotent
+        self.retry_internal_error = True
+        self.unknown_error_handler = lambda e: None
+        self.get_session_client_timeout = get_session_client_timeout
+        if max_session_acquire_timeout is not None:
+            self.get_session_client_timeout = min(self.max_session_acquire_timeout, self.get_session_client_timeout)
+
+    def with_fast_backoff(self, backoff_settings):
+        self.fast_backoff = backoff_settings
+        return self
+
+    def with_slow_backoff(self, backoff_settings):
+        self.slow_backoff = backoff_settings
+        return self
+
+
+class YdbRetryOperationSleepOpt(object):
+    def __init__(self, timeout):
+        self.timeout = timeout
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.timeout == other.timeout
+
+    def __repr__(self):
+        return "YdbRetryOperationSleepOpt(%s)" % self.timeout
+
+
+class YdbRetryOperationFinalResult(object):
+    def __init__(self, result):
+        self.result = result
+        self.exc = None
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.result == other.result and self.exc == other.exc
+
+    def __repr__(self):
+        return "YdbRetryOperationFinalResult(%s, exc=%s)" % (self.result, self.exc)
+
+    def set_exception(self, exc):
+        self.exc = exc
+
+
+def retry_operation_impl(callee, retry_settings=None, *args, **kwargs):
+    retry_settings = RetrySettings() if retry_settings is None else retry_settings
+    status = None
+
+    for attempt in range(retry_settings.max_retries + 1):
+        try:
+            result = YdbRetryOperationFinalResult(callee(*args, **kwargs))
+            yield result
+
+            if result.exc is not None:
+                raise result.exc
+
+        except issues.Error as e:
+            status = e
+            retry_settings.on_ydb_error_callback(e)
+
+            retriable_info = check_retriable_error(e, retry_settings, attempt)
+            if not retriable_info.is_retriable:
+                raise
+
+            skip_yield_error_types = [
+                issues.Aborted,
+                issues.BadSession,
+                issues.NotFound,
+                issues.InternalError,
+            ]
+
+            yield_sleep = True
+            for t in skip_yield_error_types:
+                if isinstance(e, t):
+                    yield_sleep = False
+
+            if yield_sleep:
+                yield YdbRetryOperationSleepOpt(retriable_info.sleep_timeout_seconds)
+
+        except Exception as e:
+            # you should provide your own handler you want
+            retry_settings.unknown_error_handler(e)
+            raise
+
+    raise status
+
+
+def retry_operation_sync(callee, retry_settings=None, *args, **kwargs):
+    opt_generator = retry_operation_impl(callee, retry_settings, *args, **kwargs)
+    for next_opt in opt_generator:
+        if isinstance(next_opt, YdbRetryOperationSleepOpt):
+            time.sleep(next_opt.timeout)
+        else:
+            return next_opt.result

ydb/_topic_reader/topic_reader.py

@@ -10,7 +10,7 @@
     Callable,
 )
 
-from ..table import RetrySettings
+from .._retries import RetrySettings
 from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, OffsetsRange
 
 

ydb/_topic_writer/topic_writer_asyncio.py

@@ -26,9 +26,9 @@
 from .. import (
     _apis,
     issues,
-    check_retriable_error,
-    RetrySettings,
 )
+from .._errors import check_retriable_error
+from .._retries import RetrySettings
 from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec
 from .._grpc.grpcwrapper.ydb_topic import (
     UpdateTokenRequest,

ydb/query/pool.py

@@ -1,98 +1,13 @@
-import time
 from typing import Callable
 
 from . import base
 from .session import (
     QuerySessionSync,
 )
-from .. import issues
-from .._errors import check_retriable_error
-
-
-class RetrySettings(object):
-    def __init__(
-        self,
-        max_retries: int = 10,
-        max_session_acquire_timeout: int = None,
-        on_ydb_error_callback: Callable = None,
-        idempotent: bool = False,
-    ):
-        self.max_retries = max_retries
-        self.max_session_acquire_timeout = max_session_acquire_timeout
-        self.on_ydb_error_callback = (lambda e: None) if on_ydb_error_callback is None else on_ydb_error_callback
-        self.retry_not_found = True
-        self.idempotent = idempotent
-        self.retry_internal_error = True
-        self.unknown_error_handler = lambda e: None
-
-
-class YdbRetryOperationSleepOpt:
-    def __init__(self, timeout):
-        self.timeout = timeout
-
-    def __eq__(self, other):
-        return type(self) == type(other) and self.timeout == other.timeout
-
-    def __repr__(self):
-        return "YdbRetryOperationSleepOpt(%s)" % self.timeout
-
-
-class YdbRetryOperationFinalResult:
-    def __init__(self, result):
-        self.result = result
-        self.exc = None
-
-    def __eq__(self, other):
-        return type(self) == type(other) and self.result == other.result and self.exc == other.exc
-
-    def __repr__(self):
-        return "YdbRetryOperationFinalResult(%s, exc=%s)" % (self.result, self.exc)
-
-    def set_exception(self, exc):
-        self.exc = exc
-
-
-def retry_operation_impl(callee: Callable, retry_settings: RetrySettings = None, *args, **kwargs):
-    retry_settings = RetrySettings() if retry_settings is None else retry_settings
-    status = None
-
-    for attempt in range(retry_settings.max_retries + 1):
-        try:
-            result = YdbRetryOperationFinalResult(callee(*args, **kwargs))
-            yield result
-
-            if result.exc is not None:
-                raise result.exc
-
-        except issues.Error as e:
-            status = e
-            retry_settings.on_ydb_error_callback(e)
-
-            retriable_info = check_retriable_error(e, retry_settings, attempt)
-            if not retriable_info.is_retriable:
-                raise
-
-            skip_yield_error_types = [
-                issues.Aborted,
-                issues.BadSession,
-                issues.NotFound,
-                issues.InternalError,
-            ]
-
-            yield_sleep = True
-            for t in skip_yield_error_types:
-                if isinstance(e, t):
-                    yield_sleep = False
-
-            if yield_sleep:
-                yield YdbRetryOperationSleepOpt(retriable_info.sleep_timeout_seconds)
-
-        except Exception as e:
-            # you should provide your own handler you want
-            retry_settings.unknown_error_handler(e)
-            raise
-    if status:
-        raise status
+from .._retries import (
+    RetrySettings,
+    retry_operation_sync,
+)
 
 
 class QuerySessionPool:
@@ -122,12 +37,7 @@ def wrapped_callee():
             with self.checkout() as session:
                 return callee(session, *args, **kwargs)
 
-        opt_generator = retry_operation_impl(wrapped_callee, retry_settings, *args, **kwargs)
-        for next_opt in opt_generator:
-            if isinstance(next_opt, YdbRetryOperationSleepOpt):
-                time.sleep(next_opt.timeout)
-            else:
-                return next_opt.result
+        return retry_operation_sync(wrapped_callee, retry_settings)
 
     def execute_with_retries(self, query: str, retry_settings: RetrySettings = None, *args, **kwargs):
         """Special interface to execute a one-shot queries in a safe, retriable way.
@@ -144,12 +54,7 @@ def wrapped_callee():
                 it = session.execute(query, empty_tx_control=True, *args, **kwargs)
                 return [result_set for result_set in it]
 
-            opt_generator = retry_operation_impl(wrapped_callee, retry_settings)
-            for next_opt in opt_generator:
-                if isinstance(next_opt, YdbRetryOperationSleepOpt):
-                    time.sleep(next_opt.timeout)
-                else:
-                    return next_opt.result
+            return retry_operation_sync(wrapped_callee, retry_settings)
 
 
 class SimpleQuerySessionCheckout: