Coverage for src/dataknobs_fsm/resources/pool.py: 22%

1"""Resource pooling implementation.""" 

2 

3import queue 

4import threading 

5from dataclasses import dataclass 

6from datetime import datetime 

7from typing import Any 

8 

9from dataknobs_fsm.functions.base import ResourceError 

10from dataknobs_fsm.resources.base import ( 

11 IResourceProvider, 

12 ResourceMetrics, 

13) 

14 

15 

16@dataclass 

17class PoolConfig: 

18 """Configuration for resource pools.""" 

19 

20 min_size: int = 1 

21 max_size: int = 10 

22 acquire_timeout: float = 30.0 

23 idle_timeout: float = 300.0 

24 validation_interval: float = 60.0 

25 max_lifetime: float = 3600.0 

26 

27 

28@dataclass 

29class PooledResource: 

30 """A resource in a pool with metadata.""" 

31 

32 resource: Any 

33 created_at: datetime 

34 last_used: datetime 

35 use_count: int = 0 

36 

37 def is_expired(self, max_lifetime: float) -> bool: 

38 """Check if resource has exceeded max lifetime. 

39  

40 Args: 

41 max_lifetime: Maximum lifetime in seconds. 

42  

43 Returns: 

44 True if expired. 

45 """ 

46 age = (datetime.now() - self.created_at).total_seconds() 

47 return age > max_lifetime 

48 

49 def is_idle_too_long(self, idle_timeout: float) -> bool: 

50 """Check if resource has been idle too long. 

51  

52 Args: 

53 idle_timeout: Maximum idle time in seconds. 

54  

55 Returns: 

56 True if idle too long. 

57 """ 

58 idle_time = (datetime.now() - self.last_used).total_seconds() 

59 return idle_time > idle_timeout 

60 

61 

62class ResourcePool: 

63 """Thread-safe resource pool implementation.""" 

64 

65 def __init__( 

66 self, 

67 provider: IResourceProvider, 

68 config: PoolConfig | None = None 

69 ): 

70 """Initialize the pool. 

71  

72 Args: 

73 provider: Resource provider. 

74 config: Pool configuration. 

75 """ 

76 self.provider = provider 

77 self.config = config or PoolConfig() 

78 self.metrics = ResourceMetrics() 

79 

80 self._pool = queue.Queue(maxsize=self.config.max_size) 

81 self._active_resources = set() 

82 self._lock = threading.RLock() 

83 self._closed = False 

84 self._resource_map = {} # Maps resource to PooledResource 

85 

86 # Initialize minimum resources 

87 self._initialize_pool() 

88 

89 def _initialize_pool(self) -> None: 

90 """Initialize the pool with minimum resources.""" 

91 for _ in range(self.config.min_size): 

92 try: 

93 resource = self.provider.acquire() 

94 pooled = PooledResource( 

95 resource=resource, 

96 created_at=datetime.now(), 

97 last_used=datetime.now() 

98 ) 

99 self._resource_map[id(resource)] = pooled 

100 self._pool.put(pooled) 

101 except Exception: 

102 pass # Continue with fewer resources 

103 

104 def acquire(self, timeout: float | None = None) -> Any: 

105 """Acquire a resource from the pool. 

106  

107 Args: 

108 timeout: Acquisition timeout in seconds. 

109  

110 Returns: 

111 The acquired resource. 

112  

113 Raises: 

114 ResourceError: If acquisition fails. 

115 """ 

116 if self._closed: 

117 raise ResourceError("Pool is closed", resource_name=self.provider.name, operation="acquire") 

118 

119 timeout = timeout or self.config.acquire_timeout 

120 start_time = datetime.now() 

121 

122 # Try to get from pool 

123 try: 

124 pooled = self._pool.get(timeout=timeout) 

125 

126 # Validate the resource 

127 if not self._validate_pooled_resource(pooled): 

128 # Resource is invalid, create a new one 

129 self._release_pooled_resource(pooled) 

130 return self._create_new_resource() 

131 

132 # Update metadata 

133 pooled.last_used = datetime.now() 

134 pooled.use_count += 1 

135 

136 with self._lock: 

137 self._active_resources.add(id(pooled.resource)) 

138 

139 # Track acquisition time 

140 acquisition_time = (datetime.now() - start_time).total_seconds() 

141 self.metrics.record_acquisition(acquisition_time) 

142 return pooled.resource 

143 

144 except queue.Empty: 

145 # Pool is empty, try to create new resource if under max 

146 with self._lock: 

147 if len(self._active_resources) < self.config.max_size: 

148 resource = self._create_new_resource() 

149 # Track acquisition time for newly created resource 

150 acquisition_time = (datetime.now() - start_time).total_seconds() 

151 self.metrics.record_acquisition(acquisition_time) 

152 return resource 

153 

154 # Record timeout event 

155 self.metrics.record_timeout() 

156 

157 raise ResourceError( 

158 f"Failed to acquire resource within {timeout} seconds", 

159 resource_name=self.provider.name, 

160 operation="acquire" 

161 ) from None 

162 

163 def release(self, resource: Any) -> None: 

164 """Return a resource to the pool. 

165  

166 Args: 

167 resource: The resource to return. 

168 """ 

169 if self._closed: 

170 # Pool is closed, just release the resource 

171 self.provider.release(resource) 

172 return 

173 

174 resource_id = id(resource) 

175 

176 with self._lock: 

177 if resource_id not in self._active_resources: 

178 return # Resource not from this pool 

179 

180 self._active_resources.discard(resource_id) 

181 pooled = self._resource_map.get(resource_id) 

182 

183 if pooled is None: 

184 # Create new pooled resource wrapper 

185 pooled = PooledResource( 

186 resource=resource, 

187 created_at=datetime.now(), 

188 last_used=datetime.now() 

189 ) 

190 self._resource_map[resource_id] = pooled 

191 

192 # Check if resource should be retired 

193 if (pooled.is_expired(self.config.max_lifetime) or 

194 not self.provider.validate(resource)): 

195 self._release_pooled_resource(pooled) 

196 else: 

197 # Return to pool 

198 pooled.last_used = datetime.now() 

199 try: 

200 self._pool.put_nowait(pooled) 

201 except queue.Full: 

202 # Pool is full, release the resource 

203 self._release_pooled_resource(pooled) 

204 

205 hold_time = 0.1 # Default hold time 

206 self.metrics.record_release(hold_time) 

207 

208 def _validate_pooled_resource(self, pooled: PooledResource) -> bool: 

209 """Validate a pooled resource. 

210  

211 Args: 

212 pooled: The pooled resource to validate. 

213  

214 Returns: 

215 True if valid. 

216 """ 

217 if pooled.is_expired(self.config.max_lifetime): 

218 return False 

219 

220 if pooled.is_idle_too_long(self.config.idle_timeout): 

221 return False 

222 

223 return self.provider.validate(pooled.resource) 

224 

225 def _create_new_resource(self) -> Any: 

226 """Create a new resource. 

227  

228 Returns: 

229 The new resource. 

230  

231 Raises: 

232 ResourceError: If creation fails. 

233 """ 

234 try: 

235 resource = self.provider.acquire() 

236 pooled = PooledResource( 

237 resource=resource, 

238 created_at=datetime.now(), 

239 last_used=datetime.now(), 

240 use_count=1 

241 ) 

242 

243 with self._lock: 

244 self._resource_map[id(resource)] = pooled 

245 self._active_resources.add(id(resource)) 

246 

247 # Note: Acquisition metrics are now tracked in acquire() method with timing 

248 return resource 

249 

250 except Exception as e: 

251 self.metrics.record_failure() 

252 raise ResourceError( 

253 f"Failed to create resource: {e}", 

254 resource_name=self.provider.name, 

255 operation="create" 

256 ) from e 

257 

258 def _release_pooled_resource(self, pooled: PooledResource) -> None: 

259 """Release a pooled resource. 

260  

261 Args: 

262 pooled: The pooled resource to release. 

263 """ 

264 import logging 

265 logger = logging.getLogger(__name__) 

266 

267 resource_id = id(pooled.resource) 

268 try: 

269 # Attempt to properly release the resource 

270 self.provider.release(pooled.resource) 

271 logger.debug(f"Successfully released pooled resource {resource_id} from pool {self.provider.name}") 

272 except AttributeError as e: 

273 logger.warning(f"Resource provider {self.provider.name} missing release method: {e}") 

274 except Exception as e: 

275 logger.error(f"Error releasing pooled resource {resource_id} from pool {self.provider.name}: {e}") 

276 # Track the error for debugging 

277 self.metrics.record_failure() 

278 # Still remove from map even if release failed to prevent memory leak 

279 

280 with self._lock: 

281 self._resource_map.pop(resource_id, None) 

282 

283 def size(self) -> int: 

284 """Get the current pool size. 

285  

286 Returns: 

287 Total number of resources (active + idle). 

288 """ 

289 with self._lock: 

290 return self._pool.qsize() + len(self._active_resources) 

291 

292 def available(self) -> int: 

293 """Get the number of available resources. 

294  

295 Returns: 

296 Number of idle resources in pool. 

297 """ 

298 return self._pool.qsize() 

299 

300 def close(self) -> None: 

301 """Close the pool and release all resources.""" 

302 self._closed = True 

303 

304 # Release active resources 

305 with self._lock: 

306 for resource_id in list(self._active_resources): 

307 pooled = self._resource_map.get(resource_id) 

308 if pooled: 

309 self._release_pooled_resource(pooled) 

310 self._active_resources.clear() 

311 

312 # Release pooled resources 

313 while not self._pool.empty(): 

314 try: 

315 pooled = self._pool.get_nowait() 

316 self._release_pooled_resource(pooled) 

317 except queue.Empty: 

318 break 

319 

320 self._resource_map.clear() 

321 

322 def evict_idle(self) -> int: 

323 """Evict idle resources that have exceeded timeout. 

324  

325 Returns: 

326 Number of resources evicted. 

327 """ 

328 evicted = 0 

329 temp_resources = [] 

330 

331 # Check all pooled resources 

332 while not self._pool.empty(): 

333 try: 

334 pooled = self._pool.get_nowait() 

335 if pooled.is_idle_too_long(self.config.idle_timeout): 

336 self._release_pooled_resource(pooled) 

337 evicted += 1 

338 else: 

339 temp_resources.append(pooled) 

340 except queue.Empty: 

341 break 

342 

343 # Put back the non-evicted resources 

344 for pooled in temp_resources: 

345 try: 

346 self._pool.put_nowait(pooled) 

347 except queue.Full: 

348 self._release_pooled_resource(pooled) 

349 

350 return evicted 

351 

352 def get_metrics(self) -> ResourceMetrics: 

353 """Get pool metrics. 

354  

355 Returns: 

356 Current metrics. 

357 """ 

358 return self.metrics