Coverage for /usr/local/lib/python3.10/site-packages/opal_common-0.0.0-py3.10.egg/opal_common/synchronization/hierarchical_lock.py: 0%

1import asyncio 

2from contextlib import asynccontextmanager 

3from typing import Set 

4 

5from loguru import logger 

6 

7 

8class HierarchicalLock: 

9 """A hierarchical lock for asyncio. 

10 

11 - If a path is locked, no ancestor or descendant path can be locked. 

12 - Conversely, if a child path is locked, the parent path cannot be locked 

13 until all child paths are released. 

14 """ 

15 

16 def __init__(self): 

17 # locked_paths: set of currently locked string paths 

18 self._locked_paths: Set[str] = set() 

19 # Map of tasks to their acquired locks for re-entrant protection 

20 self._task_locks: dict[asyncio.Task, Set[str]] = {} 

21 # Internal lock for synchronizing access to locked_paths 

22 self._lock = asyncio.Lock() 

23 # Condition to wake up tasks when a path is released 

24 self._cond = asyncio.Condition(self._lock) 

25 

26 @staticmethod 

27 def _is_conflicting(p1: str, p2: str) -> bool: 

28 """Check if two paths conflict with each other.""" 

29 return p1 == p2 or p1.startswith(p2) or p2.startswith(p1) 

30 

31 async def acquire(self, path: str): 

32 """Acquire the lock for the given hierarchical path. 

33 

34 If an ancestor or descendant path is locked, this will wait 

35 until it is released. 

36 """ 

37 task = asyncio.current_task() 

38 if task is None: 

39 raise RuntimeError("acquire() must be called from within a task.") 

40 

41 async with self._lock: 

42 # Prevent re-entrant locking by the same task 

43 if path in self._task_locks.get(task, set()): 

44 raise RuntimeError(f"Task {task} cannot re-acquire lock on '{path}'.") 

45 

46 # Wait until there is no conflict with existing locked paths 

47 while any(self._is_conflicting(path, lp) for lp in self._locked_paths): 

48 logger.debug( 

49 f"Found conflicting path with {path!r}, waiting for release to check again..." 

50 ) 

51 # Condition.wait() releases the lock and waits for notify_all() 

52 await self._cond.wait() 

53 

54 # Acquire the path 

55 self._locked_paths.add(path) 

56 if task not in self._task_locks: 

57 self._task_locks[task] = set() 

58 self._task_locks[task].add(path) 

59 logger.debug("Acquired lock for path: {}", path) 

60 

61 async def release(self, path: str): 

62 """Release the lock for the given path and notify waiting tasks.""" 

63 task = asyncio.current_task() 

64 if task is None: 

65 raise RuntimeError("release() must be called from within a task.") 

66 

67 async with self._lock: 

68 if path not in self._locked_paths: 

69 raise RuntimeError(f"Cannot release path '{path}' that is not locked.") 

70 

71 if path not in self._task_locks.get(task, set()): 

72 raise RuntimeError( 

73 f"Task {task} cannot release lock on '{path}' it does not hold." 

74 ) 

75 

76 # Remove the path from locked paths and task locks 

77 self._locked_paths.remove(path) 

78 self._task_locks[task].remove(path) 

79 if not self._task_locks[task]: 

80 del self._task_locks[task] 

81 

82 # Notify all tasks that something was released 

83 self._cond.notify_all() 

84 logger.debug("Released lock for path: {}", path) 

85 

86 @asynccontextmanager 

87 async def lock(self, path: str) -> "HierarchicalLock": 

88 """Acquire the lock for the given path and return a context manager.""" 

89 await self.acquire(path) 

90 try: 

91 yield self 

92 finally: 

93 await self.release(path)