celery 容灾升级
This commit is contained in:
@@ -0,0 +1,543 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Awaitable, Callable, Mapping
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_runtime import CeleryRuntimeEvent, CeleryRuntimeState
|
||||
from app.services.celery_runtime.worker_service import current_worker_identity
|
||||
from app.services.operation_log_service import log_operation_event
|
||||
from app.services.redis_registry_service import (
|
||||
RedisExecutionLockError,
|
||||
RedisExecutionLockLost,
|
||||
RedisExecutionLockUnavailable,
|
||||
datetime_to_epoch,
|
||||
get_registry_redis,
|
||||
)
|
||||
|
||||
try:
|
||||
from redis.exceptions import RedisError
|
||||
except ImportError: # pragma: no cover
|
||||
RedisError = RuntimeError # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger("video_gen")
|
||||
|
||||
DbHeartbeat = Callable[[str], Awaitable[bool]]
|
||||
|
||||
# KEYS:
|
||||
# 1 lock, 2 active hash, 3 active zset, 4 worker-instance task set,
|
||||
# 5 runtime location hash.
|
||||
#
|
||||
# ARGV:
|
||||
# 1 token, 2 lock ttl ms, 3 runtime id, 4 payload, 5 check_at,
|
||||
# 6 worker task-set ttl seconds, 7 location payload, 8 track worker flag.
|
||||
_ACQUIRE_SCRIPT = """
|
||||
if redis.call('exists', KEYS[1]) == 1 then
|
||||
return 0
|
||||
end
|
||||
redis.call('psetex', KEYS[1], ARGV[2], ARGV[1])
|
||||
redis.call('hset', KEYS[2], ARGV[3], ARGV[4])
|
||||
redis.call('zadd', KEYS[3], ARGV[5], ARGV[3])
|
||||
if ARGV[8] == '1' then
|
||||
redis.call('sadd', KEYS[4], ARGV[3])
|
||||
redis.call('expire', KEYS[4], ARGV[6])
|
||||
end
|
||||
redis.call('hset', KEYS[5], ARGV[3], ARGV[7])
|
||||
return 1
|
||||
"""
|
||||
|
||||
_HEARTBEAT_SCRIPT = """
|
||||
if redis.call('get', KEYS[1]) ~= ARGV[1] then
|
||||
return 0
|
||||
end
|
||||
redis.call('pexpire', KEYS[1], ARGV[2])
|
||||
local merged_payload = ARGV[4]
|
||||
local current_payload = redis.call('hget', KEYS[2], ARGV[3])
|
||||
if current_payload then
|
||||
local current_ok, current_obj = pcall(cjson.decode, current_payload)
|
||||
local update_ok, update_obj = pcall(cjson.decode, ARGV[4])
|
||||
if current_ok and update_ok then
|
||||
for key, value in pairs(update_obj) do
|
||||
current_obj[key] = value
|
||||
end
|
||||
merged_payload = cjson.encode(current_obj)
|
||||
end
|
||||
end
|
||||
redis.call('hset', KEYS[2], ARGV[3], merged_payload)
|
||||
redis.call('zadd', KEYS[3], ARGV[5], ARGV[3])
|
||||
if ARGV[8] == '1' then
|
||||
redis.call('sadd', KEYS[4], ARGV[3])
|
||||
redis.call('expire', KEYS[4], ARGV[6])
|
||||
end
|
||||
redis.call('hset', KEYS[5], ARGV[3], ARGV[7])
|
||||
return 1
|
||||
"""
|
||||
|
||||
_COMPLETE_SCRIPT = """
|
||||
if redis.call('get', KEYS[1]) ~= ARGV[1] then
|
||||
return 0
|
||||
end
|
||||
redis.call('del', KEYS[1])
|
||||
redis.call('hdel', KEYS[2], ARGV[2])
|
||||
redis.call('zrem', KEYS[3], ARGV[2])
|
||||
redis.call('srem', KEYS[4], ARGV[2])
|
||||
redis.call('hdel', KEYS[5], ARGV[2])
|
||||
return 1
|
||||
"""
|
||||
|
||||
|
||||
def _now_epoch() -> int:
|
||||
return datetime_to_epoch(datetime.now(timezone.utc))
|
||||
|
||||
|
||||
def _token_digest(token: str) -> str:
|
||||
return hashlib.sha256(token.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
def build_runtime_id(domain: str, owner_type: str, owner_id: str, attempt_no: int | None = None) -> str:
|
||||
attempt = max(1, int(attempt_no or 1))
|
||||
return f"{domain}:{owner_type}:{owner_id}:attempt:{attempt}"
|
||||
|
||||
|
||||
def _worker_task_set_key(worker_instance_id: str) -> str:
|
||||
return f"{settings.CELERY_RUNTIME_WORKER_TASK_SET_PREFIX}:{worker_instance_id}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RuntimeIdentity:
|
||||
domain: str
|
||||
owner_type: str
|
||||
owner_id: str
|
||||
attempt_no: int
|
||||
task_name: str
|
||||
queue: str
|
||||
registry_item_id: str | None = None
|
||||
|
||||
@property
|
||||
def runtime_id(self) -> str:
|
||||
return self.registry_item_id or build_runtime_id(
|
||||
self.domain,
|
||||
self.owner_type,
|
||||
self.owner_id,
|
||||
self.attempt_no,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CeleryRuntimeLease:
|
||||
identity: RuntimeIdentity
|
||||
lock_key: str
|
||||
hash_key: str
|
||||
zset_key: str
|
||||
token: str
|
||||
ttl_seconds: int
|
||||
heartbeat_interval_seconds: int
|
||||
payload: dict[str, Any]
|
||||
db_heartbeat: DbHeartbeat | None = None
|
||||
db_heartbeat_grace_seconds: int = 60
|
||||
_stop_event: asyncio.Event = field(default_factory=asyncio.Event, init=False, repr=False)
|
||||
_heartbeat_task: asyncio.Task[Any] | None = field(default=None, init=False, repr=False)
|
||||
_lost_error: RedisExecutionLockError | None = field(default=None, init=False, repr=False)
|
||||
|
||||
@classmethod
|
||||
async def acquire(
|
||||
cls,
|
||||
*,
|
||||
identity: RuntimeIdentity,
|
||||
lock_key: str,
|
||||
hash_key: str,
|
||||
zset_key: str,
|
||||
token: str,
|
||||
ttl_seconds: int,
|
||||
heartbeat_interval_seconds: int,
|
||||
pipeline_stage: str | None = None,
|
||||
input_hash: str | None = None,
|
||||
business_version: int | str | None = None,
|
||||
extra_payload: Mapping[str, Any] | None = None,
|
||||
db_heartbeat: DbHeartbeat | None = None,
|
||||
) -> "CeleryRuntimeLease | None":
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_REDIS_UNAVAILABLE.value,
|
||||
event_status="failed",
|
||||
source="celery",
|
||||
task_id=identity.owner_id,
|
||||
detail={"runtime_id": identity.runtime_id, "domain": identity.domain},
|
||||
)
|
||||
raise RedisExecutionLockUnavailable(f"Redis runtime unavailable: {identity.runtime_id}")
|
||||
|
||||
worker = current_worker_identity()
|
||||
now = _now_epoch()
|
||||
ttl = max(1, int(ttl_seconds or 60))
|
||||
heartbeat_interval = max(
|
||||
1,
|
||||
min(ttl - 1 if ttl > 1 else 1, int(heartbeat_interval_seconds or 30)),
|
||||
)
|
||||
check_at = now + ttl
|
||||
task_set_ttl = max(
|
||||
ttl * 2,
|
||||
int(settings.CELERY_RUNTIME_WORKER_TASK_SET_TTL_SECONDS or 86400),
|
||||
)
|
||||
task_set_key = _worker_task_set_key(worker.worker_instance_id)
|
||||
track_worker = "1" if worker.supports_targeted_recovery else "0"
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"runtime_schema_version": int(settings.CELERY_RUNTIME_SCHEMA_VERSION or 2),
|
||||
"runtime_id": identity.runtime_id,
|
||||
"domain": identity.domain,
|
||||
"task_name": identity.task_name,
|
||||
"queue": identity.queue,
|
||||
"owner_type": identity.owner_type,
|
||||
"owner_id": identity.owner_id,
|
||||
"attempt_no": identity.attempt_no,
|
||||
"business_version": business_version,
|
||||
"input_hash": input_hash,
|
||||
"worker_name": worker.worker_name,
|
||||
"worker_instance_id": worker.worker_instance_id,
|
||||
"worker_main_pid": worker.worker_main_pid,
|
||||
"execution_pid": worker.execution_pid,
|
||||
"execution_thread_id": worker.execution_thread_id,
|
||||
"worker_started_at": worker.started_at,
|
||||
"worker_identity_quality": worker.identity_quality,
|
||||
"supports_targeted_recovery": worker.supports_targeted_recovery,
|
||||
"host": worker.host,
|
||||
"host_boot_id": worker.host_boot_id,
|
||||
"celery_task_id": worker.celery_task_id,
|
||||
"lock_key": lock_key,
|
||||
"lock_token_digest": _token_digest(token),
|
||||
"runtime_state": CeleryRuntimeState.ACTIVE.value,
|
||||
"pipeline_stage": pipeline_stage,
|
||||
"started_at": now,
|
||||
"heartbeat_at": now,
|
||||
"lease_until": check_at,
|
||||
"check_at": check_at,
|
||||
"recovery_count": 0,
|
||||
}
|
||||
if extra_payload:
|
||||
payload["extra"] = dict(extra_payload)
|
||||
|
||||
location_payload = json.dumps(
|
||||
{
|
||||
"runtime_schema_version": int(settings.CELERY_RUNTIME_SCHEMA_VERSION or 2),
|
||||
"runtime_id": identity.runtime_id,
|
||||
"domain": identity.domain,
|
||||
"hash_key": hash_key,
|
||||
"zset_key": zset_key,
|
||||
"lock_key": lock_key,
|
||||
"worker_name": worker.worker_name,
|
||||
"worker_instance_id": worker.worker_instance_id,
|
||||
"supports_targeted_recovery": worker.supports_targeted_recovery,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_ACQUIRE_START.value,
|
||||
event_status="started",
|
||||
source="celery",
|
||||
task_id=identity.owner_id,
|
||||
detail={
|
||||
"runtime_id": identity.runtime_id,
|
||||
"worker_instance_id": worker.worker_instance_id,
|
||||
"execution_pid": worker.execution_pid,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
acquired = await redis.eval(
|
||||
_ACQUIRE_SCRIPT,
|
||||
5,
|
||||
lock_key,
|
||||
hash_key,
|
||||
zset_key,
|
||||
task_set_key,
|
||||
settings.CELERY_RUNTIME_LOCATION_HASH_KEY,
|
||||
token,
|
||||
ttl * 1000,
|
||||
identity.runtime_id,
|
||||
json.dumps(payload, ensure_ascii=False, default=str),
|
||||
check_at,
|
||||
task_set_ttl,
|
||||
location_payload,
|
||||
track_worker,
|
||||
)
|
||||
except (RedisError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
logger.exception("Celery runtime acquire failed. runtime_id=%s", identity.runtime_id)
|
||||
raise RedisExecutionLockUnavailable(
|
||||
f"Redis runtime acquire failed: {identity.runtime_id}: {exc}"
|
||||
) from exc
|
||||
|
||||
if not acquired:
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_LOCK_HELD.value,
|
||||
event_status="skipped",
|
||||
source="celery",
|
||||
task_id=identity.owner_id,
|
||||
detail={"runtime_id": identity.runtime_id, "lock_key": lock_key},
|
||||
)
|
||||
return None
|
||||
|
||||
lease = cls(
|
||||
identity=identity,
|
||||
lock_key=lock_key,
|
||||
hash_key=hash_key,
|
||||
zset_key=zset_key,
|
||||
token=token,
|
||||
ttl_seconds=ttl,
|
||||
heartbeat_interval_seconds=heartbeat_interval,
|
||||
payload=payload,
|
||||
db_heartbeat=db_heartbeat,
|
||||
db_heartbeat_grace_seconds=max(60, heartbeat_interval * 2),
|
||||
)
|
||||
lease._heartbeat_task = asyncio.create_task(lease._heartbeat_loop())
|
||||
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_ACQUIRED.value,
|
||||
event_status="success",
|
||||
source="celery",
|
||||
task_id=identity.owner_id,
|
||||
detail={
|
||||
"runtime_id": identity.runtime_id,
|
||||
"domain": identity.domain,
|
||||
"attempt_no": identity.attempt_no,
|
||||
"worker_instance_id": worker.worker_instance_id,
|
||||
"worker_main_pid": worker.worker_main_pid,
|
||||
"execution_pid": worker.execution_pid,
|
||||
"lock_token_suffix": token[-8:],
|
||||
"lease_until": check_at,
|
||||
},
|
||||
)
|
||||
return lease
|
||||
|
||||
async def _heartbeat_loop(self) -> None:
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._stop_event.wait(),
|
||||
timeout=self.heartbeat_interval_seconds,
|
||||
)
|
||||
return
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
try:
|
||||
await self._heartbeat_once()
|
||||
except RedisExecutionLockError as exc:
|
||||
self._lost_error = exc
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_HEARTBEAT_LOST.value,
|
||||
event_status="failed",
|
||||
source="celery",
|
||||
task_id=self.identity.owner_id,
|
||||
detail={"runtime_id": self.identity.runtime_id, "error": str(exc)},
|
||||
error=str(exc),
|
||||
)
|
||||
return
|
||||
except Exception as exc:
|
||||
self._lost_error = RedisExecutionLockUnavailable(str(exc))
|
||||
logger.exception(
|
||||
"Celery runtime heartbeat failed. runtime_id=%s",
|
||||
self.identity.runtime_id,
|
||||
)
|
||||
return
|
||||
|
||||
async def _heartbeat_once(self) -> None:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable(
|
||||
f"Redis runtime heartbeat unavailable: {self.identity.runtime_id}"
|
||||
)
|
||||
|
||||
now = _now_epoch()
|
||||
check_at = now + self.ttl_seconds
|
||||
payload = dict(self.payload)
|
||||
payload.update(
|
||||
{
|
||||
"heartbeat_at": now,
|
||||
"lease_until": check_at,
|
||||
"check_at": check_at,
|
||||
# threads 模式下线程 ID可能随下一次执行变化，但同一 lease 生命周期固定。
|
||||
"execution_pid": self.payload.get("execution_pid"),
|
||||
"execution_thread_id": self.payload.get("execution_thread_id"),
|
||||
}
|
||||
)
|
||||
|
||||
worker_instance_id = str(payload.get("worker_instance_id") or "")
|
||||
task_set_key = _worker_task_set_key(worker_instance_id)
|
||||
task_set_ttl = max(
|
||||
self.ttl_seconds * 2,
|
||||
int(settings.CELERY_RUNTIME_WORKER_TASK_SET_TTL_SECONDS or 86400),
|
||||
)
|
||||
track_worker = "1" if bool(payload.get("supports_targeted_recovery")) else "0"
|
||||
location_payload = json.dumps(
|
||||
{
|
||||
"runtime_schema_version": int(settings.CELERY_RUNTIME_SCHEMA_VERSION or 2),
|
||||
"runtime_id": self.identity.runtime_id,
|
||||
"domain": self.identity.domain,
|
||||
"hash_key": self.hash_key,
|
||||
"zset_key": self.zset_key,
|
||||
"lock_key": self.lock_key,
|
||||
"worker_name": payload.get("worker_name"),
|
||||
"worker_instance_id": worker_instance_id,
|
||||
"supports_targeted_recovery": bool(payload.get("supports_targeted_recovery")),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
|
||||
try:
|
||||
renewed = await redis.eval(
|
||||
_HEARTBEAT_SCRIPT,
|
||||
5,
|
||||
self.lock_key,
|
||||
self.hash_key,
|
||||
self.zset_key,
|
||||
task_set_key,
|
||||
settings.CELERY_RUNTIME_LOCATION_HASH_KEY,
|
||||
self.token,
|
||||
self.ttl_seconds * 1000,
|
||||
self.identity.runtime_id,
|
||||
json.dumps(payload, ensure_ascii=False, default=str),
|
||||
check_at,
|
||||
task_set_ttl,
|
||||
location_payload,
|
||||
track_worker,
|
||||
)
|
||||
except (RedisError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
raise RedisExecutionLockUnavailable(
|
||||
f"Redis runtime heartbeat failed: {self.identity.runtime_id}: {exc}"
|
||||
) from exc
|
||||
|
||||
if not renewed:
|
||||
raise RedisExecutionLockLost(
|
||||
f"Redis runtime ownership lost: {self.identity.runtime_id}"
|
||||
)
|
||||
|
||||
self.payload = payload
|
||||
if self.db_heartbeat is not None:
|
||||
started_at = int(self.payload.get("started_at") or now)
|
||||
if now - started_at < self.db_heartbeat_grace_seconds:
|
||||
return
|
||||
owned = await self.db_heartbeat(self.token)
|
||||
if not owned:
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_DB_LEASE_LOST.value,
|
||||
event_status="failed",
|
||||
source="celery",
|
||||
task_id=self.identity.owner_id,
|
||||
detail={"runtime_id": self.identity.runtime_id},
|
||||
)
|
||||
raise RedisExecutionLockLost(
|
||||
f"Database lease ownership lost: {self.identity.runtime_id}"
|
||||
)
|
||||
|
||||
async def ensure_owned(self) -> None:
|
||||
if self._lost_error is not None:
|
||||
raise self._lost_error
|
||||
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable(
|
||||
f"Redis runtime unavailable: {self.identity.runtime_id}"
|
||||
)
|
||||
|
||||
try:
|
||||
value = await redis.get(self.lock_key)
|
||||
except (RedisError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
raise RedisExecutionLockUnavailable(
|
||||
f"Redis runtime check failed: {self.identity.runtime_id}: {exc}"
|
||||
) from exc
|
||||
|
||||
if str(value or "") != self.token:
|
||||
self._lost_error = RedisExecutionLockLost(
|
||||
f"Redis runtime ownership lost: {self.identity.runtime_id}"
|
||||
)
|
||||
raise self._lost_error
|
||||
|
||||
async def close(self) -> None:
|
||||
self._stop_event.set()
|
||||
if self._heartbeat_task is not None:
|
||||
try:
|
||||
await self._heartbeat_task
|
||||
except Exception:
|
||||
logger.debug("runtime heartbeat close failed", exc_info=True)
|
||||
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
return
|
||||
|
||||
worker_instance_id = str(self.payload.get("worker_instance_id") or "")
|
||||
task_set_key = _worker_task_set_key(worker_instance_id)
|
||||
cleaned = False
|
||||
try:
|
||||
cleaned = bool(
|
||||
await redis.eval(
|
||||
_COMPLETE_SCRIPT,
|
||||
5,
|
||||
self.lock_key,
|
||||
self.hash_key,
|
||||
self.zset_key,
|
||||
task_set_key,
|
||||
settings.CELERY_RUNTIME_LOCATION_HASH_KEY,
|
||||
self.token,
|
||||
self.identity.runtime_id,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Celery runtime cleanup failed. runtime_id=%s",
|
||||
self.identity.runtime_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RUNTIME_COMPLETED.value,
|
||||
event_status="success" if cleaned else "skipped",
|
||||
source="celery",
|
||||
task_id=self.identity.owner_id,
|
||||
detail={
|
||||
"runtime_id": self.identity.runtime_id,
|
||||
"ownership_cleanup": cleaned,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def runtime_lock_exists(lock_key: str) -> bool:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable(f"Redis unavailable while checking lock: {lock_key}")
|
||||
try:
|
||||
return bool(await redis.exists(lock_key))
|
||||
except (RedisError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
raise RedisExecutionLockUnavailable(f"Redis lock check failed: {lock_key}: {exc}") from exc
|
||||
|
||||
|
||||
async def runtime_lock_values(lock_keys: list[str]) -> dict[str, str | None]:
|
||||
"""批量读取执行锁；Redis 不可用时 fail-closed。"""
|
||||
if not lock_keys:
|
||||
return {}
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while batch checking runtime locks")
|
||||
try:
|
||||
values = await redis.mget(lock_keys)
|
||||
except (RedisError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
raise RedisExecutionLockUnavailable(f"Redis batch lock check failed: {exc}") from exc
|
||||
return {
|
||||
lock_key: (str(value) if value not in (None, "") else None)
|
||||
for lock_key, value in zip(lock_keys, values)
|
||||
}
|
||||
Reference in New Issue
Block a user