Enhance Retry Logic and Configuration for Storage Control API (#787)

Author: Mahalaxmibejugam

cloudbuild/cleanup.sh

@@ -16,6 +16,6 @@ gcloud storage rm --recursive "gs://gcsfs-test-hns-${SHORT_BUILD_ID}" || true &
 gcloud storage rm --recursive "gs://gcsfs-test-zonal-${SHORT_BUILD_ID}" || true &
 gcloud storage rm --recursive "gs://gcsfs-test-standard-for-zonal-${SHORT_BUILD_ID}" || true &
 gcloud storage rm --recursive "gs://gcsfs-test-zonal-core-${SHORT_BUILD_ID}" || true &
-gcloud storage rm --recursive "gs://gcsfs-test-hns-req-pay-${SHORT_BUILD_ID}" || true &
-gcloud storage rm --recursive "gs://gcsfs-test-standard-req-pay-${SHORT_BUILD_ID}" || true &
+gcloud storage rm --recursive "gs://gcsfs-test-hns-req-pay-${SHORT_BUILD_ID}" --billing-project="${PROJECT_ID}" || true &
+gcloud storage rm --recursive "gs://gcsfs-test-standard-req-pay-${SHORT_BUILD_ID}" --billing-project="${PROJECT_ID}" || true &
 wait

cloudbuild/e2e-tests-cloudbuild.yaml

@@ -160,6 +160,7 @@ steps:
     env:
       - "SHORT_BUILD_ID=${_SHORT_BUILD_ID}"
       - "ZONE=${_ZONE}"
+      - "PROJECT_ID=${PROJECT_ID}"
     waitFor:
       [
         "run-standard-tests",

docs/source/retries.rst

@@ -38,31 +38,48 @@ For standard buckets, ``gcsfs`` uses a custom retry decorator (``retry_request``
 Hierarchical Namespace (HNS) Buckets
 ------------------------------------
 
-For HNS buckets, ``ExtendedGcsFileSystem`` utilizes the specialized Storage Control client (``StorageControlAsyncClient``) for folder-level operations (e.g., ``mkdir``, ``rename``).
+For HNS buckets, ``ExtendedGcsFileSystem`` utilizes the specialized Storage Control client (``StorageControlAsyncClient``) for control plane operations (e.g., ``mkdir``, ``rename``, ``get_storage_layout``).
 
-- These calls utilize the underlying Google Cloud Python SDK's default retry behavior. Standard ``gcsfs`` retry logic (``retry_request``) is not applied to these control plane calls.
+- These calls utilize retry configuration based on ``google.api_core.retry.AsyncRetry``.
 - **Applicable Methods:**
     - ``get_storage_layout``: Used to determine bucket type.
     - ``create_folder``: Used for ``mkdir``.
     - ``get_folder``: Used for directory metadata and existence checks.
     - ``list_folders``: Used for directory listings (``ls``).
     - ``rename_folder``: Used for moving/renaming directories (``mv``).
-- **Non-Retried Methods:** Methods like ``delete_folder`` (used for ``rmdir``) are not retried by default.
+    - ``delete_folder``: Used for deleting directories (``rmdir``, ``rm -r``).
 - **Retriable Errors:**
     - ``google.api_core.exceptions.DeadlineExceeded``
+    - ``google.api_core.exceptions.ServiceUnavailable``
     - ``google.api_core.exceptions.InternalServerError``
+    - ``google.api_core.exceptions.TooManyRequests``
     - ``google.api_core.exceptions.ResourceExhausted``
-    - ``google.api_core.exceptions.ServiceUnavailable``
     - ``google.api_core.exceptions.Unknown``
-- **Backoff Strategy:** Exponential backoff with ``initial=1.0s``, ``maximum=60.0s``, and ``multiplier=2.0``.
-- **Overall Timeout (Deadline):** 60.0s
+    - ``google.api_core.exceptions.Unauthenticated`` (when "Invalid Credentials" is in the message).
+
+- **Configuration:**
+  The retry behavior can be customized via the following parameters passed to the FileSystem instance:
+
+  - ``retry_timeout`` (float): The total deadline for the retry loop in seconds. Default: ``60.0``.
+  - ``retry_initial`` (float): The initial delay between retries in seconds. Default: ``1.0``.
+  - ``retry_maximum`` (float): The maximum delay between retries in seconds. Default: ``60.0``.
+  - ``retry_multiplier`` (float): The multiplier applied to the delay after each retry. Default: ``2.0``.
+
+  Per-attempt timeout is controlled by an internal ``STORAGE_CONTROL_RPC_TIMEOUT`` constant, currently set to ``30.0s``.
+
+Configuring Retries via fsspec
+------------------------------
+
+Since ``gcsfs`` integrates with the ``fsspec`` configuration system, these retry parameters can be set using ``fsspec`` `configuration files or environment variables <https://filesystem-spec.readthedocs.io/en/latest/features.html#configuration>`_
+
+These settings will be automatically picked up by any ``GCSFileSystem`` instance when experimental HNS support is enabled (which is the default).
 
 Rapid Storage (Zonal Buckets)
 -----------------------------
 
 For Zonal buckets, ``ZonalFile`` utilizes the specialized gRPC clients (``AsyncMultiRangeDownloader`` for reads and ``AsyncAppendableObjectWriter`` for writes).
 
-- Similar to HNS buckets, control plane operations for Zonal buckets (such as ``get_storage_layout`` or folder operations) utilize the same ``StorageControlAsyncClient`` retry mechanism described in the HNS section above.
+- Similar to HNS buckets, control plane operations for Zonal buckets (such as ``get_storage_layout`` or folder operations) utilize the same Storage Control retry mechanism described in the **Storage Control API** section above.
 - File read/write operations (data plane) for Zonal buckets utilize the underlying Google Cloud Python SDK's default retry behavior for gRPC streams. Standard ``gcsfs`` retry logic (``retry_request``) is not applied to these data plane calls.
 - **AsyncMultiRangeDownloader (MRD) Retries (Reads):**
     - **Applicable Methods:**

gcsfs/extended_gcsfs.py

@@ -20,11 +20,13 @@
 from gcsfs import __version__ as version
 from gcsfs import zb_hns_utils
 from gcsfs.core import GCSFile, GCSFileSystem
+from gcsfs.retry import DEFAULT_RETRY_CONFIG, get_storage_control_retry_config
 from gcsfs.zonal_file import ZonalFile
 
 logger = logging.getLogger("gcsfs")
 
 USER_AGENT = "python-gcsfs"
+STORAGE_CONTROL_RPC_TIMEOUT = 30.0
 
 
 class BucketType(Enum):
@@ -51,6 +53,26 @@ class ExtendedGcsFileSystem(GCSFileSystem):
     """
 
     def __init__(self, *args, finalize_on_close=False, **kwargs):
+        """
+        Parameters
+        ----------
+        finalize_on_close : bool, default False
+            By default, files in zonal buckets are left unfinalized to allow appends.
+        **kwargs : dict
+            Additional arguments passed to GCSFileSystem.
+            Supports retry configuration overrides for Storage Control API:
+            - retry_timeout: Total time to spend retrying (seconds).
+            - retry_initial: Initial delay between retries (seconds).
+            - retry_maximum: Maximum delay between retries (seconds).
+            - retry_multiplier: Multiplier for delay between retries.
+            These map to `google.api_core.retry.AsyncRetry` arguments (without 'retry_' prefix).
+        """
+        valid_keys = DEFAULT_RETRY_CONFIG.keys()
+        self.retry_config = {
+            k[6:]: v
+            for k, v in kwargs.items()
+            if k.startswith("retry_") and k[6:] in valid_keys and v is not None
+        }
         super().__init__(*args, **kwargs)
         # By default, files in zonal buckets are left unfinalized to allow appends.
         self.finalize_on_close = finalize_on_close
@@ -79,6 +101,9 @@ def _user_project(self):
             )
         return None
 
+    def _get_retry_config(self, **kwargs):
+        return get_storage_control_retry_config(self.retry_config, **kwargs)
+
     @property
     def grpc_client(self):
         if self.asynchronous and self._grpc_client is None:
@@ -141,7 +166,11 @@ async def _get_bucket_type(self, bucket):
             client = await self._get_control_plane_client()
             bucket_name_value = f"projects/_/buckets/{bucket}/storageLayout"
             logger.debug(f"get_storage_layout request for name: {bucket_name_value}")
-            response = await client.get_storage_layout(name=bucket_name_value)
+            response = await client.get_storage_layout(
+                name=bucket_name_value,
+                retry=self._get_retry_config(),
+                timeout=STORAGE_CONTROL_RPC_TIMEOUT,
+            )
 
             if response.location_type == "zone":
                 return BucketType.ZONAL_HIERARCHICAL
@@ -514,7 +543,11 @@ async def _mv(self, path1, path2, **kwargs):
 
                 logger.debug(f"rename_folder request: {request}")
                 client = await self._get_control_plane_client()
-                operation = await client.rename_folder(request=request)
+                operation = await client.rename_folder(
+                    request=request,
+                    retry=self._get_retry_config(),
+                    timeout=STORAGE_CONTROL_RPC_TIMEOUT,
+                )
                 await operation.result()
                 self._update_dircache_after_rename(path1, path2)
 
@@ -679,7 +712,11 @@ async def _create_hns_folder(self, path, bucket, key, create_parents):
         try:
             logger.debug(f"create_folder request: {request}")
             client = await self._get_control_plane_client()
-            await client.create_folder(request=request)
+            await client.create_folder(
+                request=request,
+                retry=self._get_retry_config(),
+                timeout=STORAGE_CONTROL_RPC_TIMEOUT,
+            )
             # Instead of invalidating the parent cache, update it to add the new entry.
             parent_path = self._parent(path)
             if parent_path in self.dircache:
@@ -725,7 +762,11 @@ async def _get_directory_info(self, path, bucket, key, generation):
 
                 # Verify existence using get_folder API
                 client = await self._get_control_plane_client()
-                response = await client.get_folder(request=request)
+                response = await client.get_folder(
+                    request=request,
+                    retry=self._get_retry_config(),
+                    timeout=STORAGE_CONTROL_RPC_TIMEOUT,
+                )
 
                 # If successful, return directory metadata
                 return {
@@ -798,7 +839,11 @@ async def _rmdir(self, path):
 
             logger.debug(f"delete_folder request: {request}")
             client = await self._get_control_plane_client()
-            await client.delete_folder(request=request)
+            await client.delete_folder(
+                request=request,
+                retry=self._get_retry_config(),
+                timeout=STORAGE_CONTROL_RPC_TIMEOUT,
+            )
 
             # Remove the directory from the cache and from its parent's listing.
             self.dircache.pop(path, None)
@@ -1121,7 +1166,11 @@ async def _get_all_folders(self, path, bucket, prefix=""):
         logger.debug(f"list_folders request: {request}")
 
         client = await self._get_control_plane_client()
-        async for folder in await client.list_folders(request=request):
+        async for folder in await client.list_folders(
+            request=request,
+            retry=self._get_retry_config(),
+            timeout=STORAGE_CONTROL_RPC_TIMEOUT,
+        ):
             folders.append(self._create_folder_entry(bucket, folder))
 
         return folders

gcsfs/retry.py

@@ -7,8 +7,16 @@
 import google.auth.exceptions
 import requests.exceptions
 from decorator import decorator
+from google.api_core import exceptions as api_exceptions
+from google.api_core.retry import AsyncRetry
 
 logger = logging.getLogger("gcsfs")
+DEFAULT_RETRY_CONFIG = {
+    "timeout": 60.0,
+    "initial": 1.0,
+    "maximum": 60.0,
+    "multiplier": 2.0,
+}
 
 
 class HttpError(Exception):
@@ -176,3 +184,47 @@ async def retry_request(func, retries=6, *args, **kwargs):
                 continue
             logger.exception(f"{func.__name__} non-retriable exception: {e}")
             raise e
+
+
+def _is_transient_exception(exception):
+    is_transient = isinstance(
+        exception,
+        (
+            api_exceptions.DeadlineExceeded,
+            api_exceptions.ServiceUnavailable,
+            api_exceptions.InternalServerError,
+            api_exceptions.TooManyRequests,
+            api_exceptions.ResourceExhausted,
+            api_exceptions.Unknown,
+        ),
+    )
+    if (
+        not is_transient
+        and isinstance(exception, api_exceptions.Unauthenticated)
+        and "Invalid Credentials" in str(exception)
+    ):
+        is_transient = True
+    return is_transient
+
+
+def get_storage_control_retry_config(base_config=None, **kwargs) -> AsyncRetry:
+    """
+    Returns an AsyncRetry object configured for Storage Control API calls.
+
+    Priority: kwargs (timeout, etc.) > base_config > package defaults.
+
+    Args:
+        base_config: A dict containing base settings.
+        **kwargs: Direct call-site overrides (e.g., timeout=10).
+    """
+    retry_kwargs = DEFAULT_RETRY_CONFIG.copy()
+    valid_keys = DEFAULT_RETRY_CONFIG.keys()
+    if base_config:
+        retry_kwargs.update(
+            {k: v for k, v in base_config.items() if k in valid_keys and v is not None}
+        )
+
+    overrides = {k: v for k, v in kwargs.items() if k in valid_keys and v is not None}
+    retry_kwargs.update(overrides)
+
+    return AsyncRetry(predicate=_is_transient_exception, **retry_kwargs)