Repository: apify/crawlee-python
File: src/crawlee/storages/_storage_instance_manager.py (Python)

Code before:
                # Re-check for conflicts between named and alias storages
                self._check_name_alias_conflict(
                    cls,
                    name=name,
                    alias=alias,
                    storage_client_cache_key=storage_client_cache_key,
                )

                # Create new instance
                client: KeyValueStoreClient | DatasetClient | RequestQueueClient
                client = await client_opener_coro

                metadata = await client.get_metadata()

                instance = cls(client, metadata.id, metadata.name)  # type: ignore[call-arg]
                instance_name = getattr(instance, 'name', None)

                # Cache the instance.
                # Always cache by id.
                self._cache.by_id[cls][instance.id][storage_client_cache_key] = instance

                # Cache named storage.
                if instance_name is not None:
                    self._cache.by_name[cls][instance_name][storage_client_cache_key] = instance

                # Cache unnamed storage.
                if alias is not None:
                    self._cache.by_alias[cls][alias][storage_client_cache_key] = instance

                return instance

        finally:
            # Make sure the client opener is closed.
            # If it was awaited, then closing is no operation, if it was not awaited, this is the cleanup.
            client_opener_coro.close()

    def remove_from_cache(self, storage_instance: Storage) -> None:
        """Remove a storage instance from the cache.

        Args:
            storage_instance: The storage instance to remove.
        """
        self._cache.remove_from_cache(storage_instance)

    def clear_cache(self) -> None:
        """Clear all cached storage instances."""
        self._cache = _StorageCache()

    def _get_from_cache(
        self,
        cls: type[T],

Review comment (bug):
This has to happen without interruption. Otherwise, the "cache pre-check" without a lock could fail in weird ways.

Can you add comments for our future selves that there should be no `await`s in this range?

Code after:

                # Create new instance
                client: KeyValueStoreClient | DatasetClient | RequestQueueClient
                client = await client_opener_coro

                metadata = await client.get_metadata()

                instance = cls(client, metadata.id, metadata.name)  # type: ignore[call-arg]
                instance_name = getattr(instance, 'name', None)

                # Cache the instance.
                # Note: No awaits in this section. All cache entries must be written
                # atomically to ensure pre-checks outside the lock see consistent state.

                # Always cache by id.
                self._cache.by_id[cls][instance.id][storage_client_cache_key] = instance

                # Cache named storage.
                if instance_name is not None:
                    self._cache.by_name[cls][instance_name][storage_client_cache_key] = instance

                # Cache unnamed storage.
                if alias is not None:
                    self._cache.by_alias[cls][alias][storage_client_cache_key] = instance

                return instance

        finally:
            # Make sure the client opener is closed.
            # If it was awaited, then closing is no operation, if it was not awaited, this is the cleanup.
            client_opener_coro.close()

    def remove_from_cache(self, storage_instance: Storage) -> None:
        """Remove a storage instance from the cache.

        Args:
            storage_instance: The storage instance to remove.
        """
        self._cache.remove_from_cache(storage_instance)

    def clear_cache(self) -> None:
        """Clear all cached storage instances."""
        self._cache = _StorageCache()

    def _get_from_cache(
        self,
        cls: type[T],
        *,
        id: str | None = None,
        name: str | None = None,
        alias: str | None = None,