storage protection during parallel creation

Author: Mantisus

src/crawlee/storages/_storage_instance_manager.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+from asyncio import Lock
 from collections import defaultdict
 from collections.abc import Coroutine, Hashable
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, TypeVar
+from weakref import WeakValueDictionary
 
 from crawlee._utils.raise_if_too_many_kwargs import raise_if_too_many_kwargs
 from crawlee.storage_clients._base import DatasetClient, KeyValueStoreClient, RequestQueueClient
@@ -76,6 +78,7 @@ class StorageInstanceManager:
 
     def __init__(self) -> None:
         self._cache: _StorageCache = _StorageCache()
+        self._opener_locks: WeakValueDictionary[tuple, Lock] = WeakValueDictionary()
 
     async def open_storage_instance(
         self,
@@ -120,62 +123,74 @@ async def open_storage_instance(
                 alias = self._DEFAULT_STORAGE_ALIAS
 
             # Check cache
-            if id is not None and (cached_instance := self._cache.by_id[cls][id].get(storage_client_cache_key)):
-                if isinstance(cached_instance, cls):
-                    return cached_instance
-                raise RuntimeError('Cached instance type mismatch.')
-
-            if name is not None and (cached_instance := self._cache.by_name[cls][name].get(storage_client_cache_key)):
-                if isinstance(cached_instance, cls):
-                    return cached_instance
-                raise RuntimeError('Cached instance type mismatch.')
-
-            if alias is not None and (
-                cached_instance := self._cache.by_alias[cls][alias].get(storage_client_cache_key)
+            if cached_instance := self._get_from_cache(
+                cls,
+                id=id,
+                name=name,
+                alias=alias,
+                storage_client_cache_key=storage_client_cache_key,
             ):
-                if isinstance(cached_instance, cls):
-                    return cached_instance
-                raise RuntimeError('Cached instance type mismatch.')
+                return cached_instance
 
             # Check for conflicts between named and alias storages
-            if alias and (self._cache.by_name[cls][alias].get(storage_client_cache_key)):
-                raise ValueError(
-                    f'Cannot create alias storage "{alias}" because a named storage with the same name already exists. '
-                    f'Use a different alias or drop the existing named storage first.'
-                )
-
-            if name and (self._cache.by_alias[cls][name].get(storage_client_cache_key)):
-                raise ValueError(
-                    f'Cannot create named storage "{name}" because an alias storage with the same name already exists. '
-                    f'Use a different name or drop the existing alias storage first.'
-                )
+            self._check_name_alias_conflict(
+                cls,
+                name=name,
+                alias=alias,
+                storage_client_cache_key=storage_client_cache_key,
+            )
 
             # Validate storage name
             if name is not None:
                 validate_storage_name(name)
 
-            # Create new instance
-            client: KeyValueStoreClient | DatasetClient | RequestQueueClient
-            client = await client_opener_coro
+            # Acquire lock for this opener
+            opener_lock_key = (cls, str(id or name or alias), storage_client_cache_key)
+            if not (lock := self._opener_locks.get(opener_lock_key)):
+                lock = Lock()
+                self._opener_locks[opener_lock_key] = lock
+
+            async with lock:
+                # Re-check cache inside the lock
+                if cached_instance := self._get_from_cache(
+                    cls,
+                    id=id,
+                    name=name,
+                    alias=alias,
+                    storage_client_cache_key=storage_client_cache_key,
+                ):
+                    return cached_instance
 
-            metadata = await client.get_metadata()
+                # 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,
+                )
 
-            instance = cls(client, metadata.id, metadata.name)  # type: ignore[call-arg]
-            instance_name = getattr(instance, 'name', None)
+                # Create new instance
+                client: KeyValueStoreClient | DatasetClient | RequestQueueClient
+                client = await client_opener_coro
 
-            # Cache the instance.
-            # Always cache by id.
-            self._cache.by_id[cls][instance.id][storage_client_cache_key] = instance
+                metadata = await client.get_metadata()
 
-            # Cache named storage.
-            if instance_name is not None:
-                self._cache.by_name[cls][instance_name][storage_client_cache_key] = instance
+                instance = cls(client, metadata.id, metadata.name)  # type: ignore[call-arg]
+                instance_name = getattr(instance, 'name', None)
 
-            # Cache unnamed storage.
-            if alias is not None:
-                self._cache.by_alias[cls][alias][storage_client_cache_key] = instance
+                # Cache the instance.
+                # Always cache by id.
+                self._cache.by_id[cls][instance.id][storage_client_cache_key] = instance
 
-            return 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.
@@ -193,3 +208,51 @@ def remove_from_cache(self, storage_instance: Storage) -> None:
     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,
+        storage_client_cache_key: Hashable = '',
+    ) -> T | None:
+        """Get a storage instance from the cache."""
+        if id is not None and (cached_instance := self._cache.by_id[cls][id].get(storage_client_cache_key)):
+            if isinstance(cached_instance, cls):
+                return cached_instance
+            raise RuntimeError('Cached instance type mismatch.')
+
+        if name is not None and (cached_instance := self._cache.by_name[cls][name].get(storage_client_cache_key)):
+            if isinstance(cached_instance, cls):
+                return cached_instance
+            raise RuntimeError('Cached instance type mismatch.')
+
+        if alias is not None and (cached_instance := self._cache.by_alias[cls][alias].get(storage_client_cache_key)):
+            if isinstance(cached_instance, cls):
+                return cached_instance
+            raise RuntimeError('Cached instance type mismatch.')
+
+        return None
+
+    def _check_name_alias_conflict(
+        self,
+        cls: type[T],
+        *,
+        name: str | None = None,
+        alias: str | None = None,
+        storage_client_cache_key: Hashable = '',
+    ) -> None:
+        """Check for conflicts between named and alias storages."""
+        if alias and (self._cache.by_name[cls][alias].get(storage_client_cache_key)):
+            raise ValueError(
+                f'Cannot create alias storage "{alias}" because a named storage with the same name already exists. '
+                f'Use a different alias or drop the existing named storage first.'
+            )
+
+        if name and (self._cache.by_alias[cls][name].get(storage_client_cache_key)):
+            raise ValueError(
+                f'Cannot create named storage "{name}" because an alias storage with the same name already exists. '
+                f'Use a different name or drop the existing alias storage first.'
+            )

tests/unit/storages/test_storage_instance_manager.py

@@ -1,5 +1,8 @@
+import asyncio
+import sys
 from pathlib import Path
 from typing import cast
+from unittest.mock import AsyncMock
 
 import pytest
 
@@ -128,3 +131,61 @@ async def test_preexisting_unnamed_storage_open_by_id(storage_type: type[Storage
     storage_1_again = await storage_type.open(id=storage_1.id, storage_client=storage_client)
 
     assert storage_1.id == storage_1_again.id
+
+
+@pytest.mark.skipif(sys.version_info[:3] < (3, 11), reason='asyncio.Barrier was introduced in Python 3.11.')
+async def test_concurrent_open_datasets() -> None:
+    """Test that concurrent open datasets with the same name return the same instance."""
+    from asyncio import Barrier  # type:ignore[attr-defined] # noqa: PLC0415
+
+    barrier = Barrier(2)
+
+    async def push_data(data: dict) -> None:
+        await barrier.wait()
+        dataset = await Dataset.open(name='concurrent-storage')
+        await dataset.push_data(data)
+
+    await asyncio.gather(
+        push_data({'test_1': '1'}),
+        push_data({'test_2': '2'}),
+    )
+
+    dataset = await Dataset.open(name='concurrent-storage')
+
+    items = await dataset.get_data()
+    assert len(items.items) == 2
+
+    await dataset.drop()
+
+
+@pytest.mark.skipif(sys.version_info[:3] < (3, 11), reason='asyncio.Barrier was introduced in Python 3.11.')
+async def test_concurrent_open_datasets_with_same_name_and_alias() -> None:
+    """Test that concurrent open requests for the same storage return the same instance."""
+    from asyncio import Barrier  # type:ignore[attr-defined] # noqa: PLC0415
+
+    valid_kwargs: dict[str, str | None] = {}
+
+    exception_calls = AsyncMock()
+
+    barrier = Barrier(2)
+
+    async def open_dataset(name: str | None, alias: str | None) -> None:
+        await barrier.wait()
+        try:
+            await Dataset.open(name=name, alias=alias)
+            valid_kwargs['name'] = name
+            valid_kwargs['alias'] = alias
+        except ValueError:
+            exception_calls()
+
+    await asyncio.gather(
+        open_dataset(name=None, alias='concurrent-storage'),
+        open_dataset(name='concurrent-storage', alias=None),
+    )
+
+    # Ensure that a ValueError was raised due to name/alias conflict
+    exception_calls.assert_called_once()
+
+    dataset = await Dataset.open(name=valid_kwargs.get('name'), alias=valid_kwargs.get('alias'))
+
+    await dataset.drop()