Coverage for src/dataknobs_fsm/resources/base.py: 61%

1"""Base interfaces and classes for resource management."""

2

3 from abc import ABC, abstractmethod

4 from dataclasses import dataclass

5 from datetime import datetime

6 from enum import Enum

7 from typing import Any, Dict, List, Protocol, runtime_checkable

8 from contextlib import contextmanager

9

10

11 class ResourceStatus(Enum):

12 """Status of a resource."""

13

14 IDLE = "idle"

15 ACTIVE = "active"

16 BUSY = "busy"

17 ERROR = "error"

18 CLOSED = "closed"

19 MAINTENANCE = "maintenance"

20

21

22 class ResourceHealth(Enum):

23 """Health status of a resource."""

24

25 HEALTHY = "healthy"

26 DEGRADED = "degraded"

27 UNHEALTHY = "unhealthy"

28 UNKNOWN = "unknown"

29

30

31@dataclass

32 class ResourceMetrics:

33 """Metrics for resource usage."""

34

35 total_acquisitions: int = 0

36 active_connections: int = 0

37 failed_acquisitions: int = 0

38 average_hold_time: float = 0.0

39 last_acquisition_time: datetime | None = None

40 last_release_time: datetime | None = None

41 health_check_failures: int = 0

42 last_health_check: datetime | None = None

43 average_acquisition_time: float = 0.0

44 total_timeout_events: int = 0

45 last_timeout_time: datetime | None = None

46

47 def record_acquisition(self, acquisition_time: float | None = None) -> None:

48 """Record a resource acquisition.

49

50 Args:

51 acquisition_time: Time taken to acquire the resource in seconds.

52 """

53 self.total_acquisitions += 1

54 self.active_connections += 1

55 self.last_acquisition_time = datetime.now()

56

57 # Update average acquisition time if provided

58 if acquisition_time is not None:

59 if self.average_acquisition_time == 0:

60 self.average_acquisition_time = acquisition_time

61 else:

62 # Rolling average

63 self.average_acquisition_time = (self.average_acquisition_time * 0.9) + (acquisition_time * 0.1)

64

65 def record_release(self, hold_time: float) -> None:

66 """Record a resource release.

67

68 Args:

69 hold_time: How long the resource was held in seconds.

70 """

71 self.active_connections = max(0, self.active_connections - 1)

72 self.last_release_time = datetime.now()

73

74 # Update average hold time

75 if self.average_hold_time == 0:

76 self.average_hold_time = hold_time

77 else:

78 # Rolling average

79 self.average_hold_time = (self.average_hold_time * 0.9) + (hold_time * 0.1)

80

81 def record_failure(self) -> None:

82 """Record a failed acquisition."""

83 self.failed_acquisitions += 1

84

85 def record_health_check(self, success: bool) -> None:

86 """Record a health check result.

87

88 Args:

89 success: Whether the health check passed.

90 """

91 self.last_health_check = datetime.now()

92 if not success:

93 self.health_check_failures += 1

94

95 def record_timeout(self) -> None:

96 """Record a timeout event."""

97 self.total_timeout_events += 1

98 self.last_timeout_time = datetime.now()

99 self.failed_acquisitions += 1

100

101

102@runtime_checkable

103 class IResourceProvider(Protocol):

104 """Interface for resource providers."""

105

106 def acquire(self, **kwargs) -> Any:

107 """Acquire a resource.

108

109 Args:

110 **kwargs: Provider-specific parameters.

111

112 Returns:

113 The acquired resource.

114

115 Raises:

116 ResourceError: If acquisition fails.

117 """

118 ...

119

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

121 """Release a resource.

122

123 Args:

124 resource: The resource to release.

125 """

126 ...

127

128 def validate(self, resource: Any) -> bool:

129 """Validate that a resource is still valid.

130

131 Args:

132 resource: The resource to validate.

133

134 Returns:

135 True if the resource is valid.

136 """

137 ...

138

139 def health_check(self) -> ResourceHealth:

140 """Check the health of the resource provider.

141

142 Returns:

143 The health status.

144 """

145 ...

146

147 def get_metrics(self) -> ResourceMetrics:

148 """Get resource metrics.

149

150 Returns:

151 Current metrics.

152 """

153 ...

154

155

156@runtime_checkable

157 class IResourcePool(Protocol):

158 """Interface for resource pools."""

159

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

161 """Acquire a resource from the pool.

162

163 Args:

164 timeout: Acquisition timeout in seconds.

165

166 Returns:

167 The acquired resource.

168

169 Raises:

170 ResourceError: If acquisition fails.

171 """

172 ...

173

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

175 """Return a resource to the pool.

176

177 Args:

178 resource: The resource to return.

179 """

180 ...

181

182 def size(self) -> int:

183 """Get the current pool size.

184

185 Returns:

186 Number of resources in the pool.

187 """

188 ...

189

190 def available(self) -> int:

191 """Get the number of available resources.

192

193 Returns:

194 Number of available resources.

195 """

196 ...

197

198 def close(self) -> None:

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

200 ...

201

202

203 class BaseResourceProvider(ABC):

204 """Base class for resource providers."""

205

206 def __init__(self, name: str, config: Dict[str, Any] | None = None):

207 """Initialize the provider.

208

209 Args:

210 name: Provider name.

211 config: Provider configuration.

212 """

213 self.name = name

214 self.config = config or {}

215 self.status = ResourceStatus.IDLE

216 self.metrics = ResourceMetrics()

217 self._resources: List[Any] = []

218

219 @abstractmethod

220 def acquire(self, **kwargs) -> Any:

221 """Acquire a resource.

222

223 Args:

224 **kwargs: Provider-specific parameters.

225

226 Returns:

227 The acquired resource.

228 """

229 pass

230

231 @abstractmethod

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

233 """Release a resource.

234

235 Args:

236 resource: The resource to release.

237 """

238 pass

239

240 def validate(self, resource: Any) -> bool:

241 """Validate a resource.

242

243 Args:

244 resource: The resource to validate.

245

246 Returns:

247 True if valid.

248 """

249 return resource is not None

250

251 def health_check(self) -> ResourceHealth:

252 """Check provider health.

253

254 Returns:

255 Health status.

256 """

257 if self.status == ResourceStatus.ERROR:

258 return ResourceHealth.UNHEALTHY

259 elif self.status == ResourceStatus.MAINTENANCE:

260 return ResourceHealth.DEGRADED

261 else:

262 return ResourceHealth.HEALTHY

263

264 def get_metrics(self) -> ResourceMetrics:

265 """Get provider metrics.

266

267 Returns:

268 Current metrics.

269 """

270 return self.metrics

271

272 @contextmanager

273 def resource_context(self, **kwargs):

274 """Context manager for resource acquisition.

275

276 Args:

277 **kwargs: Acquisition parameters.

278

279 Yields:

280 The acquired resource.

281 """

282 resource = None

283 start_time = datetime.now()

284 try:

285 resource = self.acquire(**kwargs)

286 self.metrics.record_acquisition()

287 yield resource

288 except Exception:

289 self.metrics.record_failure()

290 raise

291 finally:

292 if resource is not None:

293 hold_time = (datetime.now() - start_time).total_seconds()

294 self.release(resource)

295 self.metrics.record_release(hold_time)

296

297 def close(self) -> None:

298 """Close the provider and release all resources."""

299 for resource in self._resources[:]:

300 try:

301 self.release(resource)

302 except Exception:

303 pass # Best effort cleanup

304 self._resources.clear()

305 self.status = ResourceStatus.CLOSED