Merge pull request #126 from dreadnode/users/raja/fix-refresh-token-race-condition

fix: Fix S3 credential auto-refresh mechanism and extend buffer window

Author: rdheekonda

docs/sdk/artifact.mdx

@@ -244,10 +244,7 @@ ArtifactStorage
 ---------------
 
 ```python
-ArtifactStorage(
-    file_system: AbstractFileSystem,
-    credential_refresher: Callable[[], bool] | None = None,
-)
+ArtifactStorage(credential_manager: CredentialManager)
 ```
 
 Storage for artifacts with efficient handling of large files and directories.
@@ -256,35 +253,24 @@ Supports:
 - Content-based deduplication using SHA1 hashing
 - Batch uploads for directories handled by fsspec
 
-Initialize artifact storage with a file system and prefix path.
+Initialize artifact storage with credential manager.
 
 **Parameters:**
 
-* **`file_system`**
-  (`AbstractFileSystem`)
-  –FSSpec-compatible file system
-* **`credential_refresher`**
-  (`Callable[[], bool] | None`, default:
-  `None`
-  )
-  –Optional function to refresh credentials when it's about to expire
+* **`credential_manager`**
+  (`CredentialManager`)
+  –Optional credential manager for S3 operations
 
 <Accordion title="Source code in dreadnode/artifact/storage.py" icon="code">
 ```python
-def __init__(
-    self,
-    file_system: fsspec.AbstractFileSystem,
-    credential_refresher: t.Callable[[], bool] | None = None,
-):
+def __init__(self, credential_manager: CredentialManager):
     """
-    Initialize artifact storage with a file system and prefix path.
+    Initialize artifact storage with credential manager.
 
     Args:
-        file_system: FSSpec-compatible file system
-        credential_refresher: Optional function to refresh credentials when it's about to expire
+        credential_manager: Optional credential manager for S3 operations
     """
-    self._file_system = file_system
-    self._credential_refresher = credential_refresher
+    self._credential_manager: CredentialManager = credential_manager
 ```
 
 
@@ -330,23 +316,26 @@ def batch_upload_files(self, source_paths: list[str], target_paths: list[str]) -
     if not source_paths:
         return []
 
-    logger.debug("Batch uploading %d files", len(source_paths))
+    def batch_upload_operation() -> list[str]:
+        filesystem = self._credential_manager.get_filesystem()
 
-    srcs = []
-    dsts = []
+        srcs = []
+        dsts = []
 
-    for src, dst in zip(source_paths, target_paths, strict=False):
-        if not self._file_system.exists(dst):
-            srcs.append(src)
-            dsts.append(dst)
+        for src, dst in zip(source_paths, target_paths, strict=False):
+            if not filesystem.exists(dst):
+                srcs.append(src)
+                dsts.append(dst)
 
-    if srcs:
-        self._file_system.put(srcs, dsts)
-        logger.debug("Batch upload completed for %d files", len(srcs))
-    else:
-        logger.debug("All files already exist, skipping upload")
+        if srcs:
+            filesystem.put(srcs, dsts)
+            logger.info("Batch upload completed for %d files", len(srcs))
+        else:
+            logger.info("All files already exist, skipping upload")
 
-    return [str(self._file_system.unstrip_protocol(target)) for target in target_paths]
+        return [str(filesystem.unstrip_protocol(target)) for target in target_paths]
+
+    return self._credential_manager.execute_with_retry(batch_upload_operation)
 ```
 
 
@@ -391,8 +380,9 @@ def compute_file_hash(self, file_path: Path, stream_threshold_mb: int = 10) -> s
     Returns:
         First 16 chars of SHA1 hash
     """
+
     file_size = file_path.stat().st_size
-    stream_threshold = stream_threshold_mb * 1024 * 1024  # Convert MB to bytes
+    stream_threshold = stream_threshold_mb * 1024 * 1024
 
     sha1 = hashlib.sha1()  # noqa: S324 # nosec
 
@@ -478,7 +468,6 @@ Store a file in the storage system, using multipart upload for large files.
 
 <Accordion title="Source code in dreadnode/artifact/storage.py" icon="code">
 ```python
-@with_credential_refresh
 def store_file(self, file_path: Path, target_key: str) -> str:
     """
     Store a file in the storage system, using multipart upload for large files.
@@ -490,13 +479,19 @@ def store_file(self, file_path: Path, target_key: str) -> str:
     Returns:
         Full URI with protocol to the stored file
     """
-    if not self._file_system.exists(target_key):
-        self._file_system.put(str(file_path), target_key)
-        logger.debug("Artifact successfully stored at %s", target_key)
-    else:
-        logger.debug("Artifact already exists at %s, skipping upload.", target_key)
 
-    return str(self._file_system.unstrip_protocol(target_key))
+    def store_operation() -> str:
+        filesystem = self._credential_manager.get_filesystem()
+
+        if not filesystem.exists(target_key):
+            filesystem.put(str(file_path), target_key)
+            logger.info("Artifact successfully stored at %s", target_key)
+        else:
+            logger.info("Artifact already exists at %s, skipping upload.", target_key)
+
+        return str(filesystem.unstrip_protocol(target_key))
+
+    return self._credential_manager.execute_with_retry(store_operation)
 ```
 
 

docs/sdk/main.mdx

@@ -57,16 +57,14 @@ def __init__(
     self.otel_scope = otel_scope
 
     self._api: ApiClient | None = None
-
+    self._credential_manager: CredentialManager | None = None
     self._logfire = logfire.DEFAULT_LOGFIRE_INSTANCE
     self._logfire.config.ignore_no_config = True
 
     self._fs: AbstractFileSystem = LocalFileSystem(auto_mkdir=True)
     self._fs_prefix: str = ".dreadnode/storage/"
 
     self._initialized = False
-    self._credentials: UserDataCredentials | None = None
-    self._credentials_expiry: datetime | None = None
 ```
 
 
@@ -380,9 +378,7 @@ def continue_run(self, run_context: RunContext) -> RunSpan:
     return RunSpan.from_context(
         context=run_context,
         tracer=self._get_tracer(),
-        file_system=self._fs,
-        prefix_path=self._fs_prefix,
-        credential_refresher=self._refresh_storage_credentials if self._credentials else None,
+        credential_manager=self._credential_manager,  # type: ignore[arg-type]
     )
 ```
 
@@ -526,19 +522,15 @@ def initialize(self) -> None:
         #         )
         #     )
         # )
-        self._credentials = self._api.get_user_data_credentials()
-        self._credentials_expiry = self._credentials.expiration
-        resolved_endpoint = resolve_endpoint(self._credentials.endpoint)
-        self._fs = S3FileSystem(
-            key=self._credentials.access_key_id,
-            secret=self._credentials.secret_access_key,
-            token=self._credentials.session_token,
-            client_kwargs={
-                "endpoint_url": resolved_endpoint,
-                "region_name": self._credentials.region,
-            },
-        )
-        self._fs_prefix = f"{self._credentials.bucket}/{self._credentials.prefix}/"
+        if self._api is not None:
+            api = self._api
+            self._credential_manager = CredentialManager(
+                credential_fetcher=lambda: api.get_user_data_credentials()
+            )
+            self._credential_manager.initialize()
+
+            self._fs = self._credential_manager.get_filesystem()
+            self._fs_prefix = self._credential_manager.get_prefix()
 
     self._logfire = logfire.configure(
         local=not self.is_default,
@@ -1723,10 +1715,8 @@ def run(
         tracer=self._get_tracer(),
         params=params,
         tags=tags,
-        file_system=self._fs,
-        prefix_path=self._fs_prefix,
+        credential_manager=self._credential_manager,  # type: ignore[arg-type]
         autolog=autolog,
-        credential_refresher=self._refresh_storage_credentials if self._credentials else None,
     )
 ```
 

dreadnode/artifact/storage.py

@@ -4,12 +4,9 @@
 """
 
 import hashlib
-import typing as t
 from pathlib import Path
 
-import fsspec  # type: ignore[import-untyped]
-
-from dreadnode.storage_utils import with_credential_refresh
+from dreadnode.credential_manager import CredentialManager
 from dreadnode.util import logger
 
 CHUNK_SIZE = 8 * 1024 * 1024  # 8MB
@@ -24,27 +21,15 @@ class ArtifactStorage:
     - Batch uploads for directories handled by fsspec
     """
 
-    def __init__(
-        self,
-        file_system: fsspec.AbstractFileSystem,
-        credential_refresher: t.Callable[[], bool] | None = None,
-    ):
+    def __init__(self, credential_manager: CredentialManager):
         """
-        Initialize artifact storage with a file system and prefix path.
+        Initialize artifact storage with credential manager.
 
         Args:
-            file_system: FSSpec-compatible file system
-            credential_refresher: Optional function to refresh credentials when it's about to expire
+            credential_manager: Optional credential manager for S3 operations
         """
-        self._file_system = file_system
-        self._credential_refresher = credential_refresher
-
-    def _refresh_credentials_if_needed(self) -> None:
-        """Refresh credentials if refresher is available."""
-        if self._credential_refresher:
-            self._credential_refresher()
+        self._credential_manager: CredentialManager = credential_manager
 
-    @with_credential_refresh
     def store_file(self, file_path: Path, target_key: str) -> str:
         """
         Store a file in the storage system, using multipart upload for large files.
@@ -56,13 +41,19 @@ def store_file(self, file_path: Path, target_key: str) -> str:
         Returns:
             Full URI with protocol to the stored file
         """
-        if not self._file_system.exists(target_key):
-            self._file_system.put(str(file_path), target_key)
-            logger.debug("Artifact successfully stored at %s", target_key)
-        else:
-            logger.debug("Artifact already exists at %s, skipping upload.", target_key)
 
-        return str(self._file_system.unstrip_protocol(target_key))
+        def store_operation() -> str:
+            filesystem = self._credential_manager.get_filesystem()
+
+            if not filesystem.exists(target_key):
+                filesystem.put(str(file_path), target_key)
+                logger.info("Artifact successfully stored at %s", target_key)
+            else:
+                logger.info("Artifact already exists at %s, skipping upload.", target_key)
+
+            return str(filesystem.unstrip_protocol(target_key))
+
+        return self._credential_manager.execute_with_retry(store_operation)
 
     def batch_upload_files(self, source_paths: list[str], target_paths: list[str]) -> list[str]:
         """
@@ -78,23 +69,26 @@ def batch_upload_files(self, source_paths: list[str], target_paths: list[str]) -
         if not source_paths:
             return []
 
-        logger.debug("Batch uploading %d files", len(source_paths))
+        def batch_upload_operation() -> list[str]:
+            filesystem = self._credential_manager.get_filesystem()
 
-        srcs = []
-        dsts = []
+            srcs = []
+            dsts = []
 
-        for src, dst in zip(source_paths, target_paths, strict=False):
-            if not self._file_system.exists(dst):
-                srcs.append(src)
-                dsts.append(dst)
+            for src, dst in zip(source_paths, target_paths, strict=False):
+                if not filesystem.exists(dst):
+                    srcs.append(src)
+                    dsts.append(dst)
 
-        if srcs:
-            self._file_system.put(srcs, dsts)
-            logger.debug("Batch upload completed for %d files", len(srcs))
-        else:
-            logger.debug("All files already exist, skipping upload")
+            if srcs:
+                filesystem.put(srcs, dsts)
+                logger.info("Batch upload completed for %d files", len(srcs))
+            else:
+                logger.info("All files already exist, skipping upload")
 
-        return [str(self._file_system.unstrip_protocol(target)) for target in target_paths]
+            return [str(filesystem.unstrip_protocol(target)) for target in target_paths]
+
+        return self._credential_manager.execute_with_retry(batch_upload_operation)
 
     def compute_file_hash(self, file_path: Path, stream_threshold_mb: int = 10) -> str:
         """
@@ -107,8 +101,9 @@ def compute_file_hash(self, file_path: Path, stream_threshold_mb: int = 10) -> s
         Returns:
             First 16 chars of SHA1 hash
         """
+
         file_size = file_path.stat().st_size
-        stream_threshold = stream_threshold_mb * 1024 * 1024  # Convert MB to bytes
+        stream_threshold = stream_threshold_mb * 1024 * 1024
 
         sha1 = hashlib.sha1()  # noqa: S324 # nosec
 

dreadnode/constants.py

@@ -58,4 +58,4 @@
 )
 
 # Default values for the file system credential management
-FS_CREDENTIAL_REFRESH_BUFFER = 300  # 5 minutes in seconds
+FS_CREDENTIAL_REFRESH_BUFFER = 900  # 15 minutes in seconds