rest: apply client-go retry generation patches

Signed-off-by: Dr. Stefan Schimanski <[email protected]>

Author: sttts

kubernetes/aio/client/configuration.py

@@ -355,6 +355,22 @@ def __init__(
         self.retries = retries
         """Retry configuration
         """
+        self.client_go_retries = False
+        """Enable Kubernetes client-go-compatible retry semantics.
+
+           When enabled, GET and HEAD requests retry Retry-After responses.
+           The retry ceiling is read from ``retries`` when set; otherwise it
+           follows the client-go default of at most 10 retries.
+        """
+        self.client_go_retry_backoff = None
+        """Backoff for Kubernetes client-go-compatible retries.
+
+           If unset, client-go-compatible GET and HEAD retries use the
+           client-go default retry ceiling with no additional client-side
+           delay beyond Retry-After. When set, ``retries`` still overrides
+           the retry ceiling if it is not None.
+        """
+
         self.trace_configs = trace_configs
         """aiohttp.TraceConfig list forwarded to ClientSession for tracing.
         """

kubernetes/aio/client/rest.py

@@ -22,6 +22,11 @@
 import aiohttp_retry
 
 from kubernetes.aio.client.exceptions import ApiException, ApiValueError
+from kubernetes.aio.utils.retry import (
+    is_retry_after_response,
+    on_retry_after_error,
+    retry_after_backoff,
+)
 
 RESTResponseType = aiohttp.ClientResponse
 
@@ -292,6 +297,41 @@ async def request(
                 )
             pool_manager = self.retry_client
 
-        r = await pool_manager.request(**args)
+        async def read_request(check_retry_status=False):
+            response = await self.pool_manager.request(**args)
+            if check_retry_status:
+                self._raise_retry_after_response(response)
+            return response
+
+        if (
+            method in ['GET', 'HEAD']
+            and getattr(self.configuration, 'client_go_retries', False)
+        ):
+            backoff = retry_after_backoff(
+                getattr(self.configuration, 'retries', None),
+                getattr(self.configuration, 'client_go_retry_backoff', None),
+            )
+            r = await on_retry_after_error(
+                backoff, self._is_read_retryable, lambda: read_request(True))
+        else:
+            r = await pool_manager.request(**args)
 
         return RESTResponse(r)
+
+    @classmethod
+    def _is_read_retryable(cls, error):
+        return is_retry_after_response(error)
+
+    @staticmethod
+    def _retry_after_error(response):
+        error = ApiException(status=response.status, reason=response.reason)
+        error.headers = response.headers
+        return error
+
+    @classmethod
+    def _raise_retry_after_response(cls, response):
+        error = cls._retry_after_error(response)
+        if not is_retry_after_response(error):
+            return
+        response.release()
+        raise error

kubernetes/aio/test/test_generated_api.py

@@ -46,6 +46,8 @@ async def asyncSetUp(self):
             'items': [],
         }
         self.response_status = 200
+        self.response_headers = {}
+        self.responses = None
         app = web.Application()
         app.router.add_route('*', '/{path:.*}', self._handle_request)
         self.runner = web.AppRunner(app)
@@ -91,7 +93,17 @@ async def _handle_request(self, request):
             }).encode() + b'\n')
             await response.write_eof()
             return response
-        return web.json_response(self.response, status=self.response_status)
+        if self.responses is not None:
+            index = len(self.requests) - 1
+            response, status, headers = self.responses[
+                min(index, len(self.responses) - 1)
+            ]
+            return web.json_response(response, status=status, headers=headers)
+        return web.json_response(
+            self.response,
+            status=self.response_status,
+            headers=self.response_headers,
+        )
 
     async def test_bearer_alias_supports_synchronous_token_refresh(self):
         self.configuration.api_key['authorization'] = 'expired-token'
@@ -123,6 +135,19 @@ async def refresh(configuration):
             self.requests[-1][0].headers['Authorization'],
         )
 
+    async def test_client_go_retry_retries_get_retry_after_response(self):
+        self.configuration.client_go_retries = True
+        self.configuration.retries = 1
+        self.responses = [
+            ({'message': 'retry later'}, 429, {'Retry-After': '0'}),
+            (self.response, 200, {}),
+        ]
+
+        namespaces = await CoreV1Api(self.api_client).list_namespace()
+
+        self.assertEqual([], namespaces.items)
+        self.assertEqual(2, len(self.requests))
+
     async def test_delete_job_accepts_job_and_status_responses(self):
         responses = (
             (

kubernetes/client/configuration.py

@@ -357,6 +357,22 @@ def __init__(
         self.retries = retries
         """Retry configuration
         """
+        self.client_go_retries = False
+        """Enable Kubernetes client-go-compatible retry semantics.
+
+           When enabled, GET and HEAD requests retry Retry-After responses.
+           The retry ceiling is read from ``retries`` when set; otherwise it
+           follows the client-go default of at most 10 retries.
+        """
+        self.client_go_retry_backoff = None
+        """Backoff for Kubernetes client-go-compatible retries.
+
+           If unset, client-go-compatible GET and HEAD retries use the
+           client-go default retry ceiling with no additional client-side
+           delay beyond Retry-After. When set, ``retries`` still overrides
+           the retry ceiling if it is not None.
+        """
+
         # Enable client side validation
         self.client_side_validation = client_side_validation
 

kubernetes/client/rest.py

@@ -20,8 +20,14 @@
 from urllib.parse import urlparse
 
 import urllib3
+from urllib3.util.retry import Retry
 
 from kubernetes.client.exceptions import ApiException, ApiValueError
+from kubernetes.utils.retry import (
+    is_retry_after_response,
+    on_retry_after_error,
+    retry_after_backoff,
+)
 
 SUPPORTED_SOCKS_PROXIES = {"socks5", "socks5h", "socks4", "socks4a"}
 RESTResponseType = urllib3.HTTPResponse
@@ -105,6 +111,8 @@ def getheader(self, name, default=None):
 class RESTClientObject:
 
     def __init__(self, configuration) -> None:
+        self.configuration = configuration
+
         # urllib3.PoolManager will pass all kw parameters to connectionpool
         # https://github.com/shazow/urllib3/blob/f9409436f83aeb79fbaf090181cd81b784f1b8ce/urllib3/poolmanager.py#L75  # noqa: E501
         # https://github.com/shazow/urllib3/blob/f9409436f83aeb79fbaf090181cd81b784f1b8ce/urllib3/connectionpool.py#L680  # noqa: E501
@@ -217,6 +225,30 @@ def request(
                     read=_request_timeout[1]
                 )
 
+        client_go_retries = getattr(self.configuration, 'client_go_retries', False)
+        request_retries = None
+        if client_go_retries:
+            request_retries = self._urllib3_retries_without_status(
+                getattr(self.configuration, 'retries', None))
+
+        def pool_request(*args, **kwargs):
+            if request_retries is not None:
+                kwargs['retries'] = request_retries
+            return self.pool_manager.request(*args, **kwargs)
+
+        def read_request(check_retry_status=False):
+            response = pool_request(
+                method,
+                url,
+                fields={},
+                timeout=timeout,
+                headers=headers,
+                preload_content=False
+            )
+            if check_retry_status:
+                self._raise_retry_after_response(response)
+            return response
+
         try:
             # For `POST`, `PUT`, `PATCH`, `OPTIONS`, `DELETE`
             if method in ['POST', 'PUT', 'PATCH', 'OPTIONS', 'DELETE']:
@@ -247,7 +279,7 @@ def request(
                     request_body = None
                     if body is not None:
                         request_body = json.dumps(body)
-                    r = self.pool_manager.request(
+                    r = pool_request(
                         method,
                         url,
                         body=request_body,
@@ -256,7 +288,7 @@ def request(
                         preload_content=False
                     )
                 elif content_type == 'application/x-www-form-urlencoded':
-                    r = self.pool_manager.request(
+                    r = pool_request(
                         method,
                         url,
                         fields=post_params,
@@ -272,7 +304,7 @@ def request(
                     del headers['Content-Type']
                     # Ensures that dict objects are serialized
                     post_params = [(a, json.dumps(b)) if isinstance(b, dict) else (a,b) for a, b in post_params]
-                    r = self.pool_manager.request(
+                    r = pool_request(
                         method,
                         url,
                         fields=post_params,
@@ -285,7 +317,7 @@ def request(
                 # other content types than JSON when `body` argument is
                 # provided in serialized form.
                 elif isinstance(body, str) or isinstance(body, bytes):
-                    r = self.pool_manager.request(
+                    r = pool_request(
                         method,
                         url,
                         body=body,
@@ -295,7 +327,7 @@ def request(
                     )
                 elif headers['Content-Type'].startswith('text/') and isinstance(body, bool):
                     request_body = "true" if body else "false"
-                    r = self.pool_manager.request(
+                    r = pool_request(
                         method,
                         url,
                         body=request_body,
@@ -310,16 +342,57 @@ def request(
                     raise ApiException(status=0, reason=msg)
             # For `GET`, `HEAD`
             else:
-                r = self.pool_manager.request(
-                    method,
-                    url,
-                    fields={},
-                    timeout=timeout,
-                    headers=headers,
-                    preload_content=False
-                )
+                if client_go_retries:
+                    backoff = retry_after_backoff(
+                        getattr(self.configuration, 'retries', None),
+                        getattr(self.configuration, 'client_go_retry_backoff', None),
+                    )
+                    r = on_retry_after_error(
+                        backoff, self._is_read_retryable,
+                        lambda: read_request(True))
+                else:
+                    r = read_request()
         except urllib3.exceptions.SSLError as e:
             msg = "\n".join([type(e).__name__, str(e)])
             raise ApiException(status=0, reason=msg)
 
         return RESTResponse(r)
+
+    @classmethod
+    def _is_read_retryable(cls, error):
+        return is_retry_after_response(error)
+
+    @staticmethod
+    def _retry_after_error(response):
+        error = ApiException(status=response.status, reason=response.reason)
+        error.headers = response.getheaders()
+        return error
+
+    @classmethod
+    def _raise_retry_after_response(cls, response):
+        error = cls._retry_after_error(response)
+        if not is_retry_after_response(error):
+            return
+        try:
+            response.drain_conn()
+        except Exception:
+            response.release_conn()
+        raise error
+
+    @staticmethod
+    def _urllib3_retries_without_status(retries):
+        if retries is False:
+            return False
+        if retries is None:
+            retries = Retry.DEFAULT
+        elif retries is True:
+            retries = Retry.DEFAULT
+        elif isinstance(retries, int):
+            retries = Retry.from_int(retries)
+        if isinstance(retries, Retry):
+            return retries.new(
+                status=0,
+                status_forcelist=(),
+                respect_retry_after_header=False,
+            )
+        return retries

kubernetes/test/test_generated_api.py

@@ -80,10 +80,25 @@ def do_DELETE(self):
         self._respond()
 
     def _respond(self):
+        request_index = self.server.request_count
+        self.server.request_count += 1
         response = self.server.response_body
-        self.send_response(self.server.response_status)
+        response_bodies = getattr(self.server, 'response_bodies', None)
+        if response_bodies is not None:
+            response = response_bodies[min(request_index, len(response_bodies) - 1)]
+        status = self.server.response_status
+        response_statuses = getattr(self.server, 'response_statuses', None)
+        if response_statuses is not None:
+            status = response_statuses[min(request_index, len(response_statuses) - 1)]
+        response_headers = self.server.response_headers
+        response_header_sets = getattr(self.server, 'response_header_sets', None)
+        if response_header_sets is not None:
+            response_headers = response_header_sets[min(request_index, len(response_header_sets) - 1)]
+        self.send_response(status)
         self.send_header('Content-Type', 'application/json')
         self.send_header('Content-Length', str(len(response)))
+        for key, value in response_headers.items():
+            self.send_header(key, value)
         self.end_headers()
         self.wfile.write(response)
 
@@ -386,6 +401,8 @@ def setUp(self):
             ('127.0.0.1', 0), _RecordingHandler)
         self.server.response_body = b'{}'
         self.server.response_status = 200
+        self.server.response_headers = {}
+        self.server.request_count = 0
         self.thread = threading.Thread(target=self.server.serve_forever)
         self.thread.start()
         self.api_client = ApiClient(Configuration(
@@ -411,6 +428,24 @@ def test_continue_parameter_keeps_public_name(self):
             self.server.request_path,
         )
 
+    def test_client_go_retry_retries_get_retry_after_response(self):
+        self.api_client.configuration.client_go_retries = True
+        self.api_client.configuration.retries = 1
+        self.server.response_statuses = [429, 200]
+        self.server.response_bodies = [
+            b'{"message": "retry later"}',
+            b'{"items": []}',
+        ]
+        self.server.response_header_sets = [
+            {'Retry-After': '0'},
+            {},
+        ]
+
+        namespaces = CoreV1Api(self.api_client).list_namespace()
+
+        self.assertEqual([], namespaces.items)
+        self.assertEqual(2, self.server.request_count)
+
     def test_deferred_api_validation_rejects_invalid_first_call(self):
         with self.assertRaises(ValidationError) as raised:
             CoreV1Api(self.api_client).list_namespace(limit='invalid')