zulip: Optionally retry requests that hit RATE_LIMIT_HIT.

zulip/tests/test_rate_limit.py

@@ -0,0 +1,117 @@
+import json
+from typing import Any, Dict, Optional
+from unittest import TestCase
+from unittest.mock import Mock, patch
+
+import zulip
+
+
+def _make_client(**kwargs: Any) -> zulip.Client:
+    with patch.object(
+        zulip.Client,
+        "get_server_settings",
+        return_value={"zulip_version": "10.0.0", "zulip_feature_level": 0},
+    ):
+        return zulip.Client(
+            email="bot@example.com",
+            api_key="x" * 32,
+            site="https://example.zulipchat.com",
+            **kwargs,
+        )
+
+
+def _response(
+    status_code: int,
+    payload: Optional[Dict[str, Any]] = None,
+    headers: Optional[Dict[str, str]] = None,
+) -> Mock:
+    res = Mock()
+    res.status_code = status_code
+    res.headers = headers or {}
+    if payload is None:
+        res.json.side_effect = json.JSONDecodeError("Expecting value", "", 0)
+    else:
+        res.json.return_value = payload
+    return res
+
+
+class TestParseRetryAfter(TestCase):
+    def test_json_retry_after_preferred(self) -> None:
+        res = _response(
+            429,
+            {
+                "result": "error",
+                "code": "RATE_LIMIT_HIT",
+                "msg": "rate limited",
+                "retry-after": 2.5,
+            },
+            {"Retry-After": "9"},
+        )
+        self.assertEqual(zulip._parse_retry_after_seconds(res), 2.5)  # noqa: SLF001
+
+    def test_retry_after_header_fallback(self) -> None:
+        res = _response(429, {"result": "error", "code": "RATE_LIMIT_HIT", "msg": "rate limited"})
+        res.headers = {"Retry-After": "3.25"}
+        self.assertEqual(zulip._parse_retry_after_seconds(res), 3.25)  # noqa: SLF001
+
+    def test_invalid_values_default_to_one(self) -> None:
+        res = _response(429, None, {"Retry-After": "soon"})
+        self.assertEqual(zulip._parse_retry_after_seconds(res), 1.0)  # noqa: SLF001
+
+    def test_negative_values_default_to_one(self) -> None:
+        res = _response(429, {"retry-after": -4})
+        self.assertEqual(zulip._parse_retry_after_seconds(res), 1.0)  # noqa: SLF001
+
+
+class TestAutoRetryRateLimits(TestCase):
+    def test_default_does_not_retry_429(self) -> None:
+        client = _make_client()
+        client.session = Mock()
+        client.session.request.return_value = _response(
+            429,
+            {"result": "error", "code": "RATE_LIMIT_HIT", "msg": "API usage exceeded rate limit"},
+        )
+        with patch("time.sleep") as sleep:
+            result = client.do_api_query({}, "v1/messages")
+        sleep.assert_not_called()
+        self.assertEqual(result["code"], "RATE_LIMIT_HIT")
+        self.assertEqual(client.session.request.call_count, 1)
+
+    def test_retries_after_retry_after_then_succeeds(self) -> None:
+        client = _make_client(auto_retry_rate_limits=True)
+        client.session = Mock()
+        limited = _response(
+            429,
+            {
+                "result": "error",
+                "code": "RATE_LIMIT_HIT",
+                "msg": "API usage exceeded rate limit",
+                "retry-after": 1.5,
+            },
+        )
+        ok = _response(200, {"result": "success", "id": 1})
+        client.session.request.side_effect = [limited, ok]
+        with patch("time.sleep") as sleep, patch("random.uniform", return_value=0.1):
+            result = client.do_api_query({}, "v1/messages")
+        sleep.assert_called_once_with(1.6)
+        self.assertEqual(result["result"], "success")
+        self.assertEqual(client.session.request.call_count, 2)
+
+    def test_stops_after_ten_rate_limit_retries(self) -> None:
+        client = _make_client(auto_retry_rate_limits=True)
+        client.session = Mock()
+        limited = _response(
+            429,
+            {
+                "result": "error",
+                "code": "RATE_LIMIT_HIT",
+                "msg": "rate limited",
+                "retry-after": 0.01,
+            },
+        )
+        client.session.request.return_value = limited
+        with patch("time.sleep"), patch("random.uniform", return_value=0.0):
+            result = client.do_api_query({}, "v1/messages")
+        # 10 retries after the first 429 => 11 requests, then the error is returned.
+        self.assertEqual(client.session.request.call_count, 11)
+        self.assertEqual(result["code"], "RATE_LIMIT_HIT")

zulip/zulip/__init__.py

@@ -385,6 +385,31 @@ class TimeZoneMissingError(ZulipError):
     pass
 
 
+def _parse_retry_after_seconds(res: requests.Response) -> float:
+    """Return how long to wait after a 429, from JSON or the Retry-After header.
+
+    Zulip documents `retry-after` on RATE_LIMIT_HIT responses:
+    https://zulip.com/api/rest-error-handling#rate-limit-exceeded
+    """
+    wait: Optional[float] = None
+    try:
+        payload = res.json()
+        if isinstance(payload, dict) and payload.get("retry-after") is not None:
+            wait = float(payload["retry-after"])
+    except (TypeError, ValueError, json.JSONDecodeError):
+        wait = None
+    if wait is None:
+        header = res.headers.get("Retry-After")
+        if header is not None:
+            try:
+                wait = float(header)
+            except (TypeError, ValueError):
+                wait = None
+    if wait is None or wait < 0:
+        wait = 1.0
+    return wait
+
+
 class Client:
     def __init__(
         self,
@@ -393,6 +418,7 @@ def __init__(
         config_file: Optional[str] = None,
         verbose: bool = False,
         retry_on_errors: bool = True,
+        auto_retry_rate_limits: bool = False,
         site: Optional[str] = None,
         client: Optional[str] = None,
         cert_bundle: Optional[str] = None,
@@ -491,6 +517,7 @@ def __init__(
             self.base_url += "/api"
         self.base_url += "/"
         self.retry_on_errors = retry_on_errors
+        self.auto_retry_rate_limits = auto_retry_rate_limits
         self.client_name = client
 
         if insecure:
@@ -607,6 +634,7 @@ def do_api_query(
             "had_error_retry": False,
             "request": request,
             "failures": 0,
+            "rate_limit_retries": 0,
         }
 
         def error_retry(error_string: str) -> bool:
@@ -655,6 +683,23 @@ def end_error_retry(succeeded: bool) -> None:
 
                 self.has_connected = True
 
+                # Opt-in: sleep for retry-after (plus a little jitter) and
+                # retry RATE_LIMIT_HIT. Default off so callers that cannot
+                # block are unchanged. Cap retries so a tight limit cannot
+                # pin a client to the server's allowed rate forever.
+                if (
+                    res.status_code == 429
+                    and self.auto_retry_rate_limits
+                    and query_state["rate_limit_retries"] < 10
+                ):
+                    wait = _parse_retry_after_seconds(res)
+                    wait += random.uniform(0.05, 0.5)  # noqa: S311
+                    if self.verbose:
+                        print(f"Rate limit hit; retrying after {wait:.2f}s.")
+                    query_state["rate_limit_retries"] += 1
+                    time.sleep(wait)
+                    continue
+
                 # On 50x errors, try again after a short sleep
                 if str(res.status_code).startswith("5") and error_retry(
                     f" (server {res.status_code})"