Add jitter to database retry backoff

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))

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)