From fe398aa821c3a1093f87623d8617d1d668ca2100 Mon Sep 17 00:00:00 2001 From: Pigbibi <20649888+Pigbibi@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:20:38 +0800 Subject: [PATCH] fix runtime heartbeat API retries Co-Authored-By: Codex --- scripts/runtime_workflow_heartbeat.py | 45 +++++++++++++++++- tests/test_runtime_workflow_heartbeat.py | 58 ++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/scripts/runtime_workflow_heartbeat.py b/scripts/runtime_workflow_heartbeat.py index 7c46a13f..9f9df99f 100644 --- a/scripts/runtime_workflow_heartbeat.py +++ b/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]: diff --git a/tests/test_runtime_workflow_heartbeat.py b/tests/test_runtime_workflow_heartbeat.py index 56ab49bd..98dc2937 100644 --- a/tests/test_runtime_workflow_heartbeat.py +++ b/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,