110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Any, Awaitable, Callable
|
||
|
|
|
||
|
|
from yuxi.channel.runtime.backoff import BackoffConfig
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
StartupTask = Callable[[], Awaitable[Any]]
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class StartupTaskResult:
|
||
|
|
name: str
|
||
|
|
success: bool
|
||
|
|
duration_ms: float
|
||
|
|
error: str | None = None
|
||
|
|
skipped: bool = False
|
||
|
|
|
||
|
|
|
||
|
|
_StartupBackoffConfig = BackoffConfig(
|
||
|
|
base_delay=0.5,
|
||
|
|
max_delay=5.0,
|
||
|
|
exponent=1.5,
|
||
|
|
jitter=True,
|
||
|
|
jitter_factor=0.1,
|
||
|
|
max_retries=3,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class StartupTaskRunner:
|
||
|
|
"""启动任务执行器 — 支持超时、重试、并行执行与结果追踪。"""
|
||
|
|
|
||
|
|
def __init__(self, default_timeout: float = 30.0, max_retries: int = 2, backoff_config: BackoffConfig | None = None):
|
||
|
|
self.default_timeout = default_timeout
|
||
|
|
self.max_retries = max_retries
|
||
|
|
self._backoff_config = backoff_config or _StartupBackoffConfig
|
||
|
|
self._results: list[StartupTaskResult] = []
|
||
|
|
self._results_lock = asyncio.Lock()
|
||
|
|
|
||
|
|
async def run(
|
||
|
|
self,
|
||
|
|
task: StartupTask,
|
||
|
|
*,
|
||
|
|
name: str | None = None,
|
||
|
|
timeout: float | None = None,
|
||
|
|
retries: int | None = None,
|
||
|
|
) -> StartupTaskResult:
|
||
|
|
task_name = name or getattr(task, "__name__", "unknown")
|
||
|
|
timeout_val = timeout if timeout is not None else self.default_timeout
|
||
|
|
retry_count = retries if retries is not None else self.max_retries
|
||
|
|
|
||
|
|
start = time.time()
|
||
|
|
last_error: Exception | None = None
|
||
|
|
|
||
|
|
for attempt in range(retry_count + 1):
|
||
|
|
try:
|
||
|
|
await asyncio.wait_for(task(), timeout=timeout_val)
|
||
|
|
duration_ms = (time.time() - start) * 1000
|
||
|
|
result = StartupTaskResult(
|
||
|
|
name=task_name,
|
||
|
|
success=True,
|
||
|
|
duration_ms=duration_ms,
|
||
|
|
)
|
||
|
|
async with self._results_lock:
|
||
|
|
self._results.append(result)
|
||
|
|
logger.info("Startup task '%s' completed in %.2fms", task_name, duration_ms)
|
||
|
|
return result
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
last_error = asyncio.TimeoutError(f"Task '{task_name}' timed out after {timeout_val}s")
|
||
|
|
logger.warning("Startup task '%s' timed out (attempt %d/%d)", task_name, attempt + 1, retry_count + 1)
|
||
|
|
except Exception as e:
|
||
|
|
last_error = e
|
||
|
|
logger.exception("Startup task '%s' failed (attempt %d/%d)", task_name, attempt + 1, retry_count + 1)
|
||
|
|
if attempt < retry_count:
|
||
|
|
await asyncio.sleep(self._backoff_config.compute_delay(attempt))
|
||
|
|
|
||
|
|
duration_ms = (time.time() - start) * 1000
|
||
|
|
result = StartupTaskResult(
|
||
|
|
name=task_name,
|
||
|
|
success=False,
|
||
|
|
duration_ms=duration_ms,
|
||
|
|
error=str(last_error) if last_error else "Unknown error",
|
||
|
|
)
|
||
|
|
async with self._results_lock:
|
||
|
|
self._results.append(result)
|
||
|
|
return result
|
||
|
|
|
||
|
|
async def run_parallel(
|
||
|
|
self,
|
||
|
|
tasks: list[tuple[StartupTask, str]],
|
||
|
|
*,
|
||
|
|
timeout: float | None = None,
|
||
|
|
) -> list[StartupTaskResult]:
|
||
|
|
"""并行执行多个启动任务。"""
|
||
|
|
coros = [self.run(task, name=name, timeout=timeout) for task, name in tasks]
|
||
|
|
return await asyncio.gather(*coros, return_exceptions=True)
|
||
|
|
|
||
|
|
async def get_results(self) -> list[StartupTaskResult]:
|
||
|
|
async with self._results_lock:
|
||
|
|
return list(self._results)
|
||
|
|
|
||
|
|
async def clear_results(self) -> None:
|
||
|
|
async with self._results_lock:
|
||
|
|
self._results.clear()
|