Update condition.py,test_condition.py

Author: nadzhou

distributed/condition.py

@@ -3,46 +3,91 @@
 import asyncio
 import logging
 import uuid
-from collections import defaultdict
+from collections import defaultdict, deque
 
 from dask.utils import parse_timedelta
 
-from distributed.lock import Lock
 from distributed.utils import SyncMethodMixin, TimeoutError, log_errors, wait_for
 from distributed.worker import get_client
 
 logger = logging.getLogger(__name__)
 
 
 class ConditionExtension:
-    """Scheduler extension for managing Condition variable notifications
-
-    Coordinates wait/notify between distributed clients.
-    The lock itself is managed by LockExtension.
-    """
+    """Scheduler extension managing Condition lock and notifications"""
 
     def __init__(self, scheduler):
         self.scheduler = scheduler
-        # {condition_name: {waiter_id: asyncio.Event}}
-        self._waiters = defaultdict(dict)
+        # {condition_name: client_id} - who holds each lock
+        self._lock_holders = {}
+        # {condition_name: deque of (client_id, future)} - waiting to acquire
+        self._acquire_waiters = defaultdict(deque)
+        # {condition_name: {waiter_id: (client_id, Event)}} - waiting for notify
+        self._notify_waiters = defaultdict(dict)
 
         self.scheduler.handlers.update(
             {
+                "condition_acquire": self.acquire,
+                "condition_release": self.release,
                 "condition_wait": self.wait,
                 "condition_notify": self.notify,
                 "condition_notify_all": self.notify_all,
             }
         )
 
     @log_errors
-    async def wait(self, name=None, id=None, timeout=None):
-        """Register waiter and block until notified
+    async def acquire(self, name=None, client_id=None):
+        """Acquire lock - blocks until available"""
+        if name not in self._lock_holders:
+            # Lock is free
+            self._lock_holders[name] = client_id
+            return True
+
+        if self._lock_holders[name] == client_id:
+            # Already hold it (shouldn't happen in normal use)
+            return True
+
+        # Lock is held by someone else - wait our turn
+        future = asyncio.Future()
+        self._acquire_waiters[name].append((client_id, future))
+        await future
+        return True
+
+    @log_errors
+    async def release(self, name=None, client_id=None):
+        """Release lock"""
+        if name not in self._lock_holders:
+            raise RuntimeError("Released too often")
+
+        if self._lock_holders[name] != client_id:
+            raise RuntimeError("Cannot release lock held by another client")
 
-        Caller must have released the lock before calling this.
-        """
+        del self._lock_holders[name]
+
+        # Wake next waiter if any
+        waiters = self._acquire_waiters.get(name, deque())
+        while waiters:
+            next_client_id, future = waiters.popleft()
+            if not future.done():
+                self._lock_holders[name] = next_client_id
+                future.set_result(True)
+                break
+
+    @log_errors
+    async def wait(self, name=None, waiter_id=None, client_id=None, timeout=None):
+        """Release lock, wait for notify, reacquire lock"""
+        # Verify caller holds lock
+        if self._lock_holders.get(name) != client_id:
+            raise RuntimeError("wait() called without holding the lock")
+
+        # Release lock (waking next acquire waiter if any)
+        await self.release(name=name, client_id=client_id)
+
+        # Register as notify waiter
         event = asyncio.Event()
-        self._waiters[name][id] = event
+        self._notify_waiters[name][waiter_id] = (client_id, event)
 
+        # Wait for notification
         future = event.wait()
         if timeout is not None:
             future = wait_for(future, timeout)
@@ -53,65 +98,52 @@ async def wait(self, name=None, id=None, timeout=None):
         except TimeoutError:
             result = False
         finally:
-            self._waiters[name].pop(id, None)
-            if not self._waiters[name]:
-                del self._waiters[name]
+            # Cleanup waiter
+            self._notify_waiters[name].pop(waiter_id, None)
+            if not self._notify_waiters[name]:
+                del self._notify_waiters[name]
+
+            # Reacquire lock - blocks until available
+            await self.acquire(name=name, client_id=client_id)
 
         return result
 
     @log_errors
-    def notify(self, name=None, n=1):
+    def notify(self, name=None, client_id=None, n=1):
         """Wake up n waiters"""
-        waiters = self._waiters.get(name, {})
+        # Verify caller holds lock
+        if self._lock_holders.get(name) != client_id:
+            raise RuntimeError("notify() called without holding the lock")
+
+        waiters = self._notify_waiters.get(name, {})
         count = 0
-        for event in list(waiters.values())[:n]:
+        for _, (_, event) in list(waiters.items())[:n]:
             event.set()
             count += 1
         return count
 
     @log_errors
-    def notify_all(self, name=None):
+    def notify_all(self, name=None, client_id=None):
         """Wake up all waiters"""
-        waiters = self._waiters.get(name, {})
-        for event in waiters.values():
+        # Verify caller holds lock
+        if self._lock_holders.get(name) != client_id:
+            raise RuntimeError("notify_all() called without holding the lock")
+
+        waiters = self._notify_waiters.get(name, {})
+        for _, event in waiters.values():
             event.set()
         return len(waiters)
 
 
 class Condition(SyncMethodMixin):
-    """Distributed Condition Variable
-
-    Combines a Lock with wait/notify coordination across the cluster.
-
-    Parameters
-    ----------
-    name : str, optional
-        Name of the condition. Conditions with the same name share state.
-    client : Client, optional
-        Client for scheduler communication.
-
-    Examples
-    --------
-    Producer-consumer pattern:
-
-    >>> condition = Condition('data-ready')
-    >>> # Consumer
-    >>> async with condition:
-    ...     while not data_available():
-    ...         await condition.wait()
-    ...     process_data()
-
-    >>> # Producer
-    >>> async with condition:
-    ...     produce_data()
-    ...     condition.notify_all()
-    """
+    """Distributed Condition Variable"""
 
     def __init__(self, name=None, client=None):
         self.name = name or f"condition-{uuid.uuid4().hex}"
-        self.id = uuid.uuid4().hex
-        self._lock = Lock(name=f"{self.name}-lock")
+        self._waiter_id = uuid.uuid4().hex
+        self._client_id = uuid.uuid4().hex
         self._client = client
+        self._is_locked = False  # Track local state
 
     @property
     def client(self):
@@ -124,104 +156,72 @@ def client(self):
 
     @property
     def loop(self):
-        return self._lock.loop
+        return self.client.loop
 
     def _verify_running(self):
         if not self.client:
-            raise RuntimeError(
-                f"{type(self)} object not properly initialized. This can happen"
-                " if the object is being deserialized outside of the context of"
-                " a Client or Worker."
-            )
+            raise RuntimeError(f"{type(self)} object not properly initialized")
 
     async def acquire(self):
-        """Acquire the underlying lock"""
-        return await self._lock.acquire()
+        """Acquire lock"""
+        self._verify_running()
+        await self.client.scheduler.condition_acquire(
+            name=self.name, client_id=self._client_id
+        )
+        self._is_locked = True
 
     async def release(self):
-        """Release the underlying lock"""
-        await self._lock.release()
+        """Release lock"""
+        self._verify_running()
+        await self.client.scheduler.condition_release(
+            name=self.name, client_id=self._client_id
+        )
+        self._is_locked = False
 
     async def wait(self, timeout=None):
-        """Wait until notified
-
-        Must be called while lock is held. Atomically releases lock,
-        waits for notify(), then reacquires lock before returning.
-
-        Parameters
-        ----------
-        timeout : number or string or timedelta, optional
-            Maximum time to wait for notification.
-
-        Returns
-        -------
-        bool
-            True if notified, False if timeout occurred
-
-        Raises
-        ------
-        RuntimeError
-            If called without holding the lock
-        """
-        if not self._lock.locked():
+        """Wait for notification - atomically releases and reacquires lock"""
+        if not self._is_locked:
             raise RuntimeError("wait() called without holding the lock")
 
         self._verify_running()
         timeout = parse_timedelta(timeout)
 
-        # Atomically: release lock, wait for notify, reacquire lock
-        await self._lock.release()
-        try:
-            result = await self.client.scheduler.condition_wait(
-                name=self.name, id=self.id, timeout=timeout
-            )
-        finally:
-            await self._lock.acquire()
-
+        # This handles release, wait, reacquire atomically on scheduler
+        result = await self.client.scheduler.condition_wait(
+            name=self.name,
+            waiter_id=self._waiter_id,
+            client_id=self._client_id,
+            timeout=timeout,
+        )
+        # Lock is reacquired by the time this returns
         return result
 
     def notify(self, n=1):
-        """Wake up one or more waiters
-
-        Must be called while holding the lock.
-
-        Parameters
-        ----------
-        n : int, optional
-            Number of waiters to wake. Default is 1.
-
-        Returns
-        -------
-        int
-            Number of waiters actually notified
-        """
-        if not self._lock.locked():
+        """Wake up n waiters"""
+        if not self._is_locked:
             raise RuntimeError("notify() called without holding the lock")
         self._verify_running()
         return self.client.sync(
-            self.client.scheduler.condition_notify, name=self.name, n=n
+            self.client.scheduler.condition_notify,
+            name=self.name,
+            client_id=self._client_id,
+            n=n,
         )
 
     def notify_all(self):
-        """Wake up all waiters
-
-        Must be called while holding the lock.
-
-        Returns
-        -------
-        int
-            Number of waiters notified
-        """
-        if not self._lock.locked():
+        """Wake up all waiters"""
+        if not self._is_locked:
             raise RuntimeError("notify_all() called without holding the lock")
         self._verify_running()
         return self.client.sync(
-            self.client.scheduler.condition_notify_all, name=self.name
+            self.client.scheduler.condition_notify_all,
+            name=self.name,
+            client_id=self._client_id,
         )
 
     def locked(self):
-        """Return True if the lock is currently held"""
-        return self._lock.locked()
+        """Return True if lock is held by this instance"""
+        return self._is_locked
 
     async def __aenter__(self):
         await self.acquire()

distributed/tests/test_condition.py

@@ -286,16 +286,18 @@ async def test_condition_cleanup(c, s, a, b):
     """Test that condition state is cleaned up after use"""
     condition = Condition("cleanup-test")
 
-    # Check initial state - only check waiters since locks are managed by LockExtension
-    assert "cleanup-test" not in s.extensions["conditions"]._waiters
+    # Check initial state
+    assert "cleanup-test" not in s.extensions["conditions"]._lock_holders
+    assert "cleanup-test" not in s.extensions["conditions"]._notify_waiters
 
     # Use condition
     async with condition:
         condition.notify()
 
-    # Waiter state should be cleaned up
+    # State should be cleaned up
     await asyncio.sleep(0.1)
-    assert "cleanup-test" not in s.extensions["conditions"]._waiters
+    assert "cleanup-test" not in s.extensions["conditions"]._lock_holders
+    assert "cleanup-test" not in s.extensions["conditions"]._notify_waiters
 
 
 @gen_cluster(client=True)