diff --git a/src/replit/database/database.py b/src/replit/database/database.py index 83ffa62..dfdbc18 100644 --- a/src/replit/database/database.py +++ b/src/replit/database/database.py @@ -3,6 +3,7 @@ import asyncio from collections import abc import json +import secrets import threading from typing import ( Any, @@ -22,6 +23,29 @@ from requests.adapters import HTTPAdapter, Retry from urllib3.filepost import encode_multipart_formdata +_SYSTEM_RANDOM = secrets.SystemRandom() + + +class _JitteredExponentialRetry(ExponentialRetry): + """Add jitter to async exponential backoff to spread retry traffic.""" + + def get_timeout(self, attempt: int, response: Any = None) -> float: + """Return an exponential delay with up to 100% random jitter.""" + timeout = super().get_timeout(attempt, response) + jitter = _SYSTEM_RANDOM.uniform(0, timeout) + return min(timeout + jitter, self._max_timeout) + + +class _JitteredRetry(Retry): + """Add jitter to urllib3 retry backoff to spread retry traffic.""" + + def get_backoff_time(self) -> float: + """Return the usual backoff plus up to 100% random jitter.""" + timeout = super().get_backoff_time() + jitter_range = max(timeout, self.backoff_factor) + jitter = _SYSTEM_RANDOM.uniform(0, jitter_range) + return min(timeout + jitter, self.backoff_max) + def to_primitive(o: Any) -> Any: """If object is an observed object, converts to primitve, otherwise returns it. @@ -106,7 +130,7 @@ def __init__( self._get_db_url = get_db_url self._unbind = unbind - retry_options = ExponentialRetry(attempts=retry_count) + retry_options = _JitteredExponentialRetry(attempts=retry_count) self.client = RetryClient(client_session=self.sess, retry_options=retry_options) if self._get_db_url: @@ -532,7 +556,7 @@ def __init__( self.sess = requests.Session() self._get_db_url = get_db_url self._unbind = unbind - retries = Retry( + retries = _JitteredRetry( total=retry_count, backoff_factor=0.1, status_forcelist=[500, 502, 503, 504] ) self.sess.mount("http://", HTTPAdapter(max_retries=retries)) diff --git a/tests/test_database_retry.py b/tests/test_database_retry.py new file mode 100644 index 0000000..367c63e --- /dev/null +++ b/tests/test_database_retry.py @@ -0,0 +1,33 @@ +"""Tests for database retry backoff.""" + +from unittest import TestCase +from unittest.mock import patch + +from replit.database.database import _JitteredExponentialRetry, _JitteredRetry + + +class TestJitteredExponentialRetry(TestCase): + """Tests for async database retry backoff.""" + + def test_adds_jitter_to_exponential_timeout(self) -> None: + """Async retry delays include jitter while respecting their cap.""" + retry = _JitteredExponentialRetry( + attempts=3, start_timeout=1, factor=2, max_timeout=3 + ) + + with patch("replit.database.database._SYSTEM_RANDOM.uniform", return_value=1.5): + self.assertEqual(retry.get_timeout(attempt=1), 3) + + +class TestJitteredRetry(TestCase): + """Tests for synchronous database retry backoff.""" + + def test_adds_jitter_to_first_retry(self) -> None: + """The first synchronous retry is delayed by random jitter.""" + retry = _JitteredRetry(total=3, backoff_factor=0.1) + retry = retry.increment(method="GET", url="/") + + with patch( + "replit.database.database._SYSTEM_RANDOM.uniform", return_value=0.05 + ): + self.assertEqual(retry.get_backoff_time(), 0.05)