fix(fetcher): Retry transient Datadog query failures

py/tests/test_fetcher.py

@@ -1,8 +1,10 @@
 import copy
 import os
 import unittest
+from email.message import Message
 from io import StringIO
 from unittest.mock import Mock, patch
+from urllib.error import HTTPError, URLError
 
 from arroyo.backends.kafka.consumer import KafkaPayload
 from arroyo.backends.local.backend import LocalBroker
@@ -331,3 +333,59 @@ def test_main_no_unit(self, mock_time: Mock, mock_query_dd: Mock) -> None:
 
         msg2 = self.broker.consume(Partition(self.topic, 0), 1)
         assert msg2 is None
+
+
+def _http_error(code: int) -> HTTPError:
+    return HTTPError(
+        "https://api.datadoghq.com", code, "error", Message(), None
+    )
+
+
+def _ok_response() -> Mock:
+    response = Mock()
+    response.read.return_value = b'{"series": []}'
+    return response
+
+
+@patch("usageaccountant.datadog_fetcher.time.sleep")
+@patch("usageaccountant.datadog_fetcher.urlopen")
+class TestQueryDatadogRetry(unittest.TestCase):
+    def test_retries_server_error(
+        self, mock_urlopen: Mock, mock_sleep: Mock
+    ) -> None:
+        mock_urlopen.side_effect = [_http_error(503), _ok_response()]
+
+        assert ddf.query_datadog("q", 0, 1) == {"series": []}
+        assert mock_urlopen.call_count == 2
+        assert mock_sleep.call_count == 1
+
+    def test_retries_network_error(
+        self, mock_urlopen: Mock, mock_sleep: Mock
+    ) -> None:
+        mock_urlopen.side_effect = [
+            URLError("timed out"),
+            TimeoutError(),
+            _http_error(429),
+            _ok_response(),
+        ]
+
+        assert ddf.query_datadog("q", 0, 1) == {"series": []}
+        assert mock_urlopen.call_count == 4
+
+    def test_does_not_retry_client_error(
+        self, mock_urlopen: Mock, mock_sleep: Mock
+    ) -> None:
+        mock_urlopen.side_effect = [_http_error(403)]
+
+        self.assertRaises(HTTPError, ddf.query_datadog, "q", 0, 1)
+        assert mock_urlopen.call_count == 1
+        mock_sleep.assert_not_called()
+
+    def test_gives_up_after_max_attempts(
+        self, mock_urlopen: Mock, mock_sleep: Mock
+    ) -> None:
+        mock_urlopen.side_effect = _http_error(503)
+
+        self.assertRaises(HTTPError, ddf.query_datadog, "q", 0, 1)
+        assert mock_urlopen.call_count == ddf.QUERY_MAX_ATTEMPTS
+        assert mock_sleep.call_count == ddf.QUERY_MAX_ATTEMPTS - 1

py/usageaccountant/datadog_fetcher.py

@@ -2,6 +2,8 @@
 import json
 import logging
 import os
+import random
+import time
 import urllib.parse
 from typing import (
     Any,
@@ -16,6 +18,7 @@
     TypedDict,
     cast,
 )
+from urllib.error import HTTPError, URLError
 from urllib.request import Request, urlopen
 
 from usageaccountant.accumulator import (
@@ -29,6 +32,10 @@
     "DD-API-KEY": os.environ.get("DATADOG_API_KEY", ""),
 }
 
+QUERY_MAX_ATTEMPTS = 5
+QUERY_BACKOFF_BASE_SECONDS = 2.0
+QUERY_TIMEOUT_SECONDS = 30
+
 logger = logging.getLogger("fetcher")
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
@@ -196,11 +203,38 @@ def query_datadog(query: str, start_time: int, end_time: int) -> Any:
     full_url = f"{base_url}?{query_string}"
 
     request = Request(url=full_url, headers=headers)
-    response = urlopen(request).read()
-
+    for attempt in range(1, QUERY_MAX_ATTEMPTS):
+        try:
+            return _fetch_json(request)
+        except (URLError, TimeoutError, ConnectionError) as e:
+            if not _is_retryable(e):
+                raise
+            delay = QUERY_BACKOFF_BASE_SECONDS * 2 ** (attempt - 1)
+            delay += random.uniform(0, delay)
+            logger.warning(
+                f"Datadog query failed ({e}), attempt {attempt} of "
+                f"{QUERY_MAX_ATTEMPTS}; retrying in {delay:.1f}s."
+            )
+            time.sleep(delay)
+
+    return _fetch_json(request)
+
+
+def _fetch_json(request: Request) -> Any:
+    response = urlopen(request, timeout=QUERY_TIMEOUT_SECONDS).read()
     return json.loads(response)
 
 
+def _is_retryable(error: Exception) -> bool:
+    """
+    Rate limits, server errors, and network failures are transient.
+    Other HTTP errors, such as 400 or 403, will not succeed on retry.
+    """
+    if isinstance(error, HTTPError):
+        return error.code == 429 or error.code >= 500
+    return True
+
+
 def parse_and_assert_response_series(
     response: Mapping[Any, Any]
 ) -> Sequence[DatadogResponseSeries]:
@@ -447,8 +481,6 @@ def main(
     args = parser.parse_args()
 
     if args.start_time is None:
-        import time
-
         current_time = int(time.time())
         args.start_time = (
             current_time