| 268 | |
| 269 | # class SpinPool(metaclass=Singleton): |
| 270 | class SpinPool(Pool): |
| 271 | def __init__(self, client_args: ClientArgs) -> None: |
| 272 | super().__init__(client_args) |
| 273 | random.shuffle(self.queue) |
| 274 | |
| 275 | def __enter__(self, *arg, **kwargs) -> ClientArgs: |
| 276 | """Get a client from the pool and return it as a context manager.""" |
| 277 | |
| 278 | if len(self) == 0: |
| 279 | raise APIKeyPoolEmptyError("api key pool is empty") |
| 280 | |
| 281 | return self.queue[self._index].client_args |
| 282 | |
| 283 | def __exit__(self, exc_type, exc_value, exc_traceback) -> bool: |
| 284 | """Put the resource back in the pool.""" |
| 285 | |
| 286 | if exc_value is None: |
| 287 | self._index = (self._index + 1) % len(self) |
| 288 | return True |
| 289 | else: |
| 290 | if ( |
| 291 | exc_type is RateLimitError |
| 292 | and exc_value.status_code == 429 |
| 293 | and judge_quota_exceeded(exc_value.message) |
| 294 | ): |
| 295 | self._remove_cur_arg() |
| 296 | else: |
| 297 | self._index = (self._index + 1) % len(self) |
| 298 | return False |
| 299 | |
| 300 | |
| 301 | def judge_quota_exceeded(message: str) -> bool: |