Coverage for src/taipanstack/utils/concurrency.py: 100%

1""" 

2Concurrency utilities. 

3 

4Provides a bulkhead pattern concurrency limiter decorator for both 

5synchronous and asynchronous functions. Uses an `OverloadError` and 

6returns a `Result` type. 

7""" 

8 

9import asyncio 

10import functools 

11import inspect 

12import math 

13import threading 

14from collections.abc import Awaitable, Callable 

15from typing import ParamSpec, Protocol, TypeVar, cast, overload 

16 

17from taipanstack.core.result import Err, Ok, Result 

18 

19__all__ = ["OverloadError", "limit_concurrency"] 

20 

21P = ParamSpec("P") 

22T = TypeVar("T") 

23 

24 

25class OverloadError(Exception): 

26 """Exception raised when a concurrency limit is exceeded or timed out.""" 

27 

28 def __init__(self, message: str = "Concurrency limit reached") -> None: 

29 """Initialize the OverloadError. 

30 

31 Args: 

32 message: The error message to display. 

33 Defaults to "Concurrency limit reached". 

34 

35 """ 

36 super().__init__(message) 

37 

38 

39class ConcurrencyLimitDecorator(Protocol): 

40 """Protocol for the concurrency limit decorator.""" 

41 

42 @overload 

43 def __call__( 

44 self, 

45 func: Callable[P, T], 

46 ) -> Callable[P, Result[T, OverloadError]]: ... 

47 

48 @overload 

49 def __call__( 

50 self, 

51 func: Callable[P, Awaitable[T]], 

52 ) -> Callable[P, Awaitable[Result[T, OverloadError]]]: ... 

53 

54 

55async def _acquire_async_semaphore( 

56 async_semaphore: asyncio.Semaphore, 

57 timeout: float, 

58) -> Result[None, OverloadError]: 

59 """Acquire async semaphore with optional timeout.""" 

60 try: 

61 if timeout <= 0.0: 

62 if async_semaphore.locked(): 

63 return Err(OverloadError()) 

64 await async_semaphore.acquire() 

65 return Ok(None) 

66 

67 try: 

68 async with asyncio.timeout(timeout): 

69 await async_semaphore.acquire() 

70 return Ok(None) 

71 except TimeoutError: 

72 return Err(OverloadError()) 

73 except Exception as e: 

74 return Err(OverloadError(f"Resource exhaustion: {e!s}")) 

75 

76 

77def _handle_async_concurrency( 

78 func: Callable[P, Awaitable[T]], 

79 max_tasks: int, 

80 timeout: float, 

81) -> Callable[P, Awaitable[Result[T, OverloadError]]]: 

82 """Handle asynchronous concurrency limiting.""" 

83 async_semaphore = asyncio.Semaphore(max_tasks) 

84 

85 @functools.wraps(func) 

86 async def async_wrapper( 

87 *args: P.args, 

88 **kwargs: P.kwargs, 

89 ) -> Result[T, OverloadError]: 

90 acquire_result = await _acquire_async_semaphore(async_semaphore, timeout) 

91 if isinstance(acquire_result, Err): 

92 return acquire_result 

93 

94 try: 

95 return Ok(await func(*args, **kwargs)) 

96 finally: 

97 async_semaphore.release() 

98 

99 return async_wrapper 

100 

101 

102def _acquire_sync_semaphore( 

103 sync_semaphore: threading.Semaphore, 

104 timeout: float, 

105) -> Result[None, OverloadError]: 

106 """Acquire sync semaphore with optional timeout.""" 

107 try: 

108 if timeout <= 0.0: 

109 acquired = sync_semaphore.acquire(blocking=False) 

110 else: 

111 acquired = sync_semaphore.acquire(timeout=timeout) 

112 

113 if not acquired: 

114 return Err(OverloadError()) 

115 return Ok(None) 

116 except Exception as e: 

117 return Err(OverloadError(f"Resource exhaustion: {e!s}")) 

118 

119 

120def _handle_sync_concurrency( 

121 func: Callable[P, T], 

122 max_tasks: int, 

123 timeout: float, 

124) -> Callable[P, Result[T, OverloadError]]: 

125 """Handle synchronous concurrency limiting.""" 

126 sync_semaphore = threading.Semaphore(max_tasks) 

127 

128 @functools.wraps(func) 

129 def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, OverloadError]: 

130 acquire_result = _acquire_sync_semaphore(sync_semaphore, timeout) 

131 if isinstance(acquire_result, Err): 

132 return acquire_result 

133 

134 try: 

135 return Ok(func(*args, **kwargs)) 

136 finally: 

137 sync_semaphore.release() 

138 

139 return wrapper 

140 

141 

142def _validate_max_tasks(max_tasks: int) -> None: 

143 """Validate max_tasks argument.""" 

144 if not isinstance(max_tasks, int) or max_tasks <= 0: 

145 raise ValueError("max_tasks must be > 0") 

146 

147 

148def _validate_timeout(timeout: float) -> None: 

149 """Validate timeout argument.""" 

150 if ( 

151 not isinstance(timeout, (int, float)) 

152 or not math.isfinite(timeout) 

153 or timeout < 0.0 

154 ): 

155 raise ValueError("timeout must be a finite non-negative number") 

156 

157 

158def _validate_limit_concurrency_args(max_tasks: int, timeout: float) -> None: 

159 """Validate arguments for limit_concurrency.""" 

160 _validate_max_tasks(max_tasks) 

161 _validate_timeout(timeout) 

162 

163 

164def limit_concurrency( 

165 max_tasks: int, 

166 timeout: float = 0.0, 

167) -> ConcurrencyLimitDecorator: 

168 """Decorate a function to apply the bulkhead concurrency limit pattern. 

169 

170 If the maximum concurrent executions are reached, the wrapper will wait up 

171 to `timeout` seconds to acquire a execution slot. If it fails, it returns 

172 an ``Err(OverloadError)``. 

173 

174 Args: 

175 max_tasks: Maximum concurrent function executions allowed. 

176 timeout: Maximum time in seconds to wait for a slot if limit is reached. 

177 

178 Returns: 

179 Decorated function returning a ``Result[T, OverloadError]``. 

180 

181 Example: 

182 >>> @limit_concurrency(max_tasks=2, timeout=0.1) 

183 ... def process_data() -> str: 

184 ... return "data" 

185 >>> process_data() 

186 Ok('data') 

187 

188 """ 

189 _validate_limit_concurrency_args(max_tasks, timeout) 

190 

191 def decorator( 

192 func: Callable[P, T] | Callable[P, Awaitable[T]], 

193 ) -> ( 

194 Callable[P, Result[T, OverloadError]] 

195 | Callable[P, Awaitable[Result[T, OverloadError]]] 

196 ): 

197 if inspect.iscoroutinefunction(func): 

198 return _handle_async_concurrency( 

199 func, 

200 max_tasks, 

201 timeout, 

202 ) 

203 

204 return _handle_sync_concurrency( 

205 cast(Callable[P, T], func), 

206 max_tasks, 

207 timeout, 

208 ) 

209 

210 return cast(ConcurrencyLimitDecorator, decorator)