Coverage for /usr/local/lib/python3.10/site-packages/opal_common-0.0.0-py3.10.egg/opal_common/async_utils.py: 38%

1from __future__ import annotations 

2 

3import asyncio 

4import sys 

5from functools import partial 

6from typing import Any, Callable, Coroutine, Optional, Set, Tuple, TypeVar 

7 

8import loguru 

9from loguru import logger 

10 

11if sys.version_info < (3, 10): 11 ↛ 12line 11 didn't jump to line 12 because the condition on line 11 was never true

12 from typing_extensions import ParamSpec 

13else: 

14 from typing import ParamSpec 

15 

16T_result = TypeVar("T_result") 

17P_args = ParamSpec("P_args") 

18 

19 

20async def run_sync( 

21 func: Callable[P_args, T_result], *args: P_args.args, **kwargs: P_args.kwargs 

22) -> T_result: 

23 """Shorthand for running a sync function in an executor within an async 

24 context. 

25 

26 For example: 

27 def sync_function_that_takes_time_to_run(arg1, arg2): 

28 time.sleep(5) 

29 

30 async def async_function(): 

31 await run_sync(sync_function_that_takes_time_to_run, 1, arg2=5) 

32 """ 

33 return await asyncio.get_event_loop().run_in_executor( 

34 None, partial(func, *args, **kwargs) 

35 ) 

36 

37 

38class TakeANumberQueue: 

39 """Enables a task to hold a place in queue prior to having the actual item 

40 to be sent over the queue. 

41 

42 The goal is executing concurrent tasks while still processing their 

43 results by the original order of execution 

44 """ 

45 

46 class Number: 

47 def __init__(self): 

48 self._event = asyncio.Event() 

49 self._item = None 

50 

51 def put(self, item: Any): 

52 self._item = item 

53 self._event.set() 

54 

55 async def get(self) -> Any: 

56 await self._event.wait() 

57 return self._item 

58 

59 def __init__(self, logger: loguru.Logger): 

60 self._queue: asyncio.Queue | None = None 

61 self._logger = logger 

62 

63 async def take_a_number(self) -> Number: 

64 assert self._queue is not None, "Queue not initialized" 

65 n = TakeANumberQueue.Number() 

66 await self._queue.put(n) 

67 return n 

68 

69 async def get(self) -> Any: 

70 n: TakeANumberQueue.Number = await self._queue.get() 

71 return await n.get() # Wait for next in line to have a result 

72 

73 async def _handle_queue(self, handler: Coroutine): 

74 self._queue = asyncio.Queue() 

75 while True: 

76 try: 

77 item = await self.get() 

78 await handler(item) 

79 except asyncio.CancelledError: 

80 if self._logger: 

81 self._logger.debug("queue handling task cancelled") 

82 return 

83 except Exception: 

84 if self._logger: 

85 self._logger.exception("failed handling take-a-number queue item") 

86 

87 async def start_queue_handling(self, handler: Coroutine): 

88 self._handler_task = asyncio.create_task(self._handle_queue(handler)) 

89 

90 async def stop_queue_handling(self): 

91 if self._handler_task: 

92 self._handler_task.cancel() 

93 self._handler_task = None 

94 

95 

96class TasksPool: 

97 def __init__(self): 

98 self._tasks: Set[asyncio.Task] = set() 

99 self._running = True 

100 

101 def _cleanup_task(self, done_task): 

102 self._tasks.discard(done_task) 

103 

104 def add_task(self, f): 

105 if not self._running: 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 raise RuntimeError("TasksPool is already shutdown") 

107 t = asyncio.create_task(f) 

108 self._tasks.add(t) 

109 t.add_done_callback(self._cleanup_task) 

110 

111 def restart(self): 

112 """Re-arms the pool so new tasks can be added after a shutdown. 

113 

114 Must only be called after shutdown() has completed, ensuring all 

115 previous tasks have finished and their done-callbacks have run. 

116 """ 

117 self._running = True 

118 self._tasks = set() 

119 

120 async def shutdown(self, force: bool = False): 

121 """Wait for them to finish. 

122 

123 :param force: If True, cancel all tasks immediately. 

124 """ 

125 self._running = False 

126 if force: 

127 for t in self._tasks: 

128 t.cancel() 

129 

130 results = await asyncio.gather( 

131 *self._tasks, 

132 return_exceptions=True, 

133 ) 

134 for result in results: 

135 if isinstance(result, Exception): 

136 logger.exception( 

137 "Error on task during shutdown of TasksPool: {result}", 

138 result=result, 

139 ) 

140 

141 

142async def repeated_call( 

143 func: Coroutine, 

144 seconds: float, 

145 *args: Tuple[Any], 

146 logger: Optional[loguru.Logger] = None, 

147): 

148 while True: 

149 try: 

150 await func(*args) 

151 await asyncio.sleep(seconds) 

152 except asyncio.CancelledError: 

153 raise 

154 except Exception as exc: 

155 logger.exception( 

156 "Error during repeated call to {func}: {exc}", 

157 func=func, 

158 exc=exc, 

159 )