From 4fd04250c9d84de1a3db45b5ca43a55b74d79dbd Mon Sep 17 00:00:00 2001
From: Brett Holman <brett.holman@canonical.com>
Date: Thu, 1 Aug 2024 20:19:34 -0600
Subject: [PATCH] feat(url_helper): Retry on 503 error

---
 cloudinit/sources/DataSourceEc2.py | 13 ++++++++++---
 cloudinit/url_helper.py            | 27 ++++++++++++++++++---------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/cloudinit/sources/DataSourceEc2.py b/cloudinit/sources/DataSourceEc2.py
index 10837df6a0e..412f4d709b9 100644
--- a/cloudinit/sources/DataSourceEc2.py
+++ b/cloudinit/sources/DataSourceEc2.py
@@ -655,17 +655,24 @@ def _imds_exception_cb(self, msg, exception=None):
         """
         if isinstance(exception, uhelp.UrlError):
             # requests.ConnectionError will have exception.code == None
-            if exception.code and exception.code >= 400:
+            if exception.code:
                 if exception.code == 403:
                     LOG.warning(
                         "Ec2 IMDS endpoint returned a 403 error. "
                         "HTTP endpoint is disabled. Aborting."
                     )
-                else:
+                    raise exception
+                elif exception.code == 503:
+                    LOG.warning(
+                        "Ec2 IMDS endpoint returned a 503 error. "
+                        "HTTP endpoint is overloaded. Retrying."
+                    )
+                    return
+                elif exception.code >= 400:
                     LOG.warning(
                         "Fatal error while requesting Ec2 IMDSv2 API tokens"
                     )
-                raise exception
+                    raise exception
 
     def _get_headers(self, url=""):
         """Return a dict of headers for accessing a url.
diff --git a/cloudinit/url_helper.py b/cloudinit/url_helper.py
index a1d095ccb9b..e8577c93776 100644
--- a/cloudinit/url_helper.py
+++ b/cloudinit/url_helper.py
@@ -364,6 +364,7 @@ def _get_ssl_args(url, ssl_details):
 
 def readurl(
     url,
+    *,
     data=None,
     timeout=None,
     retries=0,
@@ -641,6 +642,7 @@ def dual_stack(
 
 def wait_for_url(
     urls,
+    *,
     max_wait: float = float("inf"),
     timeout: Optional[float] = None,
     status_cb: Callable = LOG.debug,  # some sources use different log levels
@@ -707,7 +709,7 @@ def timeup(max_wait: float, start_time: float, sleep_time: float = 0):
             time.monotonic() - start_time + sleep_time > max_wait
         )
 
-    def handle_url_response(response, url):
+    def handle_url_response(response, url) -> Tuple[Optional[Exception], str]:
         """Map requests response code/contents to internal "UrlError" type"""
         if not response.contents:
             reason = "empty response [%s]" % (response.code)
@@ -732,10 +734,10 @@ def handle_url_response(response, url):
 
     def read_url_handle_exceptions(
         url_reader_cb, urls, start_time, exc_cb, log_cb
-    ):
+    ) -> Tuple[str, Union[Exception, UrlResponse]]:
         """Execute request, handle response, optionally log exception"""
         reason = ""
-        url = None
+        url = ""
         try:
             url, response = url_reader_cb(urls)
             url_exc, reason = handle_url_response(response, url)
@@ -761,8 +763,9 @@ def read_url_handle_exceptions(
             # in the future, for example this is what the MAAS datasource
             # does.
             exc_cb(msg=status_msg, exception=url_exc)
+        return url, url_exc
 
-    def read_url_cb(url, timeout):
+    def read_url_cb(url: str, timeout: int) -> UrlResponse:
         return readurl(
             url,
             headers={} if headers_cb is None else headers_cb(url),
@@ -772,19 +775,21 @@ def read_url_cb(url, timeout):
             request_method=request_method,
         )
 
-    def read_url_serial(start_time, timeout, exc_cb, log_cb):
+    def read_url_serial(
+        start_time, timeout, exc_cb, log_cb
+    ) -> Optional[Tuple[str, Union[Exception, UrlResponse]]]:
         """iterate over list of urls, request each one and handle responses
         and thrown exceptions individually per url
         """
 
-        def url_reader_serial(url):
+        def url_reader_serial(url: str):
             return (url, read_url_cb(url, timeout))
 
         for url in urls:
             now = time.monotonic()
             if loop_n != 0:
                 if timeup(max_wait, start_time):
-                    return
+                    return None
                 if (
                     max_wait is not None
                     and timeout
@@ -798,8 +803,11 @@ def url_reader_serial(url):
             )
             if out:
                 return out
+        return None
 
-    def read_url_parallel(start_time, timeout, exc_cb, log_cb):
+    def read_url_parallel(
+        start_time, timeout, exc_cb, log_cb
+    ) -> Optional[Tuple[str, Union[Exception, UrlResponse]]]:
         """pass list of urls to dual_stack which sends requests in parallel
         handle response and exceptions of the first endpoint to respond
         """
@@ -838,7 +846,8 @@ def read_url_parallel(start_time, timeout, exc_cb, log_cb):
         url = do_read_url(start_time, timeout, exception_cb, status_cb)
         if url:
             address, response = url
-            return (address, response.contents)
+            if isinstance(response, UrlResponse):
+                return (address, response.contents)
 
         if timeup(max_wait, start_time, current_sleep_time):
             break