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_with_timeout( 

56 async_semaphore: asyncio.Semaphore, 

57 timeout: float, 

58) -> Result[None, OverloadError]: 

59 try: 

60 async with asyncio.timeout(timeout): 

61 await async_semaphore.acquire() 

62 return Ok(None) 

63 except TimeoutError: 

64 return Err(OverloadError()) 

65 

66 

67async def _acquire_no_timeout( 

68 async_semaphore: asyncio.Semaphore, 

69) -> Result[None, OverloadError]: 

70 if async_semaphore.locked(): 

71 return Err(OverloadError()) 

72 await async_semaphore.acquire() 

73 return Ok(None) 

74 

75 

76async def _acquire_async_semaphore( 

77 async_semaphore: asyncio.Semaphore, 

78 timeout: float, 

79) -> Result[None, OverloadError]: 

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

81 try: 

82 if timeout <= 0.0: 

83 return await _acquire_no_timeout(async_semaphore) 

84 return await _acquire_with_timeout(async_semaphore, timeout) 

85 except Exception as e: 

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

87 

88 

89def _handle_async_concurrency( 

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

91 max_tasks: int, 

92 timeout: float, 

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

94 """Handle asynchronous concurrency limiting.""" 

95 async_semaphore = asyncio.Semaphore(max_tasks) 

96 

97 @functools.wraps(func) 

98 async def async_wrapper( 

99 *args: P.args, 

100 **kwargs: P.kwargs, 

101 ) -> Result[T, OverloadError]: 

102 acquire_result = await _acquire_async_semaphore(async_semaphore, timeout) 

103 if isinstance(acquire_result, Err): 

104 return acquire_result 

105 

106 try: 

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

108 finally: 

109 async_semaphore.release() 

110 

111 return async_wrapper # type: ignore[misc] 

112 

113 

114def _acquire_sync_semaphore( 

115 sync_semaphore: threading.Semaphore, 

116 timeout: float, 

117) -> Result[None, OverloadError]: 

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

119 try: 

120 if timeout <= 0.0: 

121 acquired = sync_semaphore.acquire(blocking=False) 

122 else: 

123 acquired = sync_semaphore.acquire(timeout=timeout) 

124 

125 if not acquired: 

126 return Err(OverloadError()) 

127 return Ok(None) 

128 except Exception as e: 

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

130 

131 

132def _handle_sync_concurrency( 

133 func: Callable[P, T], 

134 max_tasks: int, 

135 timeout: float, 

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

137 """Handle synchronous concurrency limiting.""" 

138 sync_semaphore = threading.Semaphore(max_tasks) 

139 

140 @functools.wraps(func) 

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

142 acquire_result = _acquire_sync_semaphore(sync_semaphore, timeout) 

143 if isinstance(acquire_result, Err): 

144 return acquire_result 

145 

146 try: 

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

148 finally: 

149 sync_semaphore.release() 

150 

151 return wrapper 

152 

153 

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

155 """Validate max_tasks argument.""" 

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

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

158 

159 

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

161 """Validate timeout argument.""" 

162 if ( 

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

164 or not math.isfinite(timeout) 

165 or timeout < 0.0 

166 ): 

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

168 

169 

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

171 """Validate arguments for limit_concurrency.""" 

172 _validate_max_tasks(max_tasks) 

173 _validate_timeout(timeout) 

174 

175 

176def limit_concurrency( 

177 max_tasks: int, 

178 timeout: float = 0.0, 

179) -> ConcurrencyLimitDecorator: 

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

181 

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

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

184 an ``Err(OverloadError)``. 

185 

186 Args: 

187 max_tasks: Maximum concurrent function executions allowed. 

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

189 

190 Returns: 

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

192 

193 Example: 

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

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

196 ... return "data" 

197 >>> process_data() 

198 Ok('data') 

199 

200 """ 

201 _validate_limit_concurrency_args(max_tasks, timeout) 

202 

203 def decorator( 

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

205 ) -> ( 

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

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

208 ): 

209 if inspect.iscoroutinefunction(func): 

210 return _handle_async_concurrency( 

211 func, # type: ignore[misc] 

212 max_tasks, 

213 timeout, 

214 ) 

215 

216 return _handle_sync_concurrency( 

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

218 max_tasks, 

219 timeout, 

220 ) 

221 

222 return cast(ConcurrencyLimitDecorator, decorator)