celery 容灾升级
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
from app.services.celery_runtime.runtime_service import (
|
||||
CeleryRuntimeLease,
|
||||
RuntimeIdentity,
|
||||
build_runtime_id,
|
||||
runtime_lock_exists,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CeleryRuntimeLease",
|
||||
"RuntimeIdentity",
|
||||
"build_runtime_id",
|
||||
"runtime_lock_exists",
|
||||
]
|
||||
@@ -0,0 +1,409 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_runtime import CeleryRuntimeEvent
|
||||
from app.services.operation_log_service import log_operation_event
|
||||
from app.services.redis_registry_service import (
|
||||
RedisExecutionLockUnavailable,
|
||||
datetime_to_epoch,
|
||||
get_registry_redis,
|
||||
)
|
||||
|
||||
|
||||
def _worker_instance_key(worker_instance_id: str) -> str:
|
||||
return f"{settings.CELERY_RUNTIME_WORKER_INSTANCE_PREFIX}:{worker_instance_id}"
|
||||
|
||||
|
||||
def _worker_name_instances_key(worker_name: str) -> str:
|
||||
return f"{settings.CELERY_RUNTIME_WORKER_NAME_INSTANCE_ZSET_PREFIX}:{worker_name}"
|
||||
|
||||
|
||||
def _worker_task_set_key(worker_instance_id: str) -> str:
|
||||
return f"{settings.CELERY_RUNTIME_WORKER_TASK_SET_PREFIX}:{worker_instance_id}"
|
||||
|
||||
|
||||
async def startup_barrier_exists() -> bool:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while checking startup barrier")
|
||||
return bool(await redis.exists(settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY))
|
||||
|
||||
|
||||
async def set_startup_barrier() -> None:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while setting startup barrier")
|
||||
await redis.set(
|
||||
settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY,
|
||||
str(datetime_to_epoch(datetime.now(timezone.utc))),
|
||||
ex=max(30, int(settings.CELERY_RUNTIME_STARTUP_BARRIER_TTL_SECONDS or 120)),
|
||||
)
|
||||
|
||||
|
||||
async def clear_startup_barrier() -> None:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
return
|
||||
await redis.delete(settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY)
|
||||
|
||||
|
||||
async def guard_periodic_recovery(*, check_global_lock: bool = True) -> dict[str, Any] | None:
|
||||
if await startup_barrier_exists():
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RECOVERY_BARRIER_SKIPPED.value,
|
||||
event_status="skipped",
|
||||
source="recovery",
|
||||
detail={"barrier_key": settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY},
|
||||
)
|
||||
return {"skipped": "startup_barrier"}
|
||||
|
||||
if check_global_lock:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while checking global recovery lock")
|
||||
if await redis.exists(settings.CELERY_RUNTIME_GLOBAL_RECOVERY_LOCK_KEY):
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.RECOVERY_BARRIER_SKIPPED.value,
|
||||
event_status="skipped",
|
||||
source="recovery",
|
||||
detail={"global_lock_key": settings.CELERY_RUNTIME_GLOBAL_RECOVERY_LOCK_KEY},
|
||||
)
|
||||
return {"skipped": "global_recovery_lock"}
|
||||
return None
|
||||
|
||||
|
||||
async def garbage_collect_registry_pair(*, hash_key: str, zset_key: str, limit: int = 500) -> dict[str, int]:
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable during registry GC")
|
||||
|
||||
members = await redis.zrange(zset_key, 0, max(0, int(limit) - 1))
|
||||
if not members:
|
||||
return {"checked": 0, "removed": 0}
|
||||
|
||||
values = await redis.hmget(hash_key, members)
|
||||
stale = [str(member) for member, value in zip(members, values) if not value]
|
||||
if stale:
|
||||
await redis.zrem(zset_key, *stale)
|
||||
return {"checked": len(members), "removed": len(stale)}
|
||||
|
||||
|
||||
async def mark_stale_worker_instance_candidates(
|
||||
*,
|
||||
worker_name: str,
|
||||
current_worker_instance_id: str,
|
||||
supports_targeted_recovery: bool,
|
||||
emit_duplicate_log: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""只扫描同一逻辑 Worker 名称下已经失活的旧主实例。
|
||||
|
||||
同名存在其他活跃实例时保守跳过，避免滚动发布或错误重名导致互相抢占。
|
||||
旧实例中的任务只有在业务 String 锁已经不存在时才提前标记为恢复候选；
|
||||
真正接管仍由各业务恢复服务执行 DB lease/claim/attempt CAS。
|
||||
"""
|
||||
if not supports_targeted_recovery:
|
||||
return {
|
||||
"skipped": "identity_fallback",
|
||||
"checked_instances": 0,
|
||||
"stale_instances": 0,
|
||||
"marked": 0,
|
||||
}
|
||||
|
||||
normalized_name = str(worker_name or "").strip()
|
||||
normalized_current = str(current_worker_instance_id or "").strip()
|
||||
if not normalized_name or not normalized_current:
|
||||
return {
|
||||
"skipped": "invalid_identity",
|
||||
"checked_instances": 0,
|
||||
"stale_instances": 0,
|
||||
"marked": 0,
|
||||
}
|
||||
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while discovering stale worker instances")
|
||||
|
||||
index_key = _worker_name_instances_key(normalized_name)
|
||||
raw_instances = await redis.zrange(index_key, 0, -1, withscores=True)
|
||||
instances = [
|
||||
(str(member), int(float(score)))
|
||||
for member, score in raw_instances
|
||||
if str(member) != normalized_current
|
||||
]
|
||||
if not instances:
|
||||
return {
|
||||
"checked_instances": 0,
|
||||
"active_instances": 0,
|
||||
"stale_instances": 0,
|
||||
"marked": 0,
|
||||
"live": 0,
|
||||
"invalid": 0,
|
||||
}
|
||||
|
||||
keys = [_worker_instance_key(instance_id) for instance_id, _ in instances]
|
||||
values = await redis.mget(keys)
|
||||
active_instances = [
|
||||
instance_id
|
||||
for (instance_id, _), value in zip(instances, values)
|
||||
if value
|
||||
]
|
||||
if active_instances:
|
||||
if emit_duplicate_log:
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_DUPLICATE_NAME.value,
|
||||
event_status="warning",
|
||||
source="worker_ready",
|
||||
detail={
|
||||
"worker_name": normalized_name,
|
||||
"current_worker_instance_id": normalized_current,
|
||||
"other_active_worker_instance_ids": active_instances,
|
||||
},
|
||||
)
|
||||
return {
|
||||
"skipped": "duplicate_active_worker_name",
|
||||
"checked_instances": len(instances),
|
||||
"active_instances": len(active_instances),
|
||||
"stale_instances": 0,
|
||||
"marked": 0,
|
||||
}
|
||||
|
||||
# 实例 key 由主进程 heartbeat 独立维护。key 已不存在即可判定旧主实例
|
||||
# 不再活跃；任务是否可以恢复仍必须继续检查业务 String 锁和 DB fencing。
|
||||
stale_instances = [
|
||||
instance_id
|
||||
for (instance_id, _heartbeat_at), value in zip(instances, values)
|
||||
if not value
|
||||
]
|
||||
|
||||
totals: dict[str, Any] = {
|
||||
"checked_instances": len(instances),
|
||||
"active_instances": 0,
|
||||
"stale_instances": len(stale_instances),
|
||||
"checked": 0,
|
||||
"marked": 0,
|
||||
"live": 0,
|
||||
"invalid": 0,
|
||||
}
|
||||
|
||||
for stale_instance_id in stale_instances:
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_STALE_INSTANCE_FOUND.value,
|
||||
event_status="success",
|
||||
source="worker_ready",
|
||||
detail={
|
||||
"worker_name": normalized_name,
|
||||
"current_worker_instance_id": normalized_current,
|
||||
"old_worker_instance_id": stale_instance_id,
|
||||
},
|
||||
)
|
||||
result = await mark_worker_instance_runtime_candidates(
|
||||
old_worker_instance_id=stale_instance_id,
|
||||
worker_name=normalized_name,
|
||||
)
|
||||
for key in ("checked", "marked", "live", "invalid"):
|
||||
totals[key] += int(result.get(key, 0) or 0)
|
||||
|
||||
return totals
|
||||
|
||||
|
||||
async def mark_worker_instance_runtime_candidates(
|
||||
*,
|
||||
old_worker_instance_id: str,
|
||||
worker_name: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""按旧 Worker 主实例 ID 精准标记失锁任务。"""
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while marking worker recovery candidates")
|
||||
|
||||
normalized_instance = str(old_worker_instance_id or "").strip()
|
||||
if not normalized_instance:
|
||||
return {"checked": 0, "marked": 0, "live": 0, "invalid": 0}
|
||||
|
||||
task_set_key = _worker_task_set_key(normalized_instance)
|
||||
runtime_ids = [str(value) for value in await redis.smembers(task_set_key)]
|
||||
if not runtime_ids:
|
||||
return {"checked": 0, "marked": 0, "live": 0, "invalid": 0}
|
||||
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_INSTANCE_RECOVERY_START.value,
|
||||
event_status="started",
|
||||
source="worker_ready",
|
||||
detail={
|
||||
"worker_name": worker_name,
|
||||
"old_worker_instance_id": normalized_instance,
|
||||
"candidate_count": len(runtime_ids),
|
||||
},
|
||||
)
|
||||
|
||||
raw_locations = await redis.hmget(settings.CELERY_RUNTIME_LOCATION_HASH_KEY, runtime_ids)
|
||||
locations: dict[str, dict[str, Any]] = {}
|
||||
lock_keys: list[str] = []
|
||||
invalid_ids: list[str] = []
|
||||
|
||||
for runtime_id, raw in zip(runtime_ids, raw_locations):
|
||||
try:
|
||||
payload = json.loads(raw) if raw else None
|
||||
except Exception:
|
||||
payload = None
|
||||
|
||||
if (
|
||||
not isinstance(payload, dict)
|
||||
or not payload.get("lock_key")
|
||||
or not payload.get("hash_key")
|
||||
or not payload.get("zset_key")
|
||||
or str(payload.get("worker_instance_id") or "") != normalized_instance
|
||||
):
|
||||
invalid_ids.append(runtime_id)
|
||||
continue
|
||||
|
||||
locations[runtime_id] = payload
|
||||
lock_keys.append(str(payload["lock_key"]))
|
||||
|
||||
lock_values = await redis.mget(lock_keys) if lock_keys else []
|
||||
live_by_key = {key: bool(value) for key, value in zip(lock_keys, lock_values)}
|
||||
|
||||
now_epoch = datetime_to_epoch(datetime.now(timezone.utc))
|
||||
marked = 0
|
||||
live = 0
|
||||
stale_by_hash: dict[str, list[str]] = {}
|
||||
|
||||
for runtime_id, location in locations.items():
|
||||
lock_key = str(location["lock_key"])
|
||||
if live_by_key.get(lock_key):
|
||||
live += 1
|
||||
continue
|
||||
stale_by_hash.setdefault(str(location["hash_key"]), []).append(runtime_id)
|
||||
|
||||
payload_by_runtime: dict[str, dict[str, Any]] = {}
|
||||
for hash_key, hash_runtime_ids in stale_by_hash.items():
|
||||
raw_payloads = await redis.hmget(hash_key, hash_runtime_ids)
|
||||
for runtime_id, raw_payload in zip(hash_runtime_ids, raw_payloads):
|
||||
try:
|
||||
payload = json.loads(raw_payload) if raw_payload else {}
|
||||
except Exception:
|
||||
payload = {}
|
||||
payload_by_runtime[runtime_id] = payload if isinstance(payload, dict) else {}
|
||||
|
||||
pipe = redis.pipeline(transaction=False)
|
||||
for runtime_id, location in locations.items():
|
||||
lock_key = str(location["lock_key"])
|
||||
if live_by_key.get(lock_key):
|
||||
continue
|
||||
|
||||
hash_key = str(location["hash_key"])
|
||||
zset_key = str(location["zset_key"])
|
||||
payload = payload_by_runtime.get(runtime_id, {})
|
||||
payload.update(
|
||||
{
|
||||
"runtime_state": "recovery_candidate",
|
||||
"check_at": now_epoch,
|
||||
"reason": "worker_instance_inactive_lock_missing",
|
||||
"old_worker_instance_id": normalized_instance,
|
||||
}
|
||||
)
|
||||
pipe.hset(hash_key, runtime_id, json.dumps(payload, ensure_ascii=False, default=str))
|
||||
pipe.zadd(zset_key, {runtime_id: now_epoch})
|
||||
marked += 1
|
||||
|
||||
if invalid_ids:
|
||||
pipe.srem(task_set_key, *invalid_ids)
|
||||
|
||||
await pipe.execute()
|
||||
|
||||
result = {
|
||||
"checked": len(runtime_ids),
|
||||
"marked": marked,
|
||||
"live": live,
|
||||
"invalid": len(invalid_ids),
|
||||
}
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_INSTANCE_RECOVERY_DONE.value,
|
||||
event_status="success",
|
||||
source="worker_ready",
|
||||
detail={
|
||||
"worker_name": worker_name,
|
||||
"old_worker_instance_id": normalized_instance,
|
||||
**result,
|
||||
},
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
async def garbage_collect_worker_registry(*, limit: int = 500) -> dict[str, int]:
|
||||
"""清理 V2 Worker 名称索引和旧实例 Task Set 中的失配成员。"""
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable during worker registry GC")
|
||||
|
||||
checked_names = 0
|
||||
checked_instances = 0
|
||||
removed_instances = 0
|
||||
removed_task_members = 0
|
||||
now = int(time.time())
|
||||
grace = max(
|
||||
int(settings.CELERY_RUNTIME_WORKER_HEARTBEAT_TTL_SECONDS or 90),
|
||||
int(settings.CELERY_RUNTIME_WORKER_STALE_GRACE_SECONDS or 120),
|
||||
)
|
||||
pattern = f"{settings.CELERY_RUNTIME_WORKER_NAME_INSTANCE_ZSET_PREFIX}:*"
|
||||
|
||||
async for raw_index_key in redis.scan_iter(match=pattern, count=max(10, min(limit, 500))):
|
||||
if checked_names >= limit:
|
||||
break
|
||||
checked_names += 1
|
||||
index_key = str(raw_index_key)
|
||||
raw_instances = await redis.zrange(index_key, 0, -1, withscores=True)
|
||||
|
||||
for raw_instance_id, raw_score in raw_instances:
|
||||
if checked_instances >= limit:
|
||||
break
|
||||
checked_instances += 1
|
||||
instance_id = str(raw_instance_id)
|
||||
score = int(float(raw_score))
|
||||
if now - score < grace:
|
||||
continue
|
||||
if await redis.exists(_worker_instance_key(instance_id)):
|
||||
continue
|
||||
|
||||
task_set_key = _worker_task_set_key(instance_id)
|
||||
runtime_ids = [str(value) for value in await redis.smembers(task_set_key)]
|
||||
if runtime_ids:
|
||||
raw_locations = await redis.hmget(settings.CELERY_RUNTIME_LOCATION_HASH_KEY, runtime_ids)
|
||||
stale_members: list[str] = []
|
||||
for runtime_id, raw_location in zip(runtime_ids, raw_locations):
|
||||
try:
|
||||
location = json.loads(raw_location) if raw_location else None
|
||||
except Exception:
|
||||
location = None
|
||||
if (
|
||||
not isinstance(location, dict)
|
||||
or str(location.get("worker_instance_id") or "") != instance_id
|
||||
):
|
||||
stale_members.append(runtime_id)
|
||||
if stale_members:
|
||||
removed_task_members += int(await redis.srem(task_set_key, *stale_members) or 0)
|
||||
|
||||
if int(await redis.scard(task_set_key) or 0) == 0:
|
||||
await redis.delete(task_set_key)
|
||||
removed_instances += int(await redis.zrem(index_key, instance_id) or 0)
|
||||
|
||||
if int(await redis.zcard(index_key) or 0) == 0:
|
||||
await redis.delete(index_key)
|
||||
|
||||
return {
|
||||
"checked_names": checked_names,
|
||||
"checked_instances": checked_instances,
|
||||
"removed_instances": removed_instances,
|
||||
"removed_task_members": removed_task_members,
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -0,0 +1,477 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_runtime import CeleryRuntimeEvent, WorkerIdentityQuality
|
||||
from app.services.operation_log_service import log_operation_event
|
||||
from app.services.redis_registry_service import (
|
||||
RedisExecutionLockUnavailable,
|
||||
get_registry_redis,
|
||||
)
|
||||
|
||||
try:
|
||||
from celery import current_task
|
||||
except Exception: # pragma: no cover
|
||||
current_task = None # type: ignore[assignment]
|
||||
|
||||
|
||||
_PROCESS_STARTED_AT = int(time.time())
|
||||
_PROCESS_INSTANCE_ID = uuid.uuid4().hex
|
||||
_PROCESS_IDENTITY_LOCK = threading.RLock()
|
||||
_WORKER_REGISTRATION_LOCK = threading.RLock()
|
||||
_HEARTBEAT_THROTTLE_LOCK = threading.RLock()
|
||||
|
||||
_ENV_INSTANCE_TOKEN = "CELERY_WORKER_INSTANCE_TOKEN"
|
||||
_ENV_MAIN_PID = "CELERY_WORKER_MAIN_PID"
|
||||
_ENV_STARTED_AT = "CELERY_WORKER_STARTED_AT"
|
||||
_ENV_HOST_BOOT_ID = "CELERY_HOST_BOOT_ID"
|
||||
_ENV_WORKER_NAME = "CELERY_WORKER_NODE_NAME"
|
||||
_ENV_BEFORE_POOL = "CELERY_WORKER_IDENTITY_BEFORE_POOL"
|
||||
|
||||
_registered_worker: "WorkerRegistration | None" = None
|
||||
_last_heartbeat_attempt_monotonic = 0.0
|
||||
_last_stale_scan_attempt_monotonic = 0.0
|
||||
_heartbeat_failure_count = 0
|
||||
|
||||
|
||||
def _parse_optional_int(value: Any) -> int | None:
|
||||
raw = str(value or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return parsed if parsed > 0 else None
|
||||
|
||||
|
||||
def _parse_optional_float(value: Any) -> float | None:
|
||||
raw = str(value or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
return float(raw)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _boot_id() -> str:
|
||||
"""返回同一次操作系统启动期间稳定的主机启动标识。
|
||||
|
||||
Linux 直接读取内核 boot_id；Windows 使用 GetTickCount64 估算启动时间并
|
||||
生成 UUID5。无法获取时返回 unknown-boot。该字段仅用于辅助诊断，不参与
|
||||
任务锁或数据库 fencing 的最终正确性判断。
|
||||
"""
|
||||
linux_path = Path("/proc/sys/kernel/random/boot_id")
|
||||
try:
|
||||
value = linux_path.read_text(encoding="utf-8").strip()
|
||||
if value:
|
||||
return value
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if os.name == "nt":
|
||||
try:
|
||||
uptime_ms = int(ctypes.windll.kernel32.GetTickCount64()) # type: ignore[attr-defined]
|
||||
# 以 10 秒为粒度消除多个进程独立计算时的亚秒抖动。
|
||||
boot_epoch_bucket = int((time.time() - uptime_ms / 1000.0) // 10 * 10)
|
||||
source = f"windows-boot:{socket.gethostname()}:{boot_epoch_bucket}"
|
||||
return str(uuid.uuid5(uuid.NAMESPACE_OID, source))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return "unknown-boot"
|
||||
|
||||
|
||||
def initialize_worker_main_identity(
|
||||
worker_name: str | None = None,
|
||||
*,
|
||||
before_pool: bool,
|
||||
) -> None:
|
||||
"""在 Celery Worker 主进程中初始化一次实例 token。
|
||||
|
||||
before_pool=True 必须由 celeryd_init/worker_init 调用，以保证 Linux prefork
|
||||
子进程通过 fork 继承相同 token。worker_ready 只允许做 late fallback，且
|
||||
late fallback 会关闭实例级精准恢复，避免主进程与已创建子进程身份不一致。
|
||||
"""
|
||||
normalized_name = str(worker_name or "").strip()
|
||||
current_pid = os.getpid()
|
||||
|
||||
should_log = False
|
||||
with _PROCESS_IDENTITY_LOCK:
|
||||
existing_pid = _parse_optional_int(os.getenv(_ENV_MAIN_PID))
|
||||
existing_token = str(os.getenv(_ENV_INSTANCE_TOKEN, "") or "").strip()
|
||||
existing_before_pool = str(os.getenv(_ENV_BEFORE_POOL, "") or "").strip() == "1"
|
||||
|
||||
if not existing_token or existing_pid != current_pid:
|
||||
os.environ[_ENV_INSTANCE_TOKEN] = uuid.uuid4().hex
|
||||
os.environ[_ENV_MAIN_PID] = str(current_pid)
|
||||
os.environ[_ENV_STARTED_AT] = str(int(time.time()))
|
||||
os.environ[_ENV_HOST_BOOT_ID] = _boot_id()
|
||||
os.environ[_ENV_BEFORE_POOL] = "1" if before_pool else "0"
|
||||
should_log = True
|
||||
elif before_pool and not existing_before_pool:
|
||||
# celeryd_init 与 worker_init 都可能触发，幂等提升为 before-pool。
|
||||
os.environ[_ENV_BEFORE_POOL] = "1"
|
||||
should_log = True
|
||||
|
||||
if normalized_name:
|
||||
os.environ[_ENV_WORKER_NAME] = normalized_name
|
||||
|
||||
_process_identity_base.cache_clear()
|
||||
|
||||
if should_log:
|
||||
identity = current_worker_identity(worker_name_override=normalized_name or None)
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=(
|
||||
CeleryRuntimeEvent.WORKER_IDENTITY_INITIALIZED.value
|
||||
if identity.supports_targeted_recovery
|
||||
else CeleryRuntimeEvent.WORKER_IDENTITY_FALLBACK.value
|
||||
),
|
||||
event_status="success" if identity.supports_targeted_recovery else "warning",
|
||||
source="worker_init",
|
||||
detail={
|
||||
"worker_name": identity.worker_name,
|
||||
"worker_instance_id": identity.worker_instance_id,
|
||||
"worker_main_pid": identity.worker_main_pid,
|
||||
"execution_pid": identity.execution_pid,
|
||||
"host_boot_id": identity.host_boot_id,
|
||||
"identity_quality": identity.identity_quality,
|
||||
"supports_targeted_recovery": identity.supports_targeted_recovery,
|
||||
"before_pool": before_pool,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def reset_process_identity_cache() -> None:
|
||||
"""prefork 子进程启动后清理 fork 前的 Python 对象缓存。
|
||||
|
||||
环境变量中的主 Worker token 保留；每个子进程重新构建自己的执行 PID 和
|
||||
线程 ID，不会生成新的 Worker 实例 token。
|
||||
"""
|
||||
_process_identity_base.cache_clear()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProcessIdentityBase:
|
||||
worker_instance_token: str
|
||||
worker_main_pid: int | None
|
||||
worker_started_at: int
|
||||
host_boot_id: str
|
||||
identity_quality: str
|
||||
supports_targeted_recovery: bool
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _process_identity_base() -> _ProcessIdentityBase:
|
||||
token = str(os.getenv(_ENV_INSTANCE_TOKEN, "") or "").strip()
|
||||
before_pool = str(os.getenv(_ENV_BEFORE_POOL, "") or "").strip() == "1"
|
||||
worker_main_pid = _parse_optional_int(os.getenv(_ENV_MAIN_PID))
|
||||
started_at = _parse_optional_float(os.getenv(_ENV_STARTED_AT))
|
||||
host_boot_id = str(os.getenv(_ENV_HOST_BOOT_ID, "") or "").strip() or _boot_id()
|
||||
|
||||
supports_targeted_recovery = bool(token and before_pool)
|
||||
if not token:
|
||||
token = _PROCESS_INSTANCE_ID
|
||||
|
||||
execution_pid = os.getpid()
|
||||
if not supports_targeted_recovery:
|
||||
quality = WorkerIdentityQuality.FALLBACK.value
|
||||
elif worker_main_pid is None or host_boot_id == "unknown-boot":
|
||||
quality = WorkerIdentityQuality.INSTANCE_TOKEN_ONLY.value
|
||||
else:
|
||||
quality = WorkerIdentityQuality.FULL.value
|
||||
|
||||
return _ProcessIdentityBase(
|
||||
worker_instance_token=token,
|
||||
worker_main_pid=worker_main_pid,
|
||||
worker_started_at=int(started_at or _PROCESS_STARTED_AT),
|
||||
host_boot_id=host_boot_id,
|
||||
identity_quality=quality,
|
||||
supports_targeted_recovery=supports_targeted_recovery,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkerIdentity:
|
||||
worker_name: str
|
||||
worker_instance_id: str
|
||||
worker_instance_token: str
|
||||
host: str
|
||||
host_boot_id: str
|
||||
worker_main_pid: int | None
|
||||
execution_pid: int
|
||||
execution_thread_id: int
|
||||
started_at: int
|
||||
identity_quality: str
|
||||
supports_targeted_recovery: bool
|
||||
celery_task_id: str | None = None
|
||||
|
||||
|
||||
def current_worker_identity(*, worker_name_override: str | None = None) -> WorkerIdentity:
|
||||
hostname = socket.gethostname()
|
||||
celery_task_id: str | None = None
|
||||
request_worker_name = ""
|
||||
|
||||
try:
|
||||
request = getattr(current_task, "request", None)
|
||||
request_worker_name = str(getattr(request, "hostname", "") or "").strip()
|
||||
celery_task_id = str(getattr(request, "id", "") or "").strip() or None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
worker_name = (
|
||||
str(worker_name_override or "").strip()
|
||||
or request_worker_name
|
||||
or str(os.getenv(_ENV_WORKER_NAME, "") or "").strip()
|
||||
or hostname
|
||||
)
|
||||
base = _process_identity_base()
|
||||
instance_id = f"{worker_name}:{base.host_boot_id}:{base.worker_instance_token}"
|
||||
|
||||
return WorkerIdentity(
|
||||
worker_name=worker_name,
|
||||
worker_instance_id=instance_id,
|
||||
worker_instance_token=base.worker_instance_token,
|
||||
host=hostname,
|
||||
host_boot_id=base.host_boot_id,
|
||||
worker_main_pid=base.worker_main_pid,
|
||||
execution_pid=os.getpid(),
|
||||
execution_thread_id=threading.get_ident(),
|
||||
started_at=base.worker_started_at,
|
||||
identity_quality=base.identity_quality,
|
||||
supports_targeted_recovery=base.supports_targeted_recovery,
|
||||
celery_task_id=celery_task_id,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkerRegistration:
|
||||
identity: WorkerIdentity
|
||||
queues: tuple[str, ...] = field(default_factory=tuple)
|
||||
pool_type: str | None = None
|
||||
configured_concurrency: int | None = None
|
||||
|
||||
def payload(self, *, heartbeat_at: int) -> dict[str, Any]:
|
||||
return {
|
||||
"runtime_schema_version": int(settings.CELERY_RUNTIME_SCHEMA_VERSION or 2),
|
||||
"worker_name": self.identity.worker_name,
|
||||
"worker_instance_id": self.identity.worker_instance_id,
|
||||
"host": self.identity.host,
|
||||
"host_boot_id": self.identity.host_boot_id,
|
||||
"worker_main_pid": self.identity.worker_main_pid,
|
||||
"started_at": self.identity.started_at,
|
||||
"identity_quality": self.identity.identity_quality,
|
||||
"supports_targeted_recovery": self.identity.supports_targeted_recovery,
|
||||
"queues": list(self.queues),
|
||||
"pool_type": self.pool_type,
|
||||
"configured_concurrency": self.configured_concurrency,
|
||||
"heartbeat_at": heartbeat_at,
|
||||
}
|
||||
|
||||
|
||||
def _worker_instance_key(worker_instance_id: str) -> str:
|
||||
return f"{settings.CELERY_RUNTIME_WORKER_INSTANCE_PREFIX}:{worker_instance_id}"
|
||||
|
||||
|
||||
def _worker_name_instances_key(worker_name: str) -> str:
|
||||
return f"{settings.CELERY_RUNTIME_WORKER_NAME_INSTANCE_ZSET_PREFIX}:{worker_name}"
|
||||
|
||||
|
||||
def _normalize_queues(values: Any) -> tuple[str, ...]:
|
||||
if values is None:
|
||||
return ()
|
||||
if isinstance(values, str):
|
||||
values = [values]
|
||||
result: list[str] = []
|
||||
try:
|
||||
iterator = iter(values)
|
||||
except TypeError:
|
||||
return ()
|
||||
for value in iterator:
|
||||
name = str(getattr(value, "name", value) or "").strip()
|
||||
if name and name not in result:
|
||||
result.append(name)
|
||||
return tuple(result)
|
||||
|
||||
|
||||
async def register_worker_instance(
|
||||
*,
|
||||
worker_name: str,
|
||||
queues: Any = None,
|
||||
pool_type: str | None = None,
|
||||
configured_concurrency: int | None = None,
|
||||
) -> WorkerIdentity:
|
||||
"""注册 Worker 主实例；Redis 不可用时不阻塞 Worker 启动。"""
|
||||
global _registered_worker, _last_heartbeat_attempt_monotonic
|
||||
|
||||
normalized_name = str(worker_name or "").strip() or socket.gethostname()
|
||||
os.environ[_ENV_WORKER_NAME] = normalized_name
|
||||
identity = current_worker_identity(worker_name_override=normalized_name)
|
||||
registration = WorkerRegistration(
|
||||
identity=identity,
|
||||
queues=_normalize_queues(queues),
|
||||
pool_type=str(pool_type or "").strip() or None,
|
||||
configured_concurrency=_parse_optional_int(configured_concurrency),
|
||||
)
|
||||
|
||||
with _WORKER_REGISTRATION_LOCK:
|
||||
_registered_worker = registration
|
||||
|
||||
now = int(time.time())
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
raise RedisExecutionLockUnavailable("Redis unavailable while registering Celery worker instance")
|
||||
|
||||
ttl = max(30, int(settings.CELERY_RUNTIME_WORKER_HEARTBEAT_TTL_SECONDS or 90))
|
||||
index_ttl = max(ttl * 2, int(settings.CELERY_RUNTIME_WORKER_TASK_SET_TTL_SECONDS or 86400))
|
||||
payload = json.dumps(registration.payload(heartbeat_at=now), ensure_ascii=False, default=str)
|
||||
|
||||
pipe = redis.pipeline(transaction=False)
|
||||
pipe.set(_worker_instance_key(identity.worker_instance_id), payload, ex=ttl)
|
||||
pipe.zadd(_worker_name_instances_key(identity.worker_name), {identity.worker_instance_id: now})
|
||||
pipe.expire(_worker_name_instances_key(identity.worker_name), index_ttl)
|
||||
await pipe.execute()
|
||||
|
||||
with _HEARTBEAT_THROTTLE_LOCK:
|
||||
_last_heartbeat_attempt_monotonic = time.monotonic()
|
||||
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_REGISTERED.value,
|
||||
event_status="success",
|
||||
source="worker_ready",
|
||||
detail=registration.payload(heartbeat_at=now),
|
||||
)
|
||||
return identity
|
||||
|
||||
|
||||
def claim_worker_heartbeat_slot() -> bool:
|
||||
"""对 Celery heartbeat_sent 信号做本进程节流。"""
|
||||
global _last_heartbeat_attempt_monotonic
|
||||
|
||||
interval = max(5, int(settings.CELERY_RUNTIME_WORKER_HEARTBEAT_INTERVAL_SECONDS or 30))
|
||||
now = time.monotonic()
|
||||
with _HEARTBEAT_THROTTLE_LOCK:
|
||||
if now - _last_heartbeat_attempt_monotonic < interval:
|
||||
return False
|
||||
_last_heartbeat_attempt_monotonic = now
|
||||
return True
|
||||
|
||||
|
||||
def claim_worker_stale_scan_slot() -> bool:
|
||||
"""限制同一 Worker 主进程的旧实例扫描频率。"""
|
||||
global _last_stale_scan_attempt_monotonic
|
||||
|
||||
interval = max(30, int(settings.CELERY_RUNTIME_WORKER_STALE_SCAN_INTERVAL_SECONDS or 120))
|
||||
now = time.monotonic()
|
||||
with _HEARTBEAT_THROTTLE_LOCK:
|
||||
if now - _last_stale_scan_attempt_monotonic < interval:
|
||||
return False
|
||||
_last_stale_scan_attempt_monotonic = now
|
||||
return True
|
||||
|
||||
|
||||
async def heartbeat_current_worker_instance() -> bool:
|
||||
"""刷新 Worker 主实例 TTL；不刷新任何业务任务锁。"""
|
||||
global _heartbeat_failure_count
|
||||
|
||||
with _WORKER_REGISTRATION_LOCK:
|
||||
registration = _registered_worker
|
||||
if registration is None:
|
||||
return False
|
||||
|
||||
redis = await get_registry_redis()
|
||||
if redis is None:
|
||||
_heartbeat_failure_count += 1
|
||||
_log_worker_heartbeat_failure_if_needed(registration, "redis_unavailable")
|
||||
return False
|
||||
|
||||
now = int(time.time())
|
||||
ttl = max(30, int(settings.CELERY_RUNTIME_WORKER_HEARTBEAT_TTL_SECONDS or 90))
|
||||
index_ttl = max(ttl * 2, int(settings.CELERY_RUNTIME_WORKER_TASK_SET_TTL_SECONDS or 86400))
|
||||
payload = json.dumps(registration.payload(heartbeat_at=now), ensure_ascii=False, default=str)
|
||||
|
||||
try:
|
||||
pipe = redis.pipeline(transaction=False)
|
||||
pipe.set(_worker_instance_key(registration.identity.worker_instance_id), payload, ex=ttl)
|
||||
pipe.zadd(
|
||||
_worker_name_instances_key(registration.identity.worker_name),
|
||||
{registration.identity.worker_instance_id: now},
|
||||
)
|
||||
pipe.expire(_worker_name_instances_key(registration.identity.worker_name), index_ttl)
|
||||
await pipe.execute()
|
||||
except Exception as exc:
|
||||
_heartbeat_failure_count += 1
|
||||
_log_worker_heartbeat_failure_if_needed(registration, str(exc))
|
||||
return False
|
||||
|
||||
_heartbeat_failure_count = 0
|
||||
return True
|
||||
|
||||
|
||||
def _log_worker_heartbeat_failure_if_needed(registration: WorkerRegistration, error: str) -> None:
|
||||
threshold = max(1, int(settings.CELERY_RUNTIME_WORKER_HEARTBEAT_FAILURE_LOG_THRESHOLD or 3))
|
||||
if _heartbeat_failure_count != threshold and _heartbeat_failure_count % (threshold * 5) != 0:
|
||||
return
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_HEARTBEAT_LOST.value,
|
||||
event_status="failed",
|
||||
source="worker_heartbeat",
|
||||
detail={
|
||||
"worker_name": registration.identity.worker_name,
|
||||
"worker_instance_id": registration.identity.worker_instance_id,
|
||||
"failure_count": _heartbeat_failure_count,
|
||||
"error": error,
|
||||
},
|
||||
error=error,
|
||||
)
|
||||
|
||||
|
||||
async def unregister_current_worker_instance() -> None:
|
||||
"""优雅退出时删除活跃实例 key，保留名称索引供旧任务精准恢复。"""
|
||||
global _registered_worker
|
||||
|
||||
with _WORKER_REGISTRATION_LOCK:
|
||||
registration = _registered_worker
|
||||
_registered_worker = None
|
||||
if registration is None:
|
||||
return
|
||||
|
||||
redis = await get_registry_redis()
|
||||
if redis is not None:
|
||||
try:
|
||||
await redis.delete(_worker_instance_key(registration.identity.worker_instance_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
log_operation_event(
|
||||
domain="celery_runtime",
|
||||
event_type=CeleryRuntimeEvent.WORKER_SHUTDOWN.value,
|
||||
event_status="success",
|
||||
source="worker_shutdown",
|
||||
detail={
|
||||
"worker_name": registration.identity.worker_name,
|
||||
"worker_instance_id": registration.identity.worker_instance_id,
|
||||
"worker_main_pid": registration.identity.worker_main_pid,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def registered_worker_identity() -> WorkerIdentity | None:
|
||||
with _WORKER_REGISTRATION_LOCK:
|
||||
return _registered_worker.identity if _registered_worker is not None else None
|
||||
Reference in New Issue
Block a user