Kris
6cf9c8e9b6
1. 调整text_chunker.py中导入包的顺序 2. 调整config_watcher.py的typing导入顺序 3. 重构__init__.py的导入排序 4. 为CircuitBreaker新增事件回调机制，支持状态变化通知
119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from collections.abc import Awaitable, Callable
|
|
from enum import StrEnum
|
|
from typing import Any
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
class CircuitState(StrEnum):
|
|
CLOSED = "closed"
|
|
OPEN = "open"
|
|
HALF_OPEN = "half_open"
|
|
|
|
|
|
class CircuitBreakerOpenError(Exception):
|
|
pass
|
|
|
|
|
|
CircuitBreakerEventCallback = (
|
|
Callable[[str, str, dict[str, Any]], Awaitable[None]] | Callable[[str, str, dict[str, Any]], None] | None
|
|
)
|
|
|
|
|
|
class CircuitBreaker:
|
|
def __init__(
|
|
self,
|
|
failure_threshold: int = 5,
|
|
recovery_timeout: float = 60.0,
|
|
half_open_max_calls: int = 3,
|
|
channel_id: str = "",
|
|
operation: str = "default",
|
|
on_event: CircuitBreakerEventCallback = None,
|
|
):
|
|
self.failure_threshold = failure_threshold
|
|
self.recovery_timeout = recovery_timeout
|
|
self.half_open_max_calls = half_open_max_calls
|
|
self.channel_id = channel_id
|
|
self.operation = operation
|
|
self._on_event = on_event
|
|
|
|
self.state = CircuitState.CLOSED
|
|
self._failure_count = 0
|
|
self._last_failure_time: float = 0.0
|
|
self._half_open_calls = 0
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def call(self, func: Callable[[], Awaitable[Any]]) -> Any:
|
|
async with self._lock:
|
|
if self.state == CircuitState.OPEN:
|
|
elapsed = time.monotonic() - self._last_failure_time
|
|
if elapsed >= self.recovery_timeout:
|
|
self.state = CircuitState.HALF_OPEN
|
|
self._half_open_calls = 0
|
|
logger.info("Circuit breaker transitioned to HALF_OPEN")
|
|
await self._emit_event("on_half_open", {})
|
|
else:
|
|
raise CircuitBreakerOpenError(
|
|
f"Circuit breaker is OPEN, retry after {self.recovery_timeout - elapsed:.0f}s"
|
|
)
|
|
|
|
if self.state == CircuitState.HALF_OPEN:
|
|
if self._half_open_calls >= self.half_open_max_calls:
|
|
raise CircuitBreakerOpenError("Circuit breaker HALF_OPEN max calls reached")
|
|
self._half_open_calls += 1
|
|
|
|
try:
|
|
result = await func()
|
|
await self.record_success()
|
|
return result
|
|
except Exception:
|
|
await self.record_failure()
|
|
raise
|
|
|
|
async def record_success(self) -> None:
|
|
async with self._lock:
|
|
if self.state == CircuitState.HALF_OPEN:
|
|
self.state = CircuitState.CLOSED
|
|
self._failure_count = 0
|
|
self._half_open_calls = 0
|
|
logger.info("Circuit breaker closed after successful HALF_OPEN probe")
|
|
await self._emit_event("on_close", {})
|
|
elif self.state == CircuitState.CLOSED and self._failure_count > 0:
|
|
self._failure_count = 0
|
|
|
|
async def record_failure(self) -> None:
|
|
async with self._lock:
|
|
self._failure_count += 1
|
|
self._last_failure_time = time.monotonic()
|
|
|
|
if self.state == CircuitState.CLOSED and self._failure_count >= self.failure_threshold:
|
|
self.state = CircuitState.OPEN
|
|
logger.warning(
|
|
f"Circuit breaker OPEN after {self._failure_count} failures (threshold={self.failure_threshold})"
|
|
)
|
|
await self._emit_event("on_open", {"failure_count": self._failure_count})
|
|
elif self.state == CircuitState.HALF_OPEN:
|
|
self.state = CircuitState.OPEN
|
|
logger.warning("Circuit breaker re-OPENed after HALF_OPEN failure")
|
|
await self._emit_event("on_open", {"failure_count": self._failure_count, "from": "half_open"})
|
|
|
|
async def _emit_event(self, event: str, detail: dict[str, Any]) -> None:
|
|
if self._on_event is None:
|
|
return
|
|
payload = {
|
|
"channel_id": self.channel_id,
|
|
"operation": self.operation,
|
|
"state": self.state,
|
|
**detail,
|
|
}
|
|
try:
|
|
result = self._on_event(event, self.channel_id, payload) # type: ignore[call-arg]
|
|
if asyncio.iscoroutine(result):
|
|
await result
|
|
except Exception:
|
|
logger.debug(f"CircuitBreaker event callback error: {event}", exc_info=True)
|