fix runtime heartbeat API retries

scripts/runtime_workflow_heartbeat.py

@@ -7,11 +7,18 @@
 import json
 import os
 import sys
+import time
+import urllib.error
 import urllib.parse
 import urllib.request
 from typing import Any
 
 
+_GITHUB_API_MAX_ATTEMPTS = 4
+_GITHUB_API_MAX_RETRY_DELAY_SECONDS = 30.0
+_RETRYABLE_HTTP_STATUSES = frozenset({429, 500, 502, 503, 504})
+
+
 def _split_values(raw: str | None) -> list[str]:
     if not raw:
         return []
@@ -37,6 +44,16 @@ def _parse_timestamp(value: Any) -> dt.datetime | None:
     return parsed.astimezone(dt.timezone.utc)
 
 
+def _github_retry_delay(exc: urllib.error.HTTPError | urllib.error.URLError, attempt: int) -> float:
+    retry_after = exc.headers.get("Retry-After") if isinstance(exc, urllib.error.HTTPError) else None
+    if retry_after:
+        try:
+            return min(max(float(retry_after), 0.0), _GITHUB_API_MAX_RETRY_DELAY_SECONDS)
+        except ValueError:
+            pass
+    return min(float(2 ** (attempt - 1)), _GITHUB_API_MAX_RETRY_DELAY_SECONDS)
+
+
 def _github_request(url: str, token: str) -> dict[str, Any]:
     request = urllib.request.Request(
         url,
@@ -46,8 +63,32 @@ def _github_request(url: str, token: str) -> dict[str, Any]:
             "X-GitHub-Api-Version": "2022-11-28",
         },
     )
-    with urllib.request.urlopen(request, timeout=20) as response:
-        return json.loads(response.read().decode("utf-8"))
+    for attempt in range(1, _GITHUB_API_MAX_ATTEMPTS + 1):
+        retry_error: urllib.error.HTTPError | urllib.error.URLError | None = None
+        try:
+            with urllib.request.urlopen(request, timeout=20) as response:
+                return json.loads(response.read().decode("utf-8"))
+        except urllib.error.HTTPError as exc:
+            retryable = exc.code in _RETRYABLE_HTTP_STATUSES
+            if not retryable or attempt == _GITHUB_API_MAX_ATTEMPTS:
+                raise
+            retry_error = exc
+        except urllib.error.URLError as exc:
+            if attempt == _GITHUB_API_MAX_ATTEMPTS:
+                raise
+            retry_error = exc
+
+        assert retry_error is not None
+        delay = _github_retry_delay(retry_error, attempt)
+        print(
+            "GitHub API request failed "
+            f"({retry_error}); retrying in {delay:g}s "
+            f"(attempt {attempt + 1}/{_GITHUB_API_MAX_ATTEMPTS})",
+            file=sys.stderr,
+        )
+        time.sleep(delay)
+
+    raise RuntimeError("unreachable")
 
 
 def _workflow_paths(workflow: str) -> set[str]:

tests/test_runtime_workflow_heartbeat.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import datetime as dt
+import json
 import os
+import urllib.error
 import unittest
 from unittest.mock import patch
 
@@ -14,6 +16,62 @@ def _timestamp(minutes_ago: int) -> str:
 
 
 class RuntimeWorkflowHeartbeatTests(unittest.TestCase):
+    def test_github_request_retries_service_unavailable(self) -> None:
+        class FakeResponse:
+            def __enter__(self) -> FakeResponse:
+                return self
+
+            def __exit__(self, *args: object) -> None:
+                return None
+
+            def read(self) -> bytes:
+                return json.dumps({"workflow_runs": []}).encode()
+
+        unavailable = urllib.error.HTTPError(
+            "https://api.github.com/example",
+            503,
+            "Service Unavailable",
+            {"Retry-After": "0"},
+            None,
+        )
+        with patch.object(
+            heartbeat.urllib.request,
+            "urlopen",
+            side_effect=[unavailable, FakeResponse()],
+        ) as urlopen:
+            with patch.object(heartbeat.time, "sleep") as sleep:
+                result = heartbeat._github_request("https://api.github.com/example", "token-1")
+
+        self.assertEqual(result, {"workflow_runs": []})
+        self.assertEqual(urlopen.call_count, 2)
+        sleep.assert_called_once_with(0.0)
+
+    def test_github_request_does_not_retry_non_transient_http_error(self) -> None:
+        forbidden = urllib.error.HTTPError(
+            "https://api.github.com/example",
+            403,
+            "Forbidden",
+            {},
+            None,
+        )
+        with patch.object(heartbeat.urllib.request, "urlopen", side_effect=forbidden) as urlopen:
+            with patch.object(heartbeat.time, "sleep") as sleep:
+                with self.assertRaises(urllib.error.HTTPError):
+                    heartbeat._github_request("https://api.github.com/example", "token-1")
+
+        urlopen.assert_called_once()
+        sleep.assert_not_called()
+
+    def test_github_request_stops_after_bounded_network_retries(self) -> None:
+        unavailable = urllib.error.URLError("temporary DNS failure")
+        with patch.object(heartbeat.urllib.request, "urlopen", side_effect=unavailable) as urlopen:
+            with patch.object(heartbeat.time, "sleep") as sleep:
+                with self.assertRaises(urllib.error.URLError):
+                    heartbeat._github_request("https://api.github.com/example", "token-1")
+
+        self.assertEqual(urlopen.call_count, heartbeat._GITHUB_API_MAX_ATTEMPTS)
+        self.assertEqual([call.args[0] for call in sleep.call_args_list], [1.0, 2.0, 4.0])
+
     def test_repository_runs_fallback_finds_recent_runtime_success(self) -> None:
         runtime_run = {
             "id": 1,