prefect-client 3.0.0rc20__py3-none-any.whl → 3.0.1__py3-none-any.whl

prefect/locking/memory.py

@@ -0,0 +1,213 @@
+import asyncio
+import threading
+from typing import Dict, Optional, TypedDict
+
+from .protocol import LockManager
+
+
+class _LockInfo(TypedDict):
+    """
+    A dictionary containing information about a lock.
+
+    Attributes:
+        holder: The holder of the lock.
+        lock: The lock object.
+        expiration_timer: The timer for the lock expiration
+    """
+
+    holder: str
+    lock: threading.Lock
+    expiration_timer: Optional[threading.Timer]
+
+
+class MemoryLockManager(LockManager):
+    """
+    A lock manager that stores lock information in memory.
+
+    Note: because this lock manager stores data in memory, it is not suitable for
+    use in a distributed environment or across different processes.
+    """
+
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        self._locks_dict_lock = threading.Lock()
+        self._locks: Dict[str, _LockInfo] = {}
+
+    def _expire_lock(self, key: str):
+        """
+        Expire the lock for the given key.
+
+        Used as a callback for the expiration timer of a lock.
+
+        Args:
+            key: The key of the lock to expire.
+        """
+        with self._locks_dict_lock:
+            if key in self._locks:
+                lock_info = self._locks[key]
+                if lock_info["lock"].locked():
+                    lock_info["lock"].release()
+                if lock_info["expiration_timer"]:
+                    lock_info["expiration_timer"].cancel()
+                del self._locks[key]
+
+    def acquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        with self._locks_dict_lock:
+            if key not in self._locks:
+                lock = threading.Lock()
+                lock.acquire()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder, lock=lock, expiration_timer=expiration_timer
+                )
+                return True
+            elif self._locks[key]["holder"] == holder:
+                return True
+            else:
+                existing_lock_info = self._locks[key]
+
+        if acquire_timeout is not None:
+            existing_lock_acquired = existing_lock_info["lock"].acquire(
+                timeout=acquire_timeout
+            )
+        else:
+            existing_lock_acquired = existing_lock_info["lock"].acquire()
+
+        if existing_lock_acquired:
+            with self._locks_dict_lock:
+                if (
+                    expiration_timer := existing_lock_info["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder,
+                    lock=existing_lock_info["lock"],
+                    expiration_timer=expiration_timer,
+                )
+            return True
+        return False
+
+    async def aacquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        with self._locks_dict_lock:
+            if key not in self._locks:
+                lock = threading.Lock()
+                await asyncio.to_thread(lock.acquire)
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder, lock=lock, expiration_timer=expiration_timer
+                )
+                return True
+            elif self._locks[key]["holder"] == holder:
+                return True
+            else:
+                existing_lock_info = self._locks[key]
+
+        if acquire_timeout is not None:
+            existing_lock_acquired = await asyncio.to_thread(
+                existing_lock_info["lock"].acquire, timeout=acquire_timeout
+            )
+        else:
+            existing_lock_acquired = await asyncio.to_thread(
+                existing_lock_info["lock"].acquire
+            )
+
+        if existing_lock_acquired:
+            with self._locks_dict_lock:
+                if (
+                    expiration_timer := existing_lock_info["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder,
+                    lock=existing_lock_info["lock"],
+                    expiration_timer=expiration_timer,
+                )
+            return True
+        return False
+
+    def release_lock(self, key: str, holder: str) -> None:
+        with self._locks_dict_lock:
+            if key in self._locks and self._locks[key]["holder"] == holder:
+                if (
+                    expiration_timer := self._locks[key]["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                self._locks[key]["lock"].release()
+                del self._locks[key]
+            else:
+                raise ValueError(
+                    f"No lock held by {holder} for transaction with key {key}"
+                )
+
+    def is_locked(self, key: str) -> bool:
+        return key in self._locks and self._locks[key]["lock"].locked()
+
+    def is_lock_holder(self, key: str, holder: str) -> bool:
+        lock_info = self._locks.get(key)
+        return (
+            lock_info is not None
+            and lock_info["lock"].locked()
+            and lock_info["holder"] == holder
+        )
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        if lock := self._locks.get(key, {}).get("lock"):
+            if timeout is not None:
+                lock_acquired = lock.acquire(timeout=timeout)
+            else:
+                lock_acquired = lock.acquire()
+            if lock_acquired:
+                lock.release()
+            return lock_acquired
+        return True
+
+    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        if lock := self._locks.get(key, {}).get("lock"):
+            if timeout is not None:
+                lock_acquired = await asyncio.to_thread(lock.acquire, timeout=timeout)
+            else:
+                lock_acquired = await asyncio.to_thread(lock.acquire)
+            if lock_acquired:
+                lock.release()
+            return lock_acquired
+        return True

prefect/locking/protocol.py

@@ -0,0 +1,122 @@
+from typing import Optional, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class LockManager(Protocol):
+    def acquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Acquire a lock for a transaction record with the given key. Will block other
+        actors from updating this transaction record until the lock is
+        released.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+            acquire_timeout: Max number of seconds to wait for the record to become
+                available if it is locked while attempting to acquire a lock. Pass 0
+                to attempt to acquire a lock without waiting. Blocks indefinitely by
+                default.
+            hold_timeout: Max number of seconds to hold the lock for. Holds the lock
+                indefinitely by default.
+
+        Returns:
+            bool: True if the lock was successfully acquired; False otherwise.
+        """
+        ...
+
+    async def aacquire_lock(
+        self,
+        key: str,
+        holder: str,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Acquire a lock for a transaction record with the given key. Will block other
+        actors from updating this transaction record until the lock is
+        released.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+            acquire_timeout: Max number of seconds to wait for the record to become
+                available if it is locked while attempting to acquire a lock. Pass 0
+                to attempt to acquire a lock without waiting. Blocks indefinitely by
+                default.
+            hold_timeout: Max number of seconds to hold the lock for. Holds the lock
+                indefinitely by default.
+
+        Returns:
+            bool: True if the lock was successfully acquired; False otherwise.
+        """
+        ...
+
+    def release_lock(self, key: str, holder: str):
+        """
+        Releases the lock on the corresponding transaction record.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. Must match the
+                holder provided when acquiring the lock.
+        """
+        ...
+
+    def is_locked(self, key: str) -> bool:
+        """
+        Simple check to see if the corresponding record is currently locked.
+
+        Args:
+            key: Unique identifier for the transaction record.
+
+        Returns:
+            True is the record is locked; False otherwise.
+        """
+        ...
+
+    def is_lock_holder(self, key: str, holder: str) -> bool:
+        """
+        Check if the current holder is the lock holder for the transaction record.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+
+        Returns:
+            bool: True if the current holder is the lock holder; False otherwise.
+        """
+        ...
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for the corresponding transaction record to become free.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            timeout: Maximum time to wait. None means to wait indefinitely.
+
+        Returns:
+            bool: True if the lock becomes free within the timeout; False
+                otherwise.
+        """
+        ...
+
+    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for the corresponding transaction record to become free.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            timeout: Maximum time to wait. None means to wait indefinitely.
+
+        Returns:
+            bool: True if the lock becomes free within the timeout; False
+                otherwise.
+        """
+        ...

prefect/records/filesystem.py

@@ -76,7 +76,7 @@ class FileSystemRecordStore(RecordStore):
             return None
 
     def read(
-        self, key: str, holder: Optional[str] = None
+        self, key: str, holder: Optional[str] = None, timeout: Optional[float] = None
     ) -> Optional[TransactionRecord]:
         if not self.exists(key):
             return None
@@ -84,7 +84,9 @@ class FileSystemRecordStore(RecordStore):
         holder = holder or self.generate_default_holder()
 
         if self.is_locked(key) and not self.is_lock_holder(key, holder):
-            self.wait_for_lock(key)
+            unlocked = self.wait_for_lock(key, timeout=timeout)
+            if not unlocked:
+                return None
         record_data = self.records_directory.joinpath(key).read_text()
         return TransactionRecord(
             key=key, result=BaseResult.model_validate_json(record_data)

prefect/records/result_store.py

@@ -3,7 +3,7 @@ from typing import Any, Optional
 
 import pendulum
 
-from prefect.results import BaseResult, PersistedResult, ResultFactory
+from prefect.results import BaseResult, PersistedResult, ResultStore
 from prefect.transactions import IsolationLevel
 from prefect.utilities.asyncutils import run_coro_as_sync
 
@@ -11,8 +11,14 @@ from .base import RecordStore, TransactionRecord
 
 
 @dataclass
-class ResultFactoryStore(RecordStore):
-    result_factory: ResultFactory
+class ResultRecordStore(RecordStore):
+    """
+    A record store for result records.
+
+    Collocates result metadata with result data.
+    """
+
+    result_store: ResultStore
     cache: Optional[PersistedResult] = None
 
     def exists(self, key: str) -> bool:
@@ -38,8 +44,8 @@ class ResultFactoryStore(RecordStore):
             return TransactionRecord(key=key, result=self.cache)
         try:
             result = PersistedResult(
-                serializer_type=self.result_factory.serializer.type,
-                storage_block_id=self.result_factory.storage_block_id,
+                serializer_type=self.result_store.serializer.type,
+                storage_block_id=self.result_store.result_storage_block_id,
                 storage_key=key,
             )
             return TransactionRecord(key=key, result=result)
@@ -52,7 +58,7 @@ class ResultFactoryStore(RecordStore):
             # if the value is already a persisted result, write it
             result.write(_sync=True)
         elif not isinstance(result, BaseResult):
-            run_coro_as_sync(self.result_factory.create_result(obj=result, key=key))
+            run_coro_as_sync(self.result_store.create_result(obj=result, key=key))
 
     def supports_isolation_level(self, isolation_level: IsolationLevel) -> bool:
         return isolation_level == IsolationLevel.READ_COMMITTED