python-etcd3-async 0.1.0__py3-none-any.whl

etcd3/locks.py

@@ -0,0 +1,593 @@
+import asyncio
+import threading
+import time
+import uuid
+from contextlib import asynccontextmanager
+
+from etcd3 import events, exceptions
+from etcd3.trace import get_current_span, is_tracing_available
+from etcd3.trace_attributes import (
+    build_lock_attributes,
+)
+
+
+class Lock(object):
+    """
+    A distributed lock.
+
+    This can be used as a context manager, with the lock being acquired and
+    released as you would expect:
+
+    .. code-block:: python
+
+        etcd = etcd3.client()
+
+        # create a lock that expires after 20 seconds
+        with etcd.lock('/myapp/mylock', ttl=20) as lock:
+            # do something that requires the lock
+            print(lock.is_acquired())
+
+            # refresh the timeout on the lease
+            lock.refresh()
+
+    :param name: name (full path) of the lock. The full key path should be provided,
+                 e.g., '/myapp/mylock' or 'election/my-service'.
+    :type name: string or bytes
+    :param ttl: length of time for the lock to live for in seconds. The lock
+                will be released after this time elapses, unless refreshed
+    :type ttl: int
+    """
+
+    def __init__(self, name, ttl=60, etcd_client=None):
+        self.name = name
+        self.ttl = ttl
+        if etcd_client is not None:
+            self.etcd_client = etcd_client
+        self.key = name
+        self.lease = None
+        # store uuid as bytes, since it avoids having to decode each time we
+        # need to compare. Using uuid4 for better privacy (no MAC address).
+        self.uuid = uuid.uuid4().bytes
+
+    def acquire(self, timeout=10):
+        """Acquire the lock.
+
+        :params timeout: Maximum time to wait before returning. `None` means
+                         forever, any other value equal or greater than 0 is
+                         the number of seconds.
+        :returns: True if the lock has been acquired, False otherwise.
+
+        """
+        if timeout is not None:
+            deadline = time.time() + timeout
+
+        while True:
+            if self._try_acquire():
+                return True
+
+            if timeout is not None:
+                remaining_timeout = max(deadline - time.time(), 0)
+                if remaining_timeout == 0:
+                    return False
+            else:
+                remaining_timeout = None
+
+            # Wait for the lock to be deleted, then retry immediately
+            deleted = self._wait_delete_event(remaining_timeout)
+            if not deleted:
+                # Timeout or error occurred, exit loop
+                return False
+
+    def _try_acquire(self):
+        # Use a longer TTL during acquisition to avoid race condition
+        # where lease expires between creation and transaction execution
+        acquisition_ttl = max(self.ttl, 10)  # Minimum 10 seconds for safety
+        self.lease = self.etcd_client.lease(acquisition_ttl)
+
+        try:
+            success, metadata = self.etcd_client.transaction(
+                compare=[self.etcd_client.transactions.create(self.key) == 0],
+                success=[
+                    self.etcd_client.transactions.put(self.key, self.uuid, lease=self.lease)
+                ],
+                failure=[self.etcd_client.transactions.get(self.key)],
+            )
+            if success is True:
+                self.revision = metadata[0].header.revision
+                # Refresh lease immediately to ensure it has the full TTL
+                # and to verify the lease is still valid
+                try:
+                    next(self.lease.refresh())
+                except StopIteration:
+                    # If refresh fails, the lock acquisition failed
+                    self.release()
+                    return False
+                return True
+            self.revision = metadata[0][0][1].mod_revision
+            # Clean up the lease that was created but not needed
+            self.etcd_client.revoke_lease(self.lease.id)
+            self.lease = None
+            return False
+        except Exception:
+            # Clean up lease on any error to avoid resource leak
+            if self.lease is not None:
+                try:
+                    self.etcd_client.revoke_lease(self.lease.id)
+                except Exception:
+                    pass  # Ignore cleanup errors
+                self.lease = None
+            raise
+
+    def _wait_delete_event(self, timeout):
+        """Wait for the lock key to be deleted.
+
+        :param timeout: Maximum time to wait in seconds, or None for infinite.
+        :returns: True if the key was deleted, False if timeout occurred.
+
+        This method handles the race condition where the lock may be deleted
+        between the time we try to acquire it and when we create the watch.
+        """
+        # FIRST: Check if key still exists and get its current revision.
+        # This must be done BEFORE creating the watch to avoid a race window.
+        value, metadata = self.etcd_client.get(self.key)
+        if value is None:
+            # Key already deleted, return True so acquire() will retry
+            return True
+
+        # Use the current key's mod_revision to ensure we don't miss events
+        # The watch should start from the revision AFTER the key was last modified
+        current_revision = metadata.mod_revision
+        start_revision = current_revision + 1
+
+        try:
+            event_iter, cancel = self.etcd_client.watch(
+                self.key, start_revision=start_revision
+            )
+        except exceptions.WatchTimedOut:
+            return False
+
+        # Double-check after creating watch (very small window, but possible)
+        value, _ = self.etcd_client.get(self.key)
+        if value is None:
+            cancel()
+            return True
+
+        if timeout is not None:
+            timer = threading.Timer(timeout, cancel)
+            timer.start()
+        else:
+            timer = None
+
+        try:
+            for event in event_iter:
+                if isinstance(event, events.DeleteEvent):
+                    return True
+        finally:
+            # Ensure cleanup happens regardless of how we exit
+            if timer is not None:
+                timer.cancel()
+            cancel()
+
+        return False
+
+    def release(self, raise_error=False):
+        """Release the lock.
+
+        :param raise_error: If True, raise an exception on release failure.
+                           If False (default), return a status string.
+        :returns: One of 'released' (successfully released),
+                  'not_owner' (lock was held by someone else or expired),
+                  'not_found' (lock key didn't exist),
+                  or 'error' (transaction failed for another reason).
+        :raises LockReleaseError: If raise_error=True and release fails.
+        """
+        try:
+            success, responses = self.etcd_client.transaction(
+                compare=[self.etcd_client.transactions.value(self.key) == self.uuid],
+                success=[self.etcd_client.transactions.delete(self.key)],
+                failure=[self.etcd_client.transactions.get(self.key)],
+            )
+
+            if success:
+                return 'released'
+
+            # Transaction succeeded but comparison failed
+            # Check if key exists and who owns it
+            if responses and responses[0]:
+                current_value = responses[0][0][0] if responses[0][0] else None
+                if current_value is None:
+                    if raise_error:
+                        raise exceptions.LockReleaseError(
+                            "Lock key does not exist - may have expired or been released"
+                        )
+                    return 'not_found'
+                if current_value != self.uuid:
+                    if raise_error:
+                        raise exceptions.LockReleaseError(
+                            "Lock is held by another owner - cannot release"
+                        )
+                    return 'not_owner'
+
+            if raise_error:
+                raise exceptions.LockReleaseError("Unknown lock release failure")
+            return 'error'
+
+        except exceptions.Etcd3Exception:
+            if raise_error:
+                raise
+            return 'error'
+        except Exception as exc:
+            if raise_error:
+                raise exceptions.LockReleaseError(f"Lock release failed: {exc}") from exc
+            return 'error'
+
+    def refresh(self):
+        """Refresh the time to live on this lock."""
+        if self.lease is not None:
+            return self.lease.refresh()
+        else:
+            raise ValueError(
+                "No lease associated with this lock - have you acquired the lock yet?"
+            )
+
+    def is_acquired(self):
+        """Check if this lock is currently acquired."""
+        uuid, _ = self.etcd_client.get(self.key)
+
+        if uuid is None:
+            return False
+
+        return uuid == self.uuid
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.release()
+
+
+class AsyncLock(object):
+    """
+    An asynchronous distributed lock.
+
+    This can be used as an async context manager:
+
+    .. code-block:: python
+
+        etcd = etcd3.async_client()
+
+        # create a lock that expires after 20 seconds
+        async with await etcd.lock('/myapp/mylock', ttl=20) as lock:
+            # do something that requires the lock
+            print(await lock.is_acquired())
+
+            # refresh the timeout on the lease
+            await lock.refresh()
+
+    :param name: name (full path) of the lock. The full key path should be provided,
+                 e.g., '/myapp/mylock' or 'election/my-service'.
+    :type name: string or bytes
+    :param ttl: length of time for the lock to live for in seconds. The lock
+                will be released after this time elapses, unless refreshed
+    :type ttl: int
+    """
+
+    def __init__(self, name, ttl=60, etcd_client=None):
+        self.name = name
+        self.ttl = ttl
+        if etcd_client is not None:
+            self.etcd_client = etcd_client
+        self.key = name
+        self.lease = None
+        # store uuid as bytes, since it avoids having to decode each time we
+        # need to compare. Using uuid4 for better privacy (no MAC address).
+        self.uuid = uuid.uuid4().bytes
+
+    async def acquire(self, timeout=10):
+        """Acquire the lock asynchronously.
+
+        :params timeout: Maximum time to wait before returning. `None` means
+                         forever, any other value equal or greater than 0 is
+                         the number of seconds.
+        :returns: True if the lock has been acquired, False otherwise.
+
+        """
+        if timeout is not None:
+            deadline = time.time() + timeout
+
+        while True:
+            if await self._try_acquire():
+                return True
+
+            if timeout is not None:
+                remaining_timeout = max(deadline - time.time(), 0)
+                if remaining_timeout == 0:
+                    return False
+            else:
+                remaining_timeout = None
+
+            # Wait for the lock to be deleted, then retry immediately
+            deleted = await self._wait_delete_event(remaining_timeout)
+            if not deleted:
+                # Timeout or error occurred, exit loop
+                return False
+
+    def _set_lock_span_attributes(self, acquired: bool = None):
+        """Set lock-related attributes on current span.
+        
+        :param acquired: Whether the lock was acquired (optional)
+        """
+        if not is_tracing_available():
+            return
+        
+        span = get_current_span()
+        if span is None:
+            return
+        
+        attrs = build_lock_attributes(
+            name=self.key if isinstance(self.key, str) else self.key.decode('utf-8', errors='replace'),
+            uuid=self.uuid,
+            acquired=acquired
+        )
+        
+        for key, value in attrs.items():
+            if value is not None:
+                span.set_attribute(key, value)
+
+    async def _try_acquire(self):
+        # Use a longer TTL during acquisition to avoid race condition
+        # where lease expires between creation and transaction execution
+        acquisition_ttl = max(self.ttl, 10)  # Minimum 10 seconds for safety
+        self.lease = await self.etcd_client.lease(acquisition_ttl)
+
+        try:
+            success, metadata = await self.etcd_client.transaction(
+                compare=[self.etcd_client.transactions.create(self.key) == 0],
+                success=[
+                    self.etcd_client.transactions.put(self.key, self.uuid, lease=self.lease)
+                ],
+                failure=[self.etcd_client.transactions.get(self.key)],
+            )
+            if success is True:
+                self.revision = metadata[0].header.revision
+                # Refresh lease immediately to ensure it has the full TTL
+                # and to verify the lease is still valid
+                try:
+                    async for _ in self.lease.refresh():
+                        break  # Just need one response to verify lease is alive
+                except Exception:
+                    # If refresh fails, the lock acquisition failed
+                    await self.release()
+                    self._set_lock_span_attributes(acquired=False)
+                    return False
+                self._set_lock_span_attributes(acquired=True)
+                return True
+            self.revision = metadata[0][0][1].mod_revision
+            # Clean up the lease that was created but not needed
+            await self.etcd_client.revoke_lease(self.lease.id)
+            self.lease = None
+            return False
+        except Exception:
+            # Clean up lease on any error to avoid resource leak
+            if self.lease is not None:
+                try:
+                    await self.etcd_client.revoke_lease(self.lease.id)
+                except Exception:
+                    pass  # Ignore cleanup errors
+                self.lease = None
+            raise
+
+    async def _wait_delete_event(self, timeout):
+        """Wait for the lock key to be deleted.
+
+        :param timeout: Maximum time to wait in seconds, or None for infinite.
+        :returns: True if the key was deleted, False if timeout occurred.
+
+        This method handles the race condition where the lock may be deleted
+        between the time we try to acquire it and when we create the watch.
+        """
+        # FIRST: Check if key still exists and get its current revision.
+        # This must be done BEFORE creating the watch to avoid a race window.
+        value, metadata = await self.etcd_client.get(self.key)
+        if value is None:
+            # Key already deleted, return True so acquire() will retry
+            return True
+
+        # Use the current key's mod_revision to ensure we don't miss events
+        # The watch should start from the revision AFTER the key was last modified
+        current_revision = metadata.mod_revision
+        start_revision = current_revision + 1
+
+        try:
+            event_iter, cancel = await self.etcd_client.watch(
+                self.key, start_revision=start_revision
+            )
+        except exceptions.WatchTimedOut:
+            return False
+
+        # Double-check after creating watch (very small window, but possible)
+        value, _ = await self.etcd_client.get(self.key)
+        if value is None:
+            cancel()
+            return True
+
+        async def wait_for_delete():
+            """Wait for delete event"""
+            async for event in event_iter:
+                if isinstance(event, events.DeleteEvent):
+                    return True
+            return False
+
+        try:
+            if timeout is not None:
+                return await asyncio.wait_for(wait_for_delete(), timeout=timeout)
+            else:
+                return await wait_for_delete()
+        except asyncio.TimeoutError:
+            return False
+        finally:
+            # Ensure watch is cancelled
+            cancel()
+
+    async def release(self, raise_error=False):
+        """Release the lock asynchronously.
+
+        :param raise_error: If True, raise an exception on release failure.
+                           If False (default), return a status string.
+        :returns: One of 'released' (successfully released),
+                  'not_owner' (lock was held by someone else or expired),
+                  'not_found' (lock key didn't exist),
+                  or 'error' (transaction failed for another reason).
+        :raises LockReleaseError: If raise_error=True and release fails.
+        """
+        try:
+            success, responses = await self.etcd_client.transaction(
+                compare=[self.etcd_client.transactions.value(self.key) == self.uuid],
+                success=[self.etcd_client.transactions.delete(self.key)],
+                failure=[self.etcd_client.transactions.get(self.key)],
+            )
+
+            if success:
+                self._set_lock_span_attributes(acquired=False)
+                return 'released'
+
+            # Transaction succeeded but comparison failed
+            # Check if key exists and who owns it
+            if responses and responses[0]:
+                current_value = responses[0][0][0] if responses[0][0] else None
+                if current_value is None:
+                    self._set_lock_span_attributes(acquired=False)
+                    if raise_error:
+                        raise exceptions.LockReleaseError(
+                            "Lock key does not exist - may have expired or been released"
+                        )
+                    return 'not_found'
+                if current_value != self.uuid:
+                    self._set_lock_span_attributes(acquired=None)  # Owned by another
+                    if raise_error:
+                        raise exceptions.LockReleaseError(
+                            "Lock is held by another owner - cannot release"
+                        )
+                    return 'not_owner'
+
+            self._set_lock_span_attributes(acquired=None)
+            if raise_error:
+                raise exceptions.LockReleaseError("Unknown lock release failure")
+            return 'error'
+
+        except exceptions.Etcd3Exception:
+            self._set_lock_span_attributes(acquired=None)
+            if raise_error:
+                raise
+            return 'error'
+        except Exception as exc:
+            if raise_error:
+                raise exceptions.LockReleaseError(f"Lock release failed: {exc}") from exc
+            return 'error'
+
+    def refresh(self):
+        """Refresh the time to lease on this lock.
+
+        Returns an async generator that yields keep-alive responses.
+        The caller should iterate over the result to keep the lease alive.
+
+        Example:
+            async for response in lock.refresh():
+                # Lease is kept alive while iterating
+                pass
+        """
+        if self.lease is not None:
+            return self.lease.refresh()
+        else:
+            raise ValueError(
+                "No lease associated with this lock - have you acquired the lock yet?"
+            )
+
+    @asynccontextmanager
+    async def auto_refresh(self, interval=None):
+        """Async context manager that automatically refreshes the lock in background.
+
+        This context manager starts a background task that continuously refreshes
+        the lock lease while the context is active. This is useful for long-running
+        operations that need to maintain lock ownership.
+
+        Example:
+            async with lock.auto_refresh(interval=10):
+                # Lock is automatically refreshed every 10 seconds
+                await do_long_work()
+            # Lock is released when exiting the context
+
+        :param interval: Refresh interval in seconds. If None, defaults to TTL/3.
+        :type interval: float, optional
+        :yields: The lock instance
+        :rtype: AsyncLock
+        """
+        if self.lease is None:
+            raise ValueError(
+                "No lease associated with this lock - have you acquired the lock yet?"
+            )
+
+        # Default interval is TTL/3 to ensure timely refresh
+        refresh_interval = interval or (self.ttl / 3)
+        stop_event = asyncio.Event()
+
+        async def refresh_loop():
+            """Background task that keeps refreshing the lease."""
+            while not stop_event.is_set():
+                try:
+                    # Get the async generator from refresh
+                    refresh_gen = self.refresh()
+                    # Get one response from the generator to trigger keepalive
+                    try:
+                        await asyncio.wait_for(
+                            refresh_gen.__anext__(), timeout=5.0
+                        )
+                        # Successfully refreshed
+                    except StopAsyncIteration:
+                        # Generator ended unexpectedly
+                        break
+                    except asyncio.TimeoutError:
+                        # No response received, will retry
+                        pass
+                    finally:
+                        # Close the generator to release resources
+                        await refresh_gen.aclose()
+                except Exception:
+                    # If refresh fails, stop the loop
+                    break
+                # Wait before next refresh
+                try:
+                    await asyncio.wait_for(stop_event.wait(), timeout=refresh_interval)
+                except asyncio.TimeoutError:
+                    # Time for next refresh
+                    continue
+
+        # Start background refresh task
+        task = asyncio.create_task(refresh_loop())
+        try:
+            yield self
+        finally:
+            # Signal the background task to stop
+            stop_event.set()
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
+    async def is_acquired(self):
+        """Check if this lock is currently acquired asynchronously."""
+        uuid, _ = await self.etcd_client.get(self.key)
+
+        if uuid is None:
+            return False
+
+        return uuid == self.uuid
+
+    async def __aenter__(self):
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exception_type, exception_value, traceback):
+        await self.release()