Coverage for src/taipanstack/resilience/adaptive/bulkhead.py: 100%

1""" 

2Bulkhead pattern — concurrency isolation via semaphore. 

3 

4Limits the number of concurrent executions to prevent a single 

5failing dependency from consuming all available resources. 

6""" 

7 

8from __future__ import annotations 

9 

10import asyncio 

11import contextlib 

12import logging 

13import math 

14from collections.abc import Awaitable, Callable 

15from typing import ParamSpec, TypeVar 

16 

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

18 

19logger = logging.getLogger("taipanstack.resilience.adaptive.bulkhead") 

20 

21T = TypeVar("T") 

22P = ParamSpec("P") 

23 

24 

25class BulkheadFullError(Exception): 

26 """Raised when the bulkhead queue is at capacity.""" 

27 

28 def __init__(self, name: str, max_concurrent: int, max_queue: int) -> None: 

29 """Initialize BulkheadFullError. 

30 

31 Args: 

32 name: Bulkhead name. 

33 max_concurrent: Concurrency limit. 

34 max_queue: Queue limit. 

35 

36 """ 

37 self.bulkhead_name = name 

38 self.max_concurrent = max_concurrent 

39 self.max_queue = max_queue 

40 super().__init__( 

41 f"Bulkhead '{name}' is full " 

42 f"(max_concurrent={max_concurrent}, max_queue={max_queue})", 

43 ) 

44 

45 

46class Bulkhead: 

47 """Concurrency limiter using ``asyncio.Semaphore``. 

48 

49 Limits the number of concurrent executions of a callable. 

50 Excess callers are queued up to ``max_queue``; beyond that 

51 a ``BulkheadFullError`` is returned. 

52 

53 Args: 

54 name: Identifier for logging. 

55 max_concurrent: Maximum concurrent executions. 

56 max_queue: Maximum queued callers beyond concurrent limit. 

57 timeout: Seconds to wait for a permit before timing out. 

58 

59 Example: 

60 >>> bulk = Bulkhead("db", max_concurrent=5, max_queue=10) 

61 >>> result = await bulk.execute(fetch_data, user_id) 

62 

63 """ 

64 

65 def __init__( 

66 self, 

67 name: str = "default", 

68 *, 

69 max_concurrent: int = 10, 

70 max_queue: int = 50, 

71 timeout: float = 30.0, 

72 ) -> None: 

73 """Initialize the bulkhead. 

74 

75 Args: 

76 name: Bulkhead name. 

77 max_concurrent: Concurrency limit. 

78 max_queue: Queue limit. 

79 timeout: Permit acquisition timeout. 

80 

81 """ 

82 self.name = name 

83 self._max_concurrent = max_concurrent 

84 self._max_queue = max_queue 

85 if not math.isfinite(timeout) or timeout < 0: 

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

87 self._timeout = timeout 

88 self._semaphore = asyncio.Semaphore(max_concurrent) 

89 self._queued = 0 

90 self._active = 0 

91 

92 @property 

93 def available_permits(self) -> int: 

94 """Number of available concurrency permits.""" 

95 return self._max_concurrent - self._active 

96 

97 @property 

98 def queued(self) -> int: 

99 """Number of callers currently waiting in the queue.""" 

100 return self._queued 

101 

102 @property 

103 def active(self) -> int: 

104 """Number of currently executing tasks.""" 

105 return self._active 

106 

107 async def _cleanup_acquire_task( 

108 self, 

109 acquire_task: asyncio.Task[bool], 

110 ) -> None: 

111 """Helper to cancel an acquisition task and release semaphore if acquired.""" 

112 acquire_task.cancel() 

113 with contextlib.suppress(asyncio.CancelledError): 

114 await acquire_task 

115 self._semaphore.release() 

116 

117 async def _acquire_permit(self) -> Result[None, Exception]: 

118 """Wait for and acquire a concurrency permit.""" 

119 try: 

120 acquire_task = asyncio.create_task(self._semaphore.acquire()) 

121 try: 

122 await asyncio.wait_for( 

123 asyncio.shield(acquire_task), 

124 timeout=self._timeout, 

125 ) 

126 return Ok(None) 

127 except TimeoutError: 

128 await self._cleanup_acquire_task(acquire_task) 

129 return Err( 

130 TimeoutError( 

131 f"Bulkhead '{self.name}' timed out " 

132 f"after {self._timeout}s waiting for permit", 

133 ), 

134 ) 

135 except asyncio.CancelledError: 

136 await self._cleanup_acquire_task(acquire_task) 

137 raise 

138 except (RuntimeError, OSError, MemoryError) as e: 

139 return Err(RuntimeError(f"Resource exhaustion: {e!s}")) 

140 

141 async def execute( 

142 self, 

143 fn: Callable[P, Awaitable[T]], 

144 *args: P.args, 

145 **kwargs: P.kwargs, 

146 ) -> Result[T, Exception]: 

147 """Execute a callable within bulkhead limits. 

148 

149 Args: 

150 fn: Async callable to execute. 

151 *args: Positional arguments for fn. 

152 **kwargs: Keyword arguments for fn. 

153 

154 Returns: 

155 ``Ok(result)`` on success, ``Err`` on failure. 

156 

157 """ 

158 # Check queue capacity 

159 if self._queued >= self._max_queue: 

160 return Err( 

161 BulkheadFullError( 

162 self.name, 

163 self._max_concurrent, 

164 self._max_queue, 

165 ), 

166 ) 

167 

168 self._queued += 1 

169 try: 

170 permit_result = await self._acquire_permit() 

171 if isinstance(permit_result, Err): 

172 return permit_result 

173 finally: 

174 self._queued -= 1 

175 

176 # Execute within the permit 

177 self._active += 1 

178 try: 

179 result = await fn(*args, **kwargs) 

180 return Ok(result) 

181 except Exception as exc: 

182 logger.warning( 

183 "Bulkhead '%s' execution failed: %s", 

184 self.name, 

185 exc, 

186 ) 

187 return Err(exc) 

188 finally: 

189 self._active -= 1 

190 self._semaphore.release()