-
Notifications
You must be signed in to change notification settings - Fork 195
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added hierarchical lock for applying data updates (#745)
* Added hierarchical lock for applying data updates * Added docs and inline comments * Added _fetch_data to handle fetching errors * Removed line reference * Fixed test * Fixed imports * Fixed types for Python 3.9 support * Removed TaskGroup use (python compat) * Fixed import * Fixed types * Update packages/opal-client/opal_client/data/updater.py Co-authored-by: Or Weis <[email protected]> * Update packages/opal-client/opal_client/data/updater.py Co-authored-by: Or Weis <[email protected]> * Fixed pre-commit --------- Co-authored-by: Or Weis <[email protected]>
- Loading branch information
Showing
8 changed files
with
755 additions
and
258 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
85 changes: 85 additions & 0 deletions
85
packages/opal-common/opal_common/synchronization/hierarchical_lock.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| import asyncio | ||
| from contextlib import asynccontextmanager | ||
| from typing import Set | ||
|
|
||
|
|
||
| class HierarchicalLock: | ||
| """A hierarchical lock for asyncio. | ||
| - If a path is locked, no ancestor or descendant path can be locked. | ||
| - Conversely, if a child path is locked, the parent path cannot be locked | ||
| until all child paths are released. | ||
| """ | ||
|
|
||
| def __init__(self): | ||
| # locked_paths: set of currently locked string paths | ||
| self._locked_paths: Set[str] = set() | ||
| # Map of tasks to their acquired locks for re-entrant protection | ||
| self._task_locks: dict[asyncio.Task, Set[str]] = {} | ||
| # Internal lock for synchronizing access to locked_paths | ||
| self._lock = asyncio.Lock() | ||
| # Condition to wake up tasks when a path is released | ||
| self._cond = asyncio.Condition(self._lock) | ||
|
|
||
| @staticmethod | ||
| def _is_conflicting(p1: str, p2: str) -> bool: | ||
| """Check if two paths conflict with each other.""" | ||
| return p1 == p2 or p1.startswith(p2) or p2.startswith(p1) | ||
|
|
||
| async def acquire(self, path: str): | ||
| """Acquire the lock for the given hierarchical path. | ||
| If an ancestor or descendant path is locked, this will wait | ||
| until it is released. | ||
| """ | ||
| task = asyncio.current_task() | ||
| if task is None: | ||
| raise RuntimeError("acquire() must be called from within a task.") | ||
|
|
||
| async with self._lock: | ||
| # Prevent re-entrant locking by the same task | ||
| if path in self._task_locks.get(task, set()): | ||
| raise RuntimeError(f"Task {task} cannot re-acquire lock on '{path}'.") | ||
|
|
||
| # Wait until there is no conflict with existing locked paths | ||
| while any(self._is_conflicting(path, lp) for lp in self._locked_paths): | ||
| await self._cond.wait() | ||
|
|
||
| # Acquire the path | ||
| self._locked_paths.add(path) | ||
| if task not in self._task_locks: | ||
| self._task_locks[task] = set() | ||
| self._task_locks[task].add(path) | ||
|
|
||
| async def release(self, path: str): | ||
| """Release the lock for the given path and notify waiting tasks.""" | ||
| task = asyncio.current_task() | ||
| if task is None: | ||
| raise RuntimeError("release() must be called from within a task.") | ||
|
|
||
| async with self._lock: | ||
| if path not in self._locked_paths: | ||
| raise RuntimeError(f"Cannot release path '{path}' that is not locked.") | ||
|
|
||
| if path not in self._task_locks.get(task, set()): | ||
| raise RuntimeError( | ||
| f"Task {task} cannot release lock on '{path}' it does not hold." | ||
| ) | ||
|
|
||
| # Remove the path from locked paths and task locks | ||
| self._locked_paths.remove(path) | ||
| self._task_locks[task].remove(path) | ||
| if not self._task_locks[task]: | ||
| del self._task_locks[task] | ||
|
|
||
| # Notify all tasks that something was released | ||
| self._cond.notify_all() | ||
|
|
||
| @asynccontextmanager | ||
| async def lock(self, path: str) -> "HierarchicalLock": | ||
| """Acquire the lock for the given path and return a context manager.""" | ||
| await self.acquire(path) | ||
| try: | ||
| yield self | ||
| finally: | ||
| await self.release(path) |
Oops, something went wrong.