Make HTTP retry behavior configurable

planet/cli/cli.py

@@ -21,6 +21,7 @@
 import planet_auth_utils
 import planet
 from planet.cli import mosaics
+from planet.http import MAX_RETRIES, MAX_RETRY_BACKOFF, MAX_RETRY_JITTER
 
 from . import auth, cmds, collect, data, destinations, orders, subscriptions, features
 
@@ -42,21 +43,44 @@
 @planet_auth_utils.opt_client_id()
 @planet_auth_utils.opt_client_secret()
 @planet_auth_utils.opt_api_key()
+@click.option('--max-retries',
+              type=int,
+              default=MAX_RETRIES,
+              show_default=True,
+              help='Maximum number of retries of a retryable request. When '
+              'set to 0, requests are not retried.')
+@click.option('--max-retry-backoff',
+              type=float,
+              default=MAX_RETRY_BACKOFF,
+              show_default=True,
+              help='Maximum time, in seconds, to wait between retries.')
+@click.option('--max-retry-jitter',
+              type=float,
+              default=MAX_RETRY_JITTER,
+              show_default=True,
+              help='Maximum random time, in seconds, added to the wait '
+              'between retries. When set to 0, no jitter is added.')
 @cmds.translate_exceptions
 def main(ctx,
          verbosity,
          quiet,
          auth_profile,
          auth_client_id,
          auth_client_secret,
-         auth_api_key):
+         auth_api_key,
+         max_retries,
+         max_retry_backoff,
+         max_retry_jitter):
     """Planet SDK for Python CLI"""
     _configure_logging(verbosity)
 
     # ensure that ctx.obj exists and is a dict (in case `cli()` is called
     # by means other than the `if` block below)
     ctx.ensure_object(dict)
     ctx.obj['QUIET'] = quiet
+    ctx.obj['MAX_RETRIES'] = max_retries
+    ctx.obj['MAX_RETRY_BACKOFF'] = max_retry_backoff
+    ctx.obj['MAX_RETRY_JITTER'] = max_retry_jitter
 
     _configure_cli_auth_ctx(ctx,
                             auth_profile,

planet/cli/session.py

@@ -7,13 +7,22 @@ class CliSession(Session):
     """Session with CLI-specific auth and identifying header"""
 
     def __init__(self, click_ctx=None, plsdk_auth=None):
+        _plsdk_auth = None
+        _max_retries = None
+        _max_retry_backoff = None
+        _max_retry_jitter = None
+
         if click_ctx:
             _plsdk_auth = click_ctx.obj['PLSDK_AUTH']
-        else:
-            _plsdk_auth = None
+            _max_retries = click_ctx.obj.get('MAX_RETRIES')
+            _max_retry_backoff = click_ctx.obj.get('MAX_RETRY_BACKOFF')
+            _max_retry_jitter = click_ctx.obj.get('MAX_RETRY_JITTER')
 
         if plsdk_auth:
             _plsdk_auth = plsdk_auth
 
-        super().__init__(_plsdk_auth)
+        super().__init__(_plsdk_auth,
+                         max_retries=_max_retries,
+                         max_retry_backoff=_max_retry_backoff,
+                         max_retry_jitter=_max_retry_jitter)
         self._client.headers.update({'X-Planet-App': 'python-cli'})

planet/http.py

@@ -49,6 +49,7 @@
 ]
 MAX_RETRIES = 5
 MAX_RETRY_BACKOFF = 64  # seconds
+MAX_RETRY_JITTER = 1  # seconds
 
 DEFAULT_READ_TIMEOUT_SECS = 125.0
 RATE_LIMIT = 10  # per second
@@ -226,26 +227,61 @@ class Session(BaseSession):
     ...
     >>> asyncio.run(main())
 
+    ```
+
+    Example:
+    ```python
+    >>> import asyncio
+    >>> from planet import Session
+    >>>
+    >>> async def main():
+    ...     # customize the retry behavior
+    ...     async with Session(max_retries=10,
+    ...                        max_retry_backoff=32,
+    ...                        max_retry_jitter=4) as sess:
+    ...         # communicate with services here
+    ...         pass
+    ...
+    >>> asyncio.run(main())
+
     ```
     """
 
     def __init__(
         self,
         auth: Optional[AuthType] = None,
         read_timeout_secs: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        max_retry_backoff: Optional[float] = None,
+        max_retry_jitter: Optional[float] = None,
     ):
         """Initialize a Session.
 
         Parameters:
             auth: Planet server authentication.
             read_timeout_secs: Maximum time to wait for data to be received.
+            max_retries: Maximum number of retries of a retryable request.
+                Zero disables retry.
+            max_retry_backoff: Maximum time, in seconds, to wait between
+                retries.
+            max_retry_jitter: Maximum random time, in seconds, added to the
+                wait between retries. Zero disables jitter.
         """
         if auth is None:
             auth = Auth.from_user_default_session()
 
         if read_timeout_secs is None:
             read_timeout_secs = DEFAULT_READ_TIMEOUT_SECS
 
+        if max_retries is None:
+            max_retries = MAX_RETRIES
+
+        if max_retry_backoff is None:
+            max_retry_backoff = MAX_RETRY_BACKOFF
+
+        if max_retry_jitter is None:
+            max_retry_jitter = MAX_RETRY_JITTER
+
         LOGGER.info(
             f'Session read timeout set to {read_timeout_secs} seconds.')
         timeout = httpx.Timeout(10.0, read=read_timeout_secs)
@@ -270,8 +306,14 @@ async def alog_response(*args, **kwargs):
             alog_response, self._raise_for_status
         ]
 
-        self.max_retries = MAX_RETRIES
-        self.max_retry_backoff = MAX_RETRY_BACKOFF
+        self.max_retries = max_retries
+        self.max_retry_backoff = max_retry_backoff
+        self.max_retry_jitter = max_retry_jitter
+
+        LOGGER.debug(f'Session retry set to a maximum of {self.max_retries} '
+                     f'retries with a maximum backoff of '
+                     f'{self.max_retry_backoff} seconds and a maximum jitter '
+                     f'of {self.max_retry_jitter} seconds.')
 
         self._limiter = _Limiter(rate_limit=RATE_LIMIT, max_workers=MAX_ACTIVE)
         self.outcomes: Counter[str] = Counter()
@@ -364,7 +406,9 @@ async def _retry(self, func, *a, **kw):
                         LOGGER.info(f'Try {num_tries}')
                         LOGGER.info(f'Retrying: caught {type(e)}: {e}')
                         wait_time = self._calculate_wait(
-                            num_tries, self.max_retry_backoff)
+                            num_tries,
+                            self.max_retry_backoff,
+                            self.max_retry_jitter)
                         LOGGER.info(f'Retrying: sleeping {wait_time}s')
                         await asyncio.sleep(wait_time)
                 else:
@@ -374,25 +418,32 @@ async def _retry(self, func, *a, **kw):
         return resp
 
     @staticmethod
-    def _calculate_wait(num_tries, max_retry_backoff):
+    def _calculate_wait(num_tries, max_retry_backoff, max_retry_jitter=None):
         """Calculates retry wait
 
         Base wait period is calculated as a exponential based on the number of
-        tries. Then, a random jitter of up to 999ms is added to the base wait
-        to avoid waves of requests in the case of multiple requests. Finally,
-        the wait is thresholded to the maximum retry backoff.
+        tries. The base wait is thresholded to the maximum retry backoff, less
+        room for jitter. Then, a random jitter of up to the maximum retry
+        jitter is added to the base wait to avoid waves of requests in the
+        case of multiple requests.
 
-        Because threshold is applied after jitter, calculations that hit
-        threshold will not have random jitter applied, they will simply result
-        in the threshold value being returned.
+        Because the threshold is applied before jitter, waits that hit the
+        threshold are jittered just like any other wait, and the maximum retry
+        backoff is never exceeded.
 
         Ref:
         * https://docs.planet.com/develop/apis/data/#api-mechanics
         * https://cloud.google.com/iot/docs/how-tos/exponential-backoff
         """
+        if max_retry_jitter is None:
+            max_retry_jitter = MAX_RETRY_JITTER
+
+        # a backoff smaller than the jitter leaves no room for the full
+        # jitter, so the jitter is narrowed to fit within the backoff
+        jitter_secs = min(max_retry_jitter, max_retry_backoff)
+        base_wait = min(2**num_tries, max_retry_backoff - jitter_secs)
         random_number_milliseconds = random.randint(0, 999) / 1000.0
-        calc_wait = 2**num_tries + random_number_milliseconds
-        return min(calc_wait, max_retry_backoff)
+        return base_wait + jitter_secs * random_number_milliseconds
 
     async def request(self,
                       method: str,

tests/integration/test_data_cli.py

@@ -434,6 +434,25 @@ def test_data_search_cmd_item_types(mock_bundles,
         assert result.exit_code == 2
 
 
+@respx.mock
+def test_data_search_cmd_max_retries(mock_bundles):
+    """The --max-retries option is passed through to the session"""
+    route = respx.post(TEST_QUICKSEARCH_URL)
+    route.side_effect = [
+        httpx.Response(HTTPStatus.TOO_MANY_REQUESTS, json={}),
+        httpx.Response(HTTPStatus.OK, json={'features': [{
+            "key": "value"
+        }]})
+    ]
+
+    # retry disabled, so the first response is not retried
+    result = CliRunner().invoke(
+        cli.main, args=['--max-retries', '0', 'data', 'search', 'PSScene'])
+
+    assert result.exit_code == 1
+    assert route.call_count == 1
+
+
 @respx.mock
 @pytest.mark.parametrize("geom_fixture",
                          [('geom_geojson'), ('feature_geojson'),

tests/unit/test_cli_session.py

@@ -14,13 +14,14 @@
 from http import HTTPStatus
 import json
 
+import click
 import httpx
 import respx
 
 import pytest
 
 # from planet.auth import _SecretFile
-from planet import auth
+from planet import auth, http
 from planet.cli import session
 
 TEST_URL = 'mock://mock.com'
@@ -79,3 +80,42 @@ async def test_CliSession_auth_valid(test_valid_secretfile):
         # assert base64.b64decode(credentials) == b'clisessiontest:'
         credentials = received_request.headers['authorization']
         assert credentials == 'api-key clisessiontest'
+
+
+def _click_ctx(**obj):
+    """A click context carrying the given context object entries"""
+    return click.Context(click.Command('test'), obj=obj)
+
+
+@pytest.mark.anyio
+async def test_CliSession_retry_defaults(test_valid_secretfile):
+    """Retry defaults to the module-level configuration"""
+    async with session.CliSession() as sess:
+        assert sess.max_retries == http.MAX_RETRIES
+        assert sess.max_retry_backoff == http.MAX_RETRY_BACKOFF
+        assert sess.max_retry_jitter == http.MAX_RETRY_JITTER
+
+
+@pytest.mark.anyio
+async def test_CliSession_retry_from_ctx(test_valid_secretfile):
+    """Retry configuration is read from the click context"""
+    ctx = _click_ctx(PLSDK_AUTH=auth.Auth.from_key("clisessiontest"),
+                     MAX_RETRIES=2,
+                     MAX_RETRY_BACKOFF=8,
+                     MAX_RETRY_JITTER=2)
+
+    async with session.CliSession(ctx) as sess:
+        assert sess.max_retries == 2
+        assert sess.max_retry_backoff == 8
+        assert sess.max_retry_jitter == 2
+
+
+@pytest.mark.anyio
+async def test_CliSession_retry_ctx_unset(test_valid_secretfile):
+    """Retry falls back to the defaults when the context does not set it"""
+    ctx = _click_ctx(PLSDK_AUTH=auth.Auth.from_key("clisessiontest"))
+
+    async with session.CliSession(ctx) as sess:
+        assert sess.max_retries == http.MAX_RETRIES
+        assert sess.max_retry_backoff == http.MAX_RETRY_BACKOFF
+        assert sess.max_retry_jitter == http.MAX_RETRY_JITTER