Allow custom error handling (#1003)

Author: gwauge

docs/source/index.rst

@@ -154,6 +154,67 @@ Python's standard `logging framework`_.
 
 .. _logging framework: https://docs.python.org/3/library/logging.html
 
+Errors
+------
+
+The ``s3fs`` library includes a built-in mechanism to automatically retry
+operations when specific transient errors occur. You can customize this behavior
+by adding specific exception types or defining complex logic via custom handlers.
+
+Default Retryable Errors
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+By default, ``s3fs`` will retry the following exception types:
+
+- ``socket.timeout``
+- ``HTTPClientError``
+- ``IncompleteRead``
+- ``FSTimeoutError``
+- ``ResponseParserError``
+- ``aiohttp.ClientPayloadError`` (if available)
+
+Registering Custom Error Types
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To include additional exception types in the default retry logic, use the
+``add_retryable_error`` function. This is useful for simple type-based retries.
+
+.. code-block:: python
+
+    >>> class MyCustomError(Exception):
+            pass
+    >>> s3fs.add_retryable_error(MyCustomError)
+
+Implementing Custom Error Handlers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+For more complex scenarios, such as retrying based on an error message rather than
+just the type, you can register a custom error handler using ``set_custom_error_handler``.
+
+The handler should be a callable that accepts an exception instance and returns ``True``
+if the error should be retried, or ``False`` otherwise.
+
+.. code-block:: python
+
+    >>> def my_handler(e):
+            return isinstance(e, MyCustomError) and "some condition" in str(e)
+    >>> s3fs.set_custom_error_handler(my_handler)
+
+Handling AWS ClientErrors
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+``s3fs`` provides specialized handling for ``botocore.exceptions.ClientError``.
+While ``s3fs`` checks these against internal patterns (like throttling),
+you can extend this behavior using a custom handler. Note that the internal
+patterns will still be checked and handled before the custom handler.
+
+.. code-block:: python
+
+    >>> def another_handler(e):
+            return isinstance(e, ClientError) and "Throttling" in str(e)
+    >>> s3fs.set_custom_error_handler(another_handler)
+
+
 Credentials
 -----------
 

s3fs/__init__.py

@@ -1,4 +1,4 @@
-from .core import S3FileSystem, S3File
+from .core import S3FileSystem, S3File, add_retryable_error, set_custom_error_handler
 from .mapping import S3Map
 
 from ._version import get_versions

s3fs/core.py

@@ -73,6 +73,56 @@ def setup_logging(level=None):
 if ClientPayloadError is not None:
     S3_RETRYABLE_ERRORS += (ClientPayloadError,)
 
+
+def add_retryable_error(exc):
+    """
+    Add an exception type to the list of retryable S3 errors.
+
+    Parameters
+    ----------
+    exc : Exception
+        The exception type to add to the retryable errors.
+
+    Examples
+    ----------
+    >>> class MyCustomError(Exception):  # doctest: +SKIP
+    ...     pass  # doctest: +SKIP
+    >>> add_retryable_error(MyCustomError)  # doctest: +SKIP
+    """
+    global S3_RETRYABLE_ERRORS
+    S3_RETRYABLE_ERRORS += (exc,)
+
+
+CUSTOM_ERROR_HANDLER = lambda _: False
+
+
+def set_custom_error_handler(func):
+    """Set a custom error handler function for S3 retryable errors.
+
+    The function should take an exception instance as its only argument,
+    and return True if the operation should be retried, or False otherwise.
+    This can also be used for custom behavior on `ClientError` exceptions,
+    such as retrying other patterns.
+
+    Parameters
+    ----------
+    func : callable[[Exception], bool]
+        The custom error handler function.
+
+    Examples
+    ----------
+    >>> def my_handler(e):  # doctest: +SKIP
+    ...     return isinstance(e, MyCustomError) and "some condition" in str(e)  # doctest: +SKIP
+    >>> set_custom_error_handler(my_handler)  # doctest: +SKIP
+
+    >>> def another_handler(e):  # doctest: +SKIP
+    ...     return isinstance(e, ClientError) and "Throttling" in str(e)"  # doctest: +SKIP
+    >>> set_custom_error_handler(another_handler)  # doctest: +SKIP
+    """
+    global CUSTOM_ERROR_HANDLER
+    CUSTOM_ERROR_HANDLER = func
+
+
 _VALID_FILE_MODES = {"r", "w", "a", "rb", "wb", "ab"}
 
 _PRESERVE_KWARGS = [
@@ -110,29 +160,46 @@ def setup_logging(level=None):
 async def _error_wrapper(func, *, args=(), kwargs=None, retries):
     if kwargs is None:
         kwargs = {}
+    err = None
     for i in range(retries):
+        wait_time = min(1.7**i * 0.1, 15)
+
         try:
             return await func(*args, **kwargs)
         except S3_RETRYABLE_ERRORS as e:
             err = e
             logger.debug("Retryable error: %s", e)
-            await asyncio.sleep(min(1.7**i * 0.1, 15))
+            await asyncio.sleep(wait_time)
         except ClientError as e:
             logger.debug("Client error (maybe retryable): %s", e)
             err = e
-            wait_time = min(1.7**i * 0.1, 15)
-            if "SlowDown" in str(e):
-                await asyncio.sleep(wait_time)
-            elif "reduce your request rate" in str(e):
-                await asyncio.sleep(wait_time)
-            elif "XAmzContentSHA256Mismatch" in str(e):
+
+            matched = False
+            for pattern in [
+                "SlowDown",
+                "reduce your request rate",
+                "XAmzContentSHA256Mismatch",
+            ]:
+                if pattern in str(e):
+                    matched = True
+                    break
+
+            if matched:
                 await asyncio.sleep(wait_time)
             else:
-                break
+                should_retry = CUSTOM_ERROR_HANDLER(e)
+                if should_retry:
+                    await asyncio.sleep(wait_time)
+                else:
+                    break
         except Exception as e:
-            logger.debug("Nonretryable error: %s", e)
             err = e
-            break
+            should_retry = CUSTOM_ERROR_HANDLER(e)
+            if should_retry:
+                await asyncio.sleep(wait_time)
+            else:
+                logger.debug("Nonretryable error: %s", e)
+                break
 
     if "'coroutine'" in str(err):
         # aiobotocore internal error - fetch original botocore error