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

1"""Resource pooling implementation."""

2

3 import queue

4 import threading

5 from dataclasses import dataclass

6 from datetime import datetime

7 from typing import Any

8

9 from dataknobs_fsm.functions.base import ResourceError

10 from dataknobs_fsm.resources.base import (

11 IResourceProvider,

12 ResourceMetrics,

13)

14

15

16@dataclass

17 class 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

29 class 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

62 class 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