diff --git a/configs/development.ini b/configs/development.ini index a832a132..f9bb25cb 100644 --- a/configs/development.ini +++ b/configs/development.ini @@ -419,6 +419,15 @@ archive_cache.filesystem.retry_attempts = 10 use_celery = true +; Enable message deduplication — When enabled, at most one message with the same task name + payload may be present in the queue at a time. (Default: false) +celery.enable_deduplication = false + +; Redis lock store (for deduplication) — Redis connection used to hold deduplication locks. Must be different from celery.broker_url. (Default: redis/9) +celery.deduplicate_lock_store = redis://redis:6379/9 + +;Lock TTL (seconds) — Time-to-live for a deduplication lock. After expiry, another message with the same name + payload may be enqueued. (Default: 3600 seconds or 1 hour). +celery.lock_ttl_seconds = 3600 + ; path to store schedule database #celerybeat-schedule.path = diff --git a/configs/production.ini b/configs/production.ini index 7fd8afec..5beb696c 100644 --- a/configs/production.ini +++ b/configs/production.ini @@ -381,6 +381,15 @@ archive_cache.filesystem.retry_attempts = 10 use_celery = true +; Enable message deduplication — When enabled, at most one message with the same task name + payload may be present in the queue at a time. (Default: false) +celery.enable_deduplication = false + +; Redis lock store (for deduplication) — Redis connection used to hold deduplication locks. Must be different from celery.broker_url. (Default: redis/9) +celery.deduplicate_lock_store = redis://redis:6379/9 + +;Lock TTL (seconds) — Time-to-live for a deduplication lock. After expiry, another message with the same name + payload may be enqueued. (Default: 3600 seconds or 1 hour). +celery.lock_ttl_seconds = 3600 + ; path to store schedule database #celerybeat-schedule.path = diff --git a/rhodecode/config/config_maker.py b/rhodecode/config/config_maker.py index a98a6638..5f710d83 100644 --- a/rhodecode/config/config_maker.py +++ b/rhodecode/config/config_maker.py @@ -170,6 +170,9 @@ def sanitize_settings_and_apply_defaults(global_config, settings): # celery broker_url = settings_maker.make_setting("celery.broker_url", "redis://redis:6379/8", default_when_empty=True) settings_maker.make_setting("celery.result_backend", broker_url) + settings_maker.make_setting("celery.enable_deduplication", False, parser="bool") + settings_maker.make_setting("celery.deduplicate_lock_store", "redis://redis:6379/9", parser="string") + settings_maker.make_setting("celery.lock_ttl_seconds", 3600, parser="int") settings_maker.make_setting("exception_tracker.send_email", False, parser="bool") settings_maker.make_setting("exception_tracker.email_prefix", "[RHODECODE ERROR]", default_when_empty=True) diff --git a/rhodecode/config/middleware.py b/rhodecode/config/middleware.py index 45700022..9f9978ea 100644 --- a/rhodecode/config/middleware.py +++ b/rhodecode/config/middleware.py @@ -147,6 +147,9 @@ def get_celery_config(settings): return key_name[len(pref) :].replace(".", "_").lower() def type_converter(parsed_key, value): + if not isinstance(value, str): + return value + # cast to int if value.isdigit(): return int(value) diff --git a/rhodecode/integrations/types/slack.py b/rhodecode/integrations/types/slack.py index 07898b1e..8e707957 100644 --- a/rhodecode/integrations/types/slack.py +++ b/rhodecode/integrations/types/slack.py @@ -25,6 +25,7 @@ import deform.widget import colander from rhodecode import events +from rhodecode.lib.celerylib.uniqueue_tasks import RequestContextAndUniqueByPayloadTask from rhodecode.model.validation_schema.widgets import CheckboxChoiceWidgetDesc from rhodecode.translation import _ from rhodecode.lib import helpers as h @@ -140,7 +141,7 @@ class SlackIntegrationType(IntegrationTypeBase): run_task(post_text_to_slack, self.settings, slack_data) -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def post_text_to_slack(settings, slack_data: SlackData): # because JSON serialization, if we run async with celery, deserialize to SlackData if isinstance(slack_data, dict): diff --git a/rhodecode/integrations/types/webhook.py b/rhodecode/integrations/types/webhook.py index 95b0829d..820b0ad0 100644 --- a/rhodecode/integrations/types/webhook.py +++ b/rhodecode/integrations/types/webhook.py @@ -24,6 +24,7 @@ import colander import rhodecode from rhodecode import events +from rhodecode.lib.celerylib.uniqueue_tasks import RequestContextAndUniqueByPayloadTask from rhodecode.lib.colander_utils import strip_whitespace from rhodecode.model.validation_schema.widgets import CheckboxChoiceWidgetDesc from rhodecode.translation import _ @@ -36,7 +37,7 @@ from rhodecode.integrations.types.base import ( requests_retry_call, ) from rhodecode.integrations.types.handlers.webhook import WebhookDataHandler -from rhodecode.lib.celerylib import run_task, async_task, RequestContextTask +from rhodecode.lib.celerylib import run_task, async_task from rhodecode.model.validation_schema import widgets log = logging.getLogger(__name__) @@ -191,7 +192,7 @@ class WebhookIntegrationType(IntegrationTypeBase): run_task(post_to_webhook, self.settings, url_calls) -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def post_to_webhook(settings, url_calls): """ Example data:: diff --git a/rhodecode/lib/celerylib/tasks.py b/rhodecode/lib/celerylib/tasks.py index 52baefc2..2d1e1f6f 100644 --- a/rhodecode/lib/celerylib/tasks.py +++ b/rhodecode/lib/celerylib/tasks.py @@ -37,8 +37,9 @@ import rhodecode from rhodecode.apps.ai_agents.ai_service import get_ai_service from rhodecode.apps.ai_agents.models.base import Response, AIServiceError from rhodecode.lib import audit_logger, diffs, codeblocks -from rhodecode.lib.celerylib import get_logger, async_task, RequestContextTask, run_task +from rhodecode.lib.celerylib import get_logger, async_task, run_task from rhodecode.lib import hooks_base +from rhodecode.lib.celerylib.uniqueue_tasks import UniqueByPayloadTask, RequestContextAndUniqueByPayloadTask from rhodecode.lib.diffs import MAX_CONTEXT from rhodecode.lib.str_utils import safe_str from rhodecode.lib.utils import adopt_for_celery @@ -61,7 +62,7 @@ from rhodecode.model.pull_request import PullRequestModel from rhodecode.model.settings import SettingsModel -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def send_email(recipients, subject, body="", html_body="", email_config=None, extra_headers=None): """ Sends an email with defined parameters from the .ini files. @@ -142,6 +143,7 @@ def send_email(recipients, subject, body="", html_body="", email_config=None, ex return True +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def transform_legacy_email_config(email_config: dict[Any, Any] | Any, mail_server: Any | None) -> dict[ str, None | int | bool | str | Any]: return dict( @@ -177,7 +179,7 @@ def get_mailer(transformed_email_conf: dict[str, Any], original_email_conf: dict return Mailer(**transformed_email_conf) -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def create_repo(form_data, cur_user): from rhodecode.model.repo import RepoModel from rhodecode.model.user import UserModel @@ -286,7 +288,7 @@ def create_repo(form_data, cur_user): return True -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def create_repo_fork(form_data, cur_user): """ Creates a fork of repository using internal VCS methods @@ -375,7 +377,7 @@ def create_repo_fork(form_data, cur_user): return True -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def repo_maintenance(repoid): from rhodecode.lib import repo_maintenance as repo_maintenance_lib @@ -391,7 +393,7 @@ def repo_maintenance(repoid): log.debug("Repo `%s` not found or without a clone_url", repoid) -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def check_for_update(send_email_notification=True, email_recipients=None): from rhodecode.model.update import UpdateModel from rhodecode.model.notification import EmailNotificationModel @@ -450,24 +452,24 @@ def sync_last_update_for_objects(*args, **kwargs): repo_gr.update_commit_cache() -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def test_celery_exception(msg): raise Exception(f"Test exception: {msg}") -@async_task(ignore_result=True, base=RequestContextTask) +@async_task(ignore_result=True, base=RequestContextAndUniqueByPayloadTask) def sync_last_update(*args, **kwargs): sync_last_update_for_objects(*args, **kwargs) -@async_task(ignore_result=False) +@async_task(ignore_result=False, base=UniqueByPayloadTask) def beat_check(*args, **kwargs): log = get_logger(beat_check) log.info("%r: Got args: %r and kwargs %r", beat_check, args, kwargs) return time.time() -@async_task(ignore_result=True) +@async_task(ignore_result=True, base=UniqueByPayloadTask) def schedule_sync_ldap_ad_users_producer(): log = get_logger(schedule_sync_ldap_ad_users_producer) @@ -492,7 +494,7 @@ def schedule_sync_ldap_ad_users_producer(): scheduler.update_from_dict(entries) -@async_task(ignore_result=True) +@async_task(ignore_result=True, base=UniqueByPayloadTask) def unschedule_sync_ldap_ad_users_producer(): log = get_logger(unschedule_sync_ldap_ad_users_producer) @@ -508,7 +510,7 @@ def unschedule_sync_ldap_ad_users_producer(): Session().delete(existing_task) -@async_task +@async_task(base=UniqueByPayloadTask) def start_ai_code_review(pull_request_id): log = get_logger(start_ai_code_review) log.info("Starting AI code review for pull request id: %s", pull_request_id) @@ -769,7 +771,7 @@ def _get_diffset( return diffset.render_patchset(_parsed, source_ref=source_ref_id, target_ref=target_ref_id) -@async_task +@async_task(base=UniqueByPayloadTask) @adopt_for_celery def repo_size(extras): from rhodecode.lib.hooks_base import repo_size @@ -777,7 +779,7 @@ def repo_size(extras): return repo_size(extras) -@async_task +@async_task(base=UniqueByPayloadTask) @adopt_for_celery def pre_pull(extras): from rhodecode.lib.hooks_base import pre_pull @@ -785,7 +787,7 @@ def pre_pull(extras): return pre_pull(extras) -@async_task +@async_task(base=UniqueByPayloadTask) @adopt_for_celery def post_pull(extras): from rhodecode.lib.hooks_base import post_pull @@ -793,7 +795,7 @@ def post_pull(extras): return post_pull(extras) -@async_task +@async_task(base=UniqueByPayloadTask) @adopt_for_celery def pre_push(extras): from rhodecode.lib.hooks_base import pre_push @@ -801,7 +803,7 @@ def pre_push(extras): return pre_push(extras) -@async_task +@async_task(base=UniqueByPayloadTask) @adopt_for_celery def post_push(extras): from rhodecode.lib.hooks_base import post_push diff --git a/rhodecode/lib/celerylib/tests/__init__.py b/rhodecode/lib/celerylib/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/rhodecode/lib/celerylib/tests/test_uniqueue_task.py b/rhodecode/lib/celerylib/tests/test_uniqueue_task.py new file mode 100644 index 00000000..17129f00 --- /dev/null +++ b/rhodecode/lib/celerylib/tests/test_uniqueue_task.py @@ -0,0 +1,14 @@ +from rhodecode.lib.celerylib.uniqueue_tasks import stable_payload_key + + +class TestUniqueTask: + def test_stable_payload_key_generation(self): + task_name = "test_task" + task_args = (1, 2, 3) + task_kwargs = {"arg1": 10} + + expected_output = "celery:unique:test_task:463d9df63be487de5fd2073cfc26a6813713ee96c0cf9f7b5f949e12b8643511" + + for _ in range(10): + res = stable_payload_key(task_name, task_args, task_kwargs) + assert res == expected_output diff --git a/rhodecode/lib/celerylib/uniqueue_tasks.py b/rhodecode/lib/celerylib/uniqueue_tasks.py new file mode 100644 index 00000000..ba1452af --- /dev/null +++ b/rhodecode/lib/celerylib/uniqueue_tasks.py @@ -0,0 +1,139 @@ +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