Stop a rate-limited call from sleeping for hours

urllib3 sleeps whatever Retry-After says, clamped at 6 h in 2.7, and retries three times. api.hanzo.ai sends ~23,000 s for an exhausted quota, so one GET held its caller for 18 h. Client sleeps through at most 60 s; past that the answer returns at once, as ApiException with the header for generated operations and as Fault.retry_after for the six.

Co-authored-by: Hanzo Dev <[email protected]>

pkg/hanzoai/answer.py

@@ -194,18 +194,23 @@ class Fault(Exception):
     cache a policy answer nobody gave.
     """
 
-    def __init__(self, status: int, code: str = "", reason: str = "", request: str = "") -> None:
+    def __init__(self, status: int, code: str = "", reason: str = "", request: str = "", retry_after: float = 0) -> None:
         self.status = status
         #: The RFC 9457 `code`, empty where the answer carried none — including
         #: every fault the SDK raises before a request goes out.
         self.code = code
         self.reason = reason
         self.request = request
+        #: Seconds the server asked for before a retry (Retry-After), 0 when it
+        #: named none. A wait past :data:`hanzoai.client.WAIT` arrives here at once.
+        self.retry_after = retry_after
         at = str(status)
         if code:
             at += " " + code
         if request:
             at += " (request {0})".format(request)
+        if retry_after:
+            at += " (retry after {0:g}s)".format(retry_after)
         super().__init__("hanzoai: {0}: {1}".format(at, reason))
 
 
@@ -257,4 +262,4 @@ def fault(reply: Reply) -> "Fault":
     reason = text(reply.body, "detail") or text(reply.body, "title")
     if not reason:
         reason = reply.body if isinstance(reply.body, str) else "no readable body"
-    return Fault(reply.status, text(reply.body, "code"), reason, reply.request)
+    return Fault(reply.status, text(reply.body, "code"), reason, reply.request, reply.retry_after)

pkg/hanzoai/client.py

@@ -22,9 +22,14 @@ from __future__ import annotations
 
 import os
 import json
+import time
 from typing import Any, Dict, Optional
+from email.utils import parsedate_to_datetime
 from urllib.parse import urlencode
 
+from urllib3.util.retry import Retry
+from urllib3.exceptions import MaxRetryError, ResponseError
+
 from hanzoai import wire
 from hanzoai.kb import Kb
 from hanzoai.audit import Audit
@@ -40,6 +45,38 @@ from hanzoai.cloud.configuration import Configuration
 
 __all__ = ["Client"]
 
+#: The longest Retry-After a call sleeps through, in seconds. api.hanzo.ai
+#: answers an exhausted quota with the time until it resets, which is hours; a
+#: longer wait comes back to the caller at once instead.
+WAIT = 60
+
+
+class _Retry(Retry):
+    """urllib3's retries, ended at once by a server that asks for more than `WAIT`.
+
+    urllib3 sleeps for whatever Retry-After says, on each of its retries. Here a
+    longer wait ends the retries and the answer reaches the caller as it arrived:
+    a generated operation raises `ApiException`, whose `headers` carry
+    Retry-After, and a reply from :meth:`Client.send` faults with `retry_after`.
+    """
+
+    def increment(
+        self,
+        method: Optional[str] = None,
+        url: Optional[str] = None,
+        response: Any = None,
+        error: Optional[Exception] = None,
+        _pool: Any = None,
+        _stacktrace: Any = None,
+    ) -> Retry:
+        if (
+            response is not None
+            and response.status in self.RETRY_AFTER_STATUS_CODES
+            and _wait(response.headers.get("Retry-After", "")) > WAIT
+        ):
+            raise MaxRetryError(_pool, url, ResponseError("Retry-After is longer than {0}s".format(WAIT)))
+        return super().increment(method, url, response, error, _pool, _stacktrace)
+
 
 class Client(ApiClient):
     """An API client that holds an IAM credential and answers the six.
@@ -75,7 +112,9 @@ class Client(ApiClient):
         self.base = (base or os.environ.get("HANZO_BASE_URL") or BASE).rstrip("/")
         self.issuer = (issuer or os.environ.get("HANZO_ISSUER_URL") or ISSUER).rstrip("/")
         self.resource = resource or os.environ.get("HANZO_RESOURCE") or self.base
-        super().__init__(Configuration(host=self.base))
+        # urllib3 takes a Retry where the generated hint names an int.
+        retries: Any = _Retry(3, raise_on_status=False)
+        super().__init__(Configuration(host=self.base, retries=retries))
 
         if credential is None:
             # Construction never fails on a missing credential; the first call
@@ -217,4 +256,15 @@ def _reply(response: Any) -> wire.Reply:
         status=response.status,
         body=body if raw else None,
         request=response.getheader("x-request-id", "") or "",
+        retry_after=_wait(response.getheader("retry-after", "") or ""),
     )
+
+
+def _wait(header: str) -> float:
+    """Retry-After in seconds, in either form RFC 9110 allows; 0 for none, or one unreadable."""
+    if header.strip().isdigit():
+        return float(header)
+    try:
+        return max(parsedate_to_datetime(header).timestamp() - time.time(), 0.0)
+    except (TypeError, ValueError):
+        return 0.0

pkg/hanzoai/wire.py

@@ -31,6 +31,9 @@ class Reply:
     status: int
     body: Any = None
     request: str = ""
+    #: Retry-After in seconds: how long the server asked the caller to wait
+    #: before trying again, and 0 when it named no wait.
+    retry_after: float = 0
 
 
 def query(params: Optional[dict]) -> dict:

tests/test_retry.py

@@ -0,0 +1,107 @@
+"""A server that asks for a long wait gets an answer back, not a sleeping caller.
+
+api.hanzo.ai answers an exhausted quota with 429 and Retry-After set to the
+seconds until the quota resets, about 23,000. urllib3 sleeps for whatever
+Retry-After says and retries three times, so a call left to its defaults holds
+its caller for most of a day. The client sleeps through a wait of at most
+:data:`hanzoai.client.WAIT` seconds; past that the answer reaches the caller at
+once, with the wait on it.
+
+A local server answers over a real socket. `time.sleep` is recorded rather than
+slept, so a call that would block for hours fails here in milliseconds.
+"""
+
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pytest
+
+from hanzoai import Fault, Client
+from hanzoai.answer import value
+from hanzoai.cloud import AiApi
+from hanzoai.cloud.exceptions import ApiException
+
+
+class Minted:
+    """A credential that already holds its token, so no call goes to IAM."""
+
+    def token(self):
+        return "tok"
+
+    def invalidate(self):
+        pass
+
+
+@pytest.fixture
+def slept(monkeypatch):
+    """Every sleep urllib3 asks for, in order, and none of them taken."""
+    asked = []
+    monkeypatch.setattr("urllib3.util.retry.time.sleep", asked.append)
+    return asked
+
+
+@pytest.fixture
+def server():
+    """A server on a free port: the base URL, the staged answers, the paths it saw."""
+    staged, seen = [], []
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            seen.append(self.path)
+            status, headers, body = staged.pop(0) if staged else (500, {}, {"title": "nothing staged"})
+            data = json.dumps(body).encode()
+            self.send_response(status)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(data)))
+            for name, v in headers.items():
+                self.send_header(name, v)
+            self.end_headers()
+            self.wfile.write(data)
+
+        def log_message(self, *args):
+            pass
+
+    httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    threading.Thread(target=httpd.serve_forever, daemon=True).start()
+    try:
+        yield "http://127.0.0.1:{0}".format(httpd.server_address[1]), staged, seen
+    finally:
+        httpd.shutdown()
+        httpd.server_close()
+
+
+def limited(n, after, request="req_1"):
+    return [(429, {"Retry-After": after, "x-request-id": request}, {"title": "Too Many Requests"})] * n
+
+
+def test_a_long_wait_raises_at_once(server, slept):
+    """23,000 seconds is not waited: one request, no sleep, and the 429 carries the header."""
+    base, staged, seen = server
+    staged.extend(limited(4, "23000"))
+    with pytest.raises(Exception) as raised:
+        AiApi(Client(base=base, credential=Minted())).get_models()
+    assert (slept, len(seen)) == ([], 1)
+    assert isinstance(raised.value, ApiException)
+    assert (raised.value.status, raised.value.headers["Retry-After"]) == (429, "23000")
+
+
+def test_a_short_wait_is_slept_and_retried(server, slept):
+    """A wait under the ceiling is honoured as before: three sleeps, four requests, the API's error."""
+    base, staged, seen = server
+    staged.extend([(503, {"Retry-After": "1"}, {"title": "Service Unavailable"})] * 4)
+    with pytest.raises(Exception) as raised:
+        AiApi(Client(base=base, credential=Minted())).get_models()
+    assert (slept, len(seen)) == ([1, 1, 1], 4)
+    assert isinstance(raised.value, ApiException) and raised.value.status == 503
+
+
+def test_the_fault_carries_the_wait(server, slept):
+    """A capability reads the same answer through `send`, and its fault says how long to wait."""
+    base, staged, seen = server
+    staged.extend(limited(4, "23000", request="req_9"))
+    with pytest.raises(Exception) as raised:
+        value(Client(base=base, credential=Minted()).send("GET", "/v1/budget"))
+    assert (slept, len(seen)) == ([], 1)
+    assert isinstance(raised.value, Fault)
+    assert (raised.value.status, raised.value.retry_after, raised.value.request) == (429, 23000, "req_9")