simplify again

remove unused code

typing

one more

Author: zzstoatzz

src/integrations/prefect-redis/prefect_redis/locking.py

@@ -19,7 +19,8 @@ class RedisLockManager(LockManager):
         username: The username to use when connecting to the Redis server
         password: The password to use when connecting to the Redis server
         ssl: Whether to use SSL when connecting to the Redis server
-        # client and async_client are initialized lazily
+        client: The Redis client used to communicate with the Redis server
+        async_client: The asynchronous Redis client used to communicate with the Redis server
 
     Example:
         Use with a cache policy:
@@ -64,11 +65,13 @@ def __init__(
         self.username = username
         self.password = password
         self.ssl = ssl
-        self.client: Optional[Redis] = None
-        self.async_client: Optional[AsyncRedis] = None
+        # Clients are initialized by _init_clients
+        self.client: Redis
+        self.async_client: AsyncRedis
+        self._init_clients()  # Initialize clients here
         self._locks: dict[str, Lock | AsyncLock] = {}
 
-    # ---------- pickle helpers ----------
+    # ---------- pickling ----------
     def __getstate__(self) -> dict[str, Any]:
         return {
             k: getattr(self, k)
@@ -77,31 +80,28 @@ def __getstate__(self) -> dict[str, Any]:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__.update(state)
-        self.client = None
-        self.async_client = None
+        self._init_clients()  # Re-initialize clients here
         self._locks = {}
 
     # ------------------------------------
 
-    def _ensure_clients(self) -> None:
-        if self.client is None:
-            self.client = Redis(
-                host=self.host,
-                port=self.port,
-                db=self.db,
-                username=self.username,
-                password=self.password,
-                ssl=self.ssl,
-            )
-        if self.async_client is None:
-            self.async_client = AsyncRedis(
-                host=self.host,
-                port=self.port,
-                db=self.db,
-                username=self.username,
-                password=self.password,
-                ssl=self.ssl,
-            )
+    def _init_clients(self) -> None:
+        self.client = Redis(
+            host=self.host,
+            port=self.port,
+            db=self.db,
+            username=self.username,
+            password=self.password,
+            ssl=self.ssl,
+        )
+        self.async_client = AsyncRedis(
+            host=self.host,
+            port=self.port,
+            db=self.db,
+            username=self.username,
+            password=self.password,
+            ssl=self.ssl,
+        )
 
     @staticmethod
     def _lock_name_for_key(key: str) -> str:
@@ -114,16 +114,24 @@ def acquire_lock(
         acquire_timeout: Optional[float] = None,
         hold_timeout: Optional[float] = None,
     ) -> bool:
-        self._ensure_clients()
+        """
+        Acquires a lock synchronously.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock.
+            acquire_timeout: Maximum time to wait for the lock to be acquired.
+            hold_timeout: Maximum time to hold the lock.
+
+        Returns:
+            True if the lock was acquired, False otherwise.
+        """
         lock_name = self._lock_name_for_key(key)
         lock = self._locks.get(lock_name)
 
-        if lock is not None and self.is_lock_holder(
-            key, holder
-        ):  # is_lock_holder will also call _ensure_clients
+        if lock is not None and self.is_lock_holder(key, holder):
             return True
         else:
-            # If lock is None, or not held by current holder, create/acquire new one.
             lock = Lock(
                 self.client, lock_name, timeout=hold_timeout, thread_local=False
             )
@@ -139,7 +147,19 @@ async def aacquire_lock(
         acquire_timeout: Optional[float] = None,
         hold_timeout: Optional[float] = None,
     ) -> bool:
-        self._ensure_clients()
+        """
+        Acquires a lock asynchronously.
+
+        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.
+            acquire_timeout: Maximum time to wait for the lock to be acquired.
+            hold_timeout: Maximum time to hold the lock.
+
+        Returns:
+            True if the lock was acquired, False otherwise.
+        """
         lock_name = self._lock_name_for_key(key)
         lock = self._locks.get(lock_name)
 
@@ -149,8 +169,10 @@ async def aacquire_lock(
             else:
                 lock = None
 
+        # Handles the case where a lock might have been released during a task retry
+        # If the lock doesn't exist in Redis at all, this method will succeed even if
+        # the holder ID doesn't match the original holder.
         if lock is None:
-            assert self.async_client is not None, "Async client should be initialized"
             new_lock = AsyncLock(
                 self.async_client, lock_name, timeout=hold_timeout, thread_local=False
             )
@@ -164,7 +186,21 @@ async def aacquire_lock(
         return False
 
     def release_lock(self, key: str, holder: str) -> None:
-        self._ensure_clients()
+        """
+        Releases the lock on the corresponding transaction record.
+
+        Handles the case where a lock might have been released during a task retry
+        If the lock doesn't exist in Redis at all, this method will succeed even if
+        the holder ID doesn't match the original holder.
+
+        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.
+
+        Raises:
+            ValueError: If the lock is held by a different holder.
+        """
         lock_name = self._lock_name_for_key(key)
         lock = self._locks.get(lock_name)
 
@@ -173,67 +209,37 @@ def release_lock(self, key: str, holder: str) -> None:
             del self._locks[lock_name]
             return
 
-        if not self.is_locked(key):  # is_locked calls _ensure_clients
+        # If the lock doesn't exist in Redis at all, it's already been released
+        if not self.is_locked(key):
             if lock_name in self._locks:
                 del self._locks[lock_name]
             return
 
+        # We have a real conflict - lock exists in Redis but with a different holder
         raise ValueError(f"No lock held by {holder} for transaction with key {key}")
 
-    async def arelease_lock(self, key: str, holder: str) -> None:  # Added async version
-        self._ensure_clients()
-        lock_name = self._lock_name_for_key(key)
-        lock = self._locks.get(lock_name)
-
-        if lock is not None and isinstance(
-            lock, AsyncLock
-        ):  # Still need to check if it *is* an AsyncLock to call await .owned()
-            if await lock.owned() and lock.local.token == holder.encode():
-                await lock.release()
-                del self._locks[lock_name]
-                return
-
-        # Check if the lock key exists on the server at all.
-        if not AsyncLock(self.async_client, lock_name).locked():
-            # If the lock doesn't exist on the server, it's already effectively released.
-            # Clean up from self._locks if it was there but holder didn't match.
-            if lock_name in self._locks:
-                del self._locks[lock_name]
-            return
-
-        raise ValueError(
-            f"No lock held by {holder} for transaction with key {key} (async)"
-        )
-
     def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
-        self._ensure_clients()
         lock_name = self._lock_name_for_key(key)
-        lock = Lock(self.client, lock_name)  # Create a temporary lock for waiting
+        lock = Lock(self.client, lock_name)
         lock_freed = lock.acquire(blocking_timeout=timeout)
         if lock_freed:
             lock.release()
         return lock_freed
 
     async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
-        self._ensure_clients()
         lock_name = self._lock_name_for_key(key)
-        assert self.async_client is not None, "Async client should be initialized"
-        lock = AsyncLock(
-            self.async_client, lock_name
-        )  # Create a temporary lock for waiting
+        lock = AsyncLock(self.async_client, lock_name)
         lock_freed = await lock.acquire(blocking_timeout=timeout)
         if lock_freed:
             await lock.release()
         return lock_freed
 
     def is_locked(self, key: str) -> bool:
-        self._ensure_clients()
         lock_name = self._lock_name_for_key(key)
-        lock = Lock(self.client, lock_name)  # Create a temporary lock for checking
+        lock = Lock(self.client, lock_name)
         return lock.locked()
 
     def is_lock_holder(self, key: str, holder: str) -> bool:
-        self._ensure_clients()  # Ensures self.client is available if needed by _locks access logic
         lock_name = self._lock_name_for_key(key)
         lock = self._locks.get(lock_name)
         if lock is None:

src/integrations/prefect-redis/tests/test_locking.py

@@ -1,3 +1,4 @@
+import asyncio
 import pickle
 import queue
 import threading
@@ -261,33 +262,42 @@ async def test_pickle_unpickle_and_use_lock_manager(
     ):
         """
         Tests that RedisLockManager can be pickled, unpickled, and then used successfully,
-        ensuring lazy client initialization works.
+        ensuring clients are re-initialized correctly on unpickle.
         """
-        # Initial state checks for the provided fixture instance
-        # (clients should be None due to our __init__ and __setstate__ if it were unpickled,
-        # or because _ensure_clients hasn't been called on a fresh instance)
-        assert lock_manager.client is None, "Initial client should be None"
-        assert lock_manager.async_client is None, "Initial async_client should be None"
+        # With the new __init__, clients are initialized immediately.
+        # So, no initial check for them being None.
+
+        # Store original client IDs for comparison after unpickling
+        original_client_id = id(lock_manager.client)
+        original_async_client_id = id(lock_manager.async_client)
 
         # Pickle and unpickle
         pickled_manager = pickle.dumps(lock_manager)
         unpickled_manager: RedisLockManager = pickle.loads(pickled_manager)
 
-        # Verify state after unpickling (clients should be None due to __setstate__)
-        assert unpickled_manager.client is None, (
-            "Client should be None after unpickling"
+        # Verify state after unpickling: clients should be NEW instances due to __setstate__ calling _init_clients()
+        assert unpickled_manager.client is not None, (
+            "Client should be re-initialized after unpickling, not None"
         )
-        assert unpickled_manager.async_client is None, (
-            "Async client should be None after unpickling"
+        assert unpickled_manager.async_client is not None, (
+            "Async client should be re-initialized after unpickling, not None"
         )
-        # Accessing _locks directly is for testing internals.
+
+        assert id(unpickled_manager.client) != original_client_id, (
+            "Client should be a NEW instance after unpickling"
+        )
+        assert id(unpickled_manager.async_client) != original_async_client_id, (
+            "Async client should be a NEW instance after unpickling"
+        )
+
+        # _locks should be an empty dict after unpickling due to __setstate__
         assert (
             hasattr(unpickled_manager, "_locks")
             and isinstance(getattr(unpickled_manager, "_locks"), dict)
             and not getattr(unpickled_manager, "_locks")
         ), "_locks should be an empty dict after unpickling"
 
-        # Test synchronous operations (should trigger _ensure_clients)
+        # Test synchronous operations (clients are already initialized)
         sync_key = "test_sync_pickle_key"
         sync_holder = "sync_pickle_holder"
 
@@ -307,29 +317,40 @@ async def test_pickle_unpickle_and_use_lock_manager(
         # Test asynchronous operations (should trigger _ensure_clients for async_client)
         async_key = "test_async_pickle_key"
         async_holder = "async_pickle_holder"
+        hold_timeout_seconds = 0.2  # Use a short hold timeout for testing expiration
 
         acquired_async = await unpickled_manager.aacquire_lock(
-            async_key, holder=async_holder, acquire_timeout=1, hold_timeout=5
+            async_key,
+            holder=async_holder,
+            acquire_timeout=1,
+            hold_timeout=hold_timeout_seconds,
         )
         assert acquired_async, "Should acquire async lock after unpickling"
         assert unpickled_manager.async_client is not None, (
             "Async client should be initialized after async use"
         )
 
-        # Verify holder by re-acquiring (should succeed) and then releasing.
+        # Verify holder by re-acquiring (should succeed as it's the same holder and lock is fresh)
         assert await unpickled_manager.aacquire_lock(
-            async_key, holder=async_holder, acquire_timeout=1, hold_timeout=5
+            async_key,
+            holder=async_holder,
+            acquire_timeout=1,
+            hold_timeout=hold_timeout_seconds,
         ), "Re-acquiring same async lock should succeed"
 
-        await unpickled_manager.arelease_lock(async_key, async_holder)
+        # Wait for the lock to expire based on hold_timeout
+        await asyncio.sleep(
+            hold_timeout_seconds + 0.1
+        )  # Wait a bit longer than the timeout
 
-        # Verify it's released by trying to acquire with a different holder.
+        # Verify it's released (expired) by trying to acquire with a different holder.
         new_async_holder = "new_async_pickle_holder"
         acquired_by_new = await unpickled_manager.aacquire_lock(
-            async_key, holder=new_async_holder, acquire_timeout=1, hold_timeout=5
+            async_key,
+            holder=new_async_holder,
+            acquire_timeout=1,
+            hold_timeout=hold_timeout_seconds,
         )
         assert acquired_by_new, (
-            "Should acquire async lock with new holder after release"
+            "Should acquire async lock with new holder after original lock expires"
         )
-        if acquired_by_new:  # Cleanup if acquired
-            await unpickled_manager.arelease_lock(async_key, new_async_holder)