Retry 409 write.contention in the v1 and v2 clients (aito-core #1619)

aito/_write_contention.py

@@ -0,0 +1,62 @@
+"""The retry policy for a starved write: 409 ``write.contention`` (aito-core #1619)
+
+A large write that loses to concurrent writes on the same table gives up without
+committing anything and answers 409 with the code ``write.contention`` and a
+``Retry-After`` header. Nothing was written, so the request is safe to repeat; the v1 and
+v2 clients both do, by this one policy:
+
+- keyed on the CODE, never on 409 alone: the server's other 409s (an env being migrated, an
+  old binary format that needs a repair, a release or branch conflict) are other conditions;
+- the server's Retry-After (seconds; 1 when absent or not a number) is the floor, plus full
+  jitter up to ``BASE * 2**n`` seconds (at most ``CAP``), so clients starved by the same
+  churn do not all retry on the same beat;
+- a capped number of retries, set per client.
+
+Standard library only.
+"""
+
+import random
+from typing import Any, Optional
+
+#: the error code: v1 carries it in the body's ``error``, v2 in ``data.code``
+WRITE_CONTENTION = 'write.contention'
+#: full jitter on top of Retry-After: up to BASE * 2**n seconds, at most CAP
+BACKOFF_BASE = 0.5
+BACKOFF_CAP = 8.0
+#: the clients' default number of retries
+DEFAULT_RETRIES = 3
+
+
+def code_of(body: Any) -> Optional[str]:
+    """the error code in a parsed error body, v1 (``error``) or v2 (``data.code``)"""
+    if not isinstance(body, dict):
+        return None
+    data = body.get('data')
+    if isinstance(data, dict) and isinstance(data.get('code'), str):
+        return data['code']
+    error = body.get('error')
+    return error if isinstance(error, str) else None
+
+
+def retry_after(headers: Any) -> float:
+    """the Retry-After header in seconds; 1 when absent or not a number (an HTTP date included)"""
+    raw = None
+    for name, value in (headers or {}).items():
+        if str(name).lower() == 'retry-after':
+            raw = value
+    try:
+        return max(0.0, float(raw))
+    except (TypeError, ValueError):
+        return 1.0
+
+
+def delay(headers: Any, retry: int) -> float:
+    """seconds to wait before retry number ``retry`` (0-based)"""
+    return retry_after(headers) + random.uniform(0.0, min(BACKOFF_CAP, BACKOFF_BASE * 2 ** retry))
+
+
+def check_retries(retries: int) -> int:
+    """a client's write_contention_retries setting, validated"""
+    if retries < 0:
+        raise ValueError(f'write_contention_retries must be 0 or more, got {retries}')
+    return retries

aito/v1/client.py

@@ -8,6 +8,7 @@
 
 import asyncio
 import logging
+import time
 import warnings
 from typing import Dict, List, Union, Tuple, Optional, TYPE_CHECKING
 
@@ -16,6 +17,7 @@
 if TYPE_CHECKING:  # pragma: no cover - import-time typing only
     from aiohttp import ClientSession
 
+from aito import _write_contention as contention
 from aito.exceptions import BaseError
 from aito.utils._generic_utils import instance_url_problem
 from .requests import AitoRequest, BaseRequest, GetVersionRequest
@@ -51,13 +53,16 @@ class RequestError(Error):
     """An error occurred when sending a request to the Aito instance
 
     """
-    def __init__(self, request_obj: AitoRequest, error: Exception):
+    def __init__(self, request_obj: AitoRequest, error: Exception, retries: int = 0):
         """
 
         :param request_obj: the request object
         :type request_obj: AitoRequest
         :param error: the error
         :type error: Exception
+        :param retries: how many times the request was repeated before this error (a 409
+            ``write.contention`` is retried), defaults to 0
+        :type retries: int
         """
         self.request_obj = request_obj
         self.error = error
@@ -72,7 +77,9 @@ def __init__(self, request_obj: AitoRequest, error: Exception):
             error_msg = error.message
         else:
             error_msg = str(error)
-        super().__init__(f'failed to {request_obj}: {error_msg}')
+        self.retries = retries
+        suffix = f' (retried {retries} times)' if retries else ''
+        super().__init__(f'failed to {request_obj}: {error_msg}{suffix}')
 
 
 class AitoClient:
@@ -101,7 +108,8 @@ def __init__(
             instance_url: str,
             api_key: str,
             check_credentials: bool = True,
-            raise_for_status: bool = True
+            raise_for_status: bool = True,
+            write_contention_retries: int = contention.DEFAULT_RETRIES
     ):
         """
 
@@ -113,7 +121,13 @@ def __init__(
         :type check_credentials: bool
         :param raise_for_status: automatically raise RequestError for each failed response, defaults to True
         :type raise_for_status: bool
+        :param write_contention_retries: how many times to repeat a request the server answered with 409
+            ``write.contention``: a write that lost to concurrent writes on the same table and committed
+            nothing, so repeating it is safe. Each retry waits the server's ``Retry-After`` plus exponential
+            jitter (at most 8 s on top). 0 turns it off. Other 409s are never retried. Defaults to 3
+        :type write_contention_retries: int
         :raises BaseError: an error occurred during the creation of AitoClient
+        :raises ValueError: write_contention_retries is negative
 
         >>> aito_client = AitoClient(your_instance_url, your_api_key) # doctest: +SKIP
         >>> # Change the API key to READ-WRITE or READ-ONLY
@@ -126,6 +140,7 @@ def __init__(
         self.instance_url = instance_url.strip("/")
         self.api_key = api_key
         self.raise_for_status = raise_for_status
+        self.write_contention_retries = contention.check_retries(write_contention_retries)
         self.instance_version = None
         if check_credentials:
             try:
@@ -272,23 +287,43 @@ def request(
                 request_obj = AitoRequest.make_request(method, endpoint, query)
             else:
                 raise TypeError("'request() requires either 'request_obj' or 'method' and 'endpoint'")
-        try:
-            resp = requestslib.request(
-                method=request_obj.method,
-                url=self.instance_url + request_obj.endpoint,
-                headers=self.headers,
-                json=request_obj.query
-            )
-            resp.raise_for_status()
-            json_resp = resp.json()
-        except Exception as e:
-            req_err = RequestError(request_obj, e)
-            _raise = raise_for_status if raise_for_status is not None else self.raise_for_status
-            if _raise:
-                raise req_err
-            else:
-                return req_err
-        return request_obj.response_cls(json_resp)
+        retries = 0
+        while True:
+            try:
+                resp = requestslib.request(
+                    method=request_obj.method,
+                    url=self.instance_url + request_obj.endpoint,
+                    headers=self.headers,
+                    json=request_obj.query
+                )
+                # the body is only parsed here for a 409, so a normal response is parsed once
+                if resp.status_code == 409 and self._is_write_contention(
+                        resp.status_code, _json_or_none(resp), retries):
+                    delay = contention.delay(resp.headers, retries)
+                    retries += 1
+                    LOG.info(f'{request_obj}: write.contention, retry {retries} of '
+                             f'{self.write_contention_retries} in {delay:.1f}s')
+                    time.sleep(delay)
+                    continue
+                resp.raise_for_status()
+                json_resp = resp.json()
+            except Exception as e:
+                req_err = RequestError(request_obj, e, retries)
+                _raise = raise_for_status if raise_for_status is not None else self.raise_for_status
+                if _raise:
+                    raise req_err
+                else:
+                    return req_err
+            return request_obj.response_cls(json_resp)
+
+    def _is_write_contention(self, status: int, body, retries: int) -> bool:
+        """a 409 whose code is write.contention, with retries left (aito-core #1619)
+
+        Keyed on the code: the server's other 409s (an env being migrated, an old binary
+        format that needs a repair, a release or branch conflict) are other conditions.
+        """
+        return (status == 409 and retries < self.write_contention_retries
+                and contention.code_of(body) == contention.WRITE_CONTENTION)
 
     async def async_request(
             self, session: 'ClientSession', *,
@@ -321,22 +356,38 @@ async def async_request(
                 request_obj = AitoRequest.make_request(method, endpoint, query)
             else:
                 raise TypeError("request() requires either 'request_obj' or 'method' and 'endpoint'")
-        try:
-            async with session.request(
-                    method=request_obj.method,
-                    url=self.instance_url + request_obj.endpoint,
-                    json=request_obj.query,
-                    headers=self.headers,
-                    raise_for_status=True
-            ) as resp:
-                return request_obj.response_cls(await resp.json())
-        except Exception as e:
-            req_err = RequestError(request_obj, e)
-            _raise = raise_for_status if raise_for_status is not None else self.raise_for_status
-            if _raise:
-                raise req_err
-            else:
-                return req_err
+        retries = 0
+        while True:
+            delay = None
+            try:
+                async with session.request(
+                        method=request_obj.method,
+                        url=self.instance_url + request_obj.endpoint,
+                        json=request_obj.query,
+                        headers=self.headers,
+                        raise_for_status=False
+                ) as resp:
+                    if resp.status == 409:
+                        try:
+                            body = await resp.json(content_type=None)
+                        except Exception:
+                            body = None
+                        if self._is_write_contention(resp.status, body, retries):
+                            delay = contention.delay(resp.headers, retries)
+                    if delay is None:
+                        resp.raise_for_status()
+                        return request_obj.response_cls(await resp.json())
+            except Exception as e:
+                req_err = RequestError(request_obj, e, retries)
+                _raise = raise_for_status if raise_for_status is not None else self.raise_for_status
+                if _raise:
+                    raise req_err
+                else:
+                    return req_err
+            retries += 1
+            LOG.info(f'async {request_obj}: write.contention, retry {retries} of '
+                     f'{self.write_contention_retries} in {delay:.1f}s')
+            await asyncio.sleep(delay)
 
     async def bounded_async_request(
             self, semaphore: asyncio.Semaphore, *args, **kwargs
@@ -424,3 +475,11 @@ async def run():
         loop = asyncio.get_event_loop()
         responses = loop.run_until_complete(run())
         return responses
+
+
+def _json_or_none(resp) -> Optional[Union[Dict, List]]:
+    """the parsed JSON body, or None when there is none"""
+    try:
+        return resp.json()
+    except ValueError:
+        return None

aito/v2/client.py

@@ -18,10 +18,12 @@
 """
 
 import logging
+import time
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import requests as requestslib
 
+from aito import _write_contention as contention
 from aito.local.profiles import resolve_credentials
 from aito.utils._generic_utils import instance_url_problem
 from .errors import AitoV2Error
@@ -32,6 +34,10 @@
 
 LOG = logging.getLogger('AitoClientV2')
 
+#: The error code of a write that gave up under concurrent same-table writes without
+#: committing anything (aito-core #1619): safe to repeat, so the client retries it.
+WRITE_CONTENTION = contention.WRITE_CONTENTION
+
 #: environment names may not start with these — the engine reserves them
 _RESERVED_ENV_PREFIXES = ('_', 'env.', 'release.')
 
@@ -71,8 +77,15 @@ class AitoClientV2:
         mostly network. Keep the callback cheap and non-throwing: it runs
         inline, and an exception in it would surface as a failed request
     :type on_response: Optional[Callable[[Any, str], None]]
+    :param write_contention_retries: how many times to repeat a request the server answered
+        with 409 ``write.contention``: a write that lost to concurrent writes on the same table
+        and committed nothing, so repeating it is safe. Each retry waits the server's
+        ``Retry-After`` plus exponential jitter (at most 8 s on top). 0 turns it off; the
+        error is then raised at once. Other 409s are never retried.
+    :type write_contention_retries: int
     :raises ValueError: the instance URL has no scheme or host, the environment name is
-        one the engine reserves, or no instance URL and key could be resolved
+        one the engine reserves, no instance URL and key could be resolved, or
+        ``write_contention_retries`` is negative
         (``NoCredentialsError``)
     :raises AitoV2Error: the credentials could not be verified
 
@@ -93,7 +106,9 @@ def __init__(
             timeout: float = 30.0,
             check_credentials: bool = True,
             on_response: Optional[Callable[[Any, str], None]] = None,
+            write_contention_retries: int = contention.DEFAULT_RETRIES,
     ):
+        contention.check_retries(write_contention_retries)
         if on_warning not in _ON_WARNING_CHOICES:
             raise ValueError(
                 f"invalid on_warning '{on_warning}', expected one of {'|'.join(_ON_WARNING_CHOICES)}")
@@ -113,6 +128,7 @@ def __init__(
         self.on_warning = on_warning
         self.timeout = timeout
         self.on_response = on_response
+        self.write_contention_retries = write_contention_retries
         # A pooled session keeps the TCP+TLS connection alive across calls. A
         # fresh connection per request pays the handshake every time, which
         # dominates the per-call wall-clock against a shared instance.
@@ -200,30 +216,47 @@ def request(
         if data is not None:
             headers = {**headers, 'Content-Type': content_type or 'application/octet-stream'}
             body = {'data': data}
-        try:
-            resp = self._session.request(
-                method=method, url=url, params=params, headers=headers,
-                timeout=self.timeout if timeout is None else timeout, **body,
-            )
-        except requestslib.RequestException as e:
-            hint = ''
-            if isinstance(e, requestslib.ConnectionError) and any(
-                    h in self.instance_url for h in ('://127.0.0.1', '://localhost')):
-                hint = ' (a local Aito: is it running? `aito status`, `aito start`)'
-            raise AitoV2Error(f'Aito v2 request failed: {method} {path}: {e}{hint}') from e
-
-        if self.on_response is not None:
-            # Before the status check, so a caller observing timings still sees
-            # the calls that failed — those are the ones worth looking at.
-            self.on_response(resp, path)
-
-        try:
-            parsed = resp.json()
-        except ValueError:
-            parsed = None
-
-        if resp.status_code >= 400:
-            raise AitoV2Error.from_response(resp.status_code, resp.text, parsed)
+        retries = 0
+        while True:
+            try:
+                resp = self._session.request(
+                    method=method, url=url, params=params, headers=headers,
+                    timeout=self.timeout if timeout is None else timeout, **body,
+                )
+            except requestslib.RequestException as e:
+                hint = ''
+                if isinstance(e, requestslib.ConnectionError) and any(
+                        h in self.instance_url for h in ('://127.0.0.1', '://localhost')):
+                    hint = ' (a local Aito: is it running? `aito status`, `aito start`)'
+                raise AitoV2Error(f'Aito v2 request failed: {method} {path}: {e}{hint}') from e
+
+            if self.on_response is not None:
+                # Before the status check, so a caller observing timings still sees
+                # the calls that failed — those are the ones worth looking at.
+                self.on_response(resp, path)
+
+            try:
+                parsed = resp.json()
+            except ValueError:
+                parsed = None
+
+            if resp.status_code < 400:
+                break
+            error = AitoV2Error.from_response(resp.status_code, resp.text, parsed)
+            # Keyed on the code, never on 409 alone: the server's other 409s (an env being
+            # migrated, an old binary format that needs a repair, a release or branch
+            # conflict) are different conditions, and some must never be repeated.
+            if error.code != WRITE_CONTENTION:
+                raise error
+            if retries >= self.write_contention_retries:
+                if retries:
+                    error.args = (f'{error.args[0]} (retried {retries} times)',)
+                raise error
+            delay = contention.delay(getattr(resp, 'headers', None), retries)
+            retries += 1
+            LOG.info('%s %s: write.contention, retry %d of %d in %.1fs',
+                     method, path, retries, self.write_contention_retries, delay)
+            time.sleep(delay)
         if parsed is None:
             raise AitoV2Error(
                 f'Aito v2 returned a non-JSON body for {method} {path}: {resp.text[:200]}',