add a retry_on_error decoration for gardenlinux.oci

Author: yeoldegrove

src/gardenlinux/oci/helper.py

@@ -1,7 +1,14 @@
+import functools
 import json
+import logging
 import os
 import re
+import time
 
+from requests.exceptions import RequestException, HTTPError
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 
 def write_dict_to_json_file(input, output_path):
     if os.path.exists(output_path):
@@ -22,3 +29,74 @@ def get_uri_for_digest(uri, digest):
     """
     base_uri = re.split(r"[@:]", uri, maxsplit=1)[0]
     return f"{base_uri}@{digest}"
+
+
+def retry_on_error(
+    max_retries=3,
+    initial_delay=1,
+    backoff_factor=2,
+    retryable_status_codes=(502, 503, 504, 429),
+    retryable_exceptions=(RequestException,),
+):
+    """
+    A decorator for retrying functions that might fail due to transient network issues.
+
+    Args:
+        max_retries: Maximum number of retry attempts
+        initial_delay: Initial delay between retries in seconds
+        backoff_factor: Factor by which the delay increases with each retry
+        retryable_status_codes: HTTP status codes that trigger a retry
+        retryable_exceptions: Exception types that trigger a retry
+
+    Returns:
+        Decorated function with retry logic
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            delay = initial_delay
+            last_exception = None
+
+            for retry_count in range(max_retries + 1):
+                try:
+                    if retry_count > 0:
+                        logger.info(
+                            f"Retry attempt {retry_count}/{max_retries} for {func.__name__}"
+                        )
+
+                    response = func(*args, **kwargs)
+
+                    # Check for retryable status codes in the response
+                    if (
+                        hasattr(response, "status_code")
+                        and response.status_code in retryable_status_codes
+                    ):
+                        status_code = response.status_code
+                        logger.warning(
+                            f"Received status code {status_code} from {func.__name__}, retrying..."
+                        )
+                        last_exception = HTTPError(f"HTTP Error {status_code}")
+                    else:
+                        # Success, return the response
+                        return response
+
+                except retryable_exceptions as e:
+                    logger.warning(f"Request failed in {func.__name__}: {str(e)}")
+                    last_exception = e
+
+                # Don't sleep if this was the last attempt
+                if retry_count < max_retries:
+                    sleep_time = delay * (backoff_factor**retry_count)
+                    logger.info(f"Waiting {sleep_time:.2f} seconds before retry")
+                    time.sleep(sleep_time)
+
+            # If we got here, all retries failed
+            logger.error(f"All {max_retries} retries failed for {func.__name__}")
+            if last_exception:
+                raise last_exception
+            return None
+
+        return wrapper
+
+    return decorator

src/gardenlinux/oci/registry.py

@@ -38,6 +38,7 @@
     calculate_sha256,
     verify_sha256,
 )
+from .helper import retry_on_error
 from python_gardenlinux_lib.features.parse_features import get_oci_metadata_from_fileset
 from .schemas import (
     EmptyIndex,
@@ -659,6 +660,22 @@ def create_layer(
         }
         return layer
 
+    @retry_on_error(max_retries=3, initial_delay=2, backoff_factor=2)
+    def upload_blob(self, file_path, container, metadata=None):
+        """
+        Upload a blob to the registry with retry logic for network errors.
+
+        Args:
+            file_path: Path to the file to upload
+            container: Container object
+            metadata: Optional metadata for the blob
+
+        Returns:
+            Response from the upload
+        """
+        # Call the parent class's upload_blob method
+        return super().upload_blob(file_path, container, metadata)
+
     def push_from_dir(
         self,
         architecture: str,