139 lines
5 KiB
Python
139 lines
5 KiB
Python
import logging
|
|
import json
|
|
import uuid
|
|
import hashlib
|
|
import redis
|
|
from celery import Task
|
|
|
|
from rhodecode import CONFIG
|
|
from rhodecode.lib.celerylib import RequestContextTask
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def stable_payload_key(task_name: str, args: tuple, kwargs: dict) -> str:
|
|
"""Deterministic key from (name, args, kwargs) ignoring any private kwargs."""
|
|
clean_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
|
|
payload = json.dumps([task_name, args, clean_kwargs], sort_keys=True, separators=(",", ":"))
|
|
digest = hashlib.sha256(payload.encode()).hexdigest()
|
|
return f"celery:unique:{task_name}:{digest}"
|
|
|
|
|
|
class UniqueByPayloadTask(Task):
|
|
"""Base task that deduplicates by (task name + args + kwargs) at enqueue time."""
|
|
|
|
abstract = True
|
|
|
|
_redis_url = None
|
|
_lock_ttl_seconds = None
|
|
_redis_connection = None
|
|
_deduplication_enabled = None
|
|
_redis_available = None
|
|
|
|
@classmethod
|
|
def redis_url(cls):
|
|
if cls._redis_url is None:
|
|
cls._redis_url = CONFIG["celery.deduplicate_lock_store"]
|
|
return cls._redis_url
|
|
|
|
@classmethod
|
|
def _redis(cls):
|
|
if cls._redis_connection is None:
|
|
log.debug("Connecting to Redis at %s", cls.redis_url())
|
|
cls._redis_connection = redis.from_url(cls.redis_url())
|
|
return cls._redis_connection
|
|
|
|
@property
|
|
def deduplication_enabled(self):
|
|
if self._deduplication_enabled is None:
|
|
self._deduplication_enabled = CONFIG["celery.enable_deduplication"]
|
|
return self._deduplication_enabled
|
|
|
|
@property
|
|
def lock_ttl_seconds(self):
|
|
if self._lock_ttl_seconds is None:
|
|
self._lock_ttl_seconds = CONFIG["celery.lock_ttl_seconds"]
|
|
return self._lock_ttl_seconds
|
|
|
|
def _unique_key(self, args, kwargs) -> str:
|
|
return stable_payload_key(self.name, args, kwargs)
|
|
|
|
def apply_async(
|
|
self, args=None, kwargs=None, task_id=None, producer=None, link=None, link_error=None, shadow=None, **options
|
|
):
|
|
"""
|
|
If another identical (name+payload) task is queued/running,
|
|
return its AsyncResult instead of enqueueing a new message.
|
|
"""
|
|
|
|
self._validate_config_and_redis_available()
|
|
|
|
if not self.deduplication_enabled:
|
|
log.debug("Deduplication disabled; skipping unique key check")
|
|
return super().apply_async(args, kwargs, task_id, producer, link, link_error, shadow, **options)
|
|
|
|
args = args or ()
|
|
kwargs = kwargs or {}
|
|
|
|
key = self._unique_key(args, kwargs)
|
|
r = self._redis()
|
|
|
|
# Generate a task_id up front so we can store/return it consistently.
|
|
task_id = uuid.uuid4().hex
|
|
|
|
# Atomically claim the key if absent
|
|
claimed = r.set(key, task_id, nx=True, ex=self.lock_ttl_seconds)
|
|
log.debug("Task %s unique key claim: %s", key, claimed)
|
|
if not claimed:
|
|
# Someone else already queued (or is running) the same task.
|
|
from celery.result import AsyncResult
|
|
|
|
existing = r.get(key)
|
|
log.debug("Task %s already queued or running; returning existing result", key)
|
|
return AsyncResult((existing or b"").decode() or None, app=self._get_app())
|
|
|
|
options["headers"] = options.get("headers", {})
|
|
options["headers"].update({"unique_key": key})
|
|
|
|
return super().apply_async(args, kwargs, task_id, producer, link, link_error, shadow, **options)
|
|
|
|
def after_return(self, status, retval, task_id, args, kwargs, einfo):
|
|
if not self.deduplication_enabled:
|
|
log.debug("Deduplication disabled; skipping unique key check")
|
|
return super().after_return(status, retval, task_id, args, kwargs, einfo)
|
|
|
|
key = None
|
|
try:
|
|
key = getattr(self.request, "headers", {}).get("unique_key")
|
|
log.debug("Deleting lock %s", key)
|
|
if key:
|
|
self._redis().delete(key)
|
|
except Exception as e:
|
|
# Don't crash the worker if Redis is momentarily unavailable.
|
|
if key:
|
|
log.warning("Failed to delete lock %s: %s", key, e)
|
|
|
|
return super().after_return(status, retval, task_id, args, kwargs, einfo)
|
|
|
|
def _validate_config_and_redis_available(self):
|
|
if (
|
|
CONFIG.get("celery.enable_deduplication") is None
|
|
or CONFIG.get("celery.deduplicate_lock_store") is None
|
|
or CONFIG.get("celery.lock_ttl_seconds") is None
|
|
):
|
|
log.warning("Missing config; unique task deduplication disabled")
|
|
self._deduplication_enabled = False
|
|
self._redis_available = False
|
|
|
|
if self._redis_available is None:
|
|
try:
|
|
self._redis()
|
|
self._redis_available = True
|
|
except Exception as e:
|
|
log.warning("Failed to connect to Redis: %s; unique task deduplication disabled", e)
|
|
self._deduplication_enabled = False
|
|
self._redis_available = False
|
|
|
|
|
|
class RequestContextAndUniqueByPayloadTask(RequestContextTask, UniqueByPayloadTask):
|
|
abstract = True
|