diff --git a/api/.importlinter b/api/.importlinter index d53f0c2e1c9d74..bc880da786a821 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1222,6 +1222,7 @@ source_modules = services.trigger.subscription_query_service services.trigger.subscription_service services.trigger.subscription_renewal_service + services.trigger.subscription_refresh_service services.trigger.subscription_binding_service forbidden_modules = controllers diff --git a/api/extensions/application_services/workspace.py b/api/extensions/application_services/workspace.py index aab411d998f4f3..77d5ddb49b5b1c 100644 --- a/api/extensions/application_services/workspace.py +++ b/api/extensions/application_services/workspace.py @@ -16,7 +16,11 @@ from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository from repositories.trigger.oauth_client_repository import TriggerOAuthClientRepository from repositories.trigger.subscription_builder_repository import RedisTriggerSubscriptionBuilderRepository -from repositories.trigger.subscription_repository import RedisTriggerSubscriptionLocks, TriggerSubscriptionRepository +from repositories.trigger.subscription_repository import ( + RedisTriggerSubscriptionLocks, + RedisTriggerSubscriptionRefreshReservations, + TriggerSubscriptionRepository, +) from repositories.workspace.resource_cleanup_repository import WorkspaceResourceCleanupRepository from repositories.workspace.workspace_repository import WorkspaceRepository from services.account.adapters import RedisInvitationTokenStore @@ -40,8 +44,13 @@ TriggerSubscriptionBindingQuery, TriggerSubscriptionBindingService, ) -from services.trigger.subscription_gateway import PluginTriggerSubscriptionGateway, invalidate_subscription_properties +from services.trigger.subscription_gateway import ( + PluginTriggerSubscriptionGateway, + enqueue_subscription_refresh, + invalidate_subscription_properties, +) from services.trigger.subscription_query_service import TriggerSubscriptionQueryService +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshService from services.trigger.subscription_renewal_service import TriggerSubscriptionRenewalService from services.trigger.subscription_service import TriggerSubscriptionService from services.trigger.trigger_request_service import TriggerHttpRequestCachingService @@ -88,6 +97,7 @@ class WorkspaceServices: create_trigger_binding_service: Callable[[Session], TriggerSubscriptionBindingService] trigger_subscription_management: TriggerSubscriptionService trigger_subscription_renewal: TriggerSubscriptionRenewalService + trigger_subscription_refresh: TriggerSubscriptionRefreshService trigger_builders: TriggerSubscriptionBuilderService trigger_endpoints: TriggerEndpointService trigger_events: TriggerEventService @@ -153,6 +163,14 @@ def build_workspace_services( trigger_subscriptions = TriggerSubscriptionQueryService( credentials=credentials, schemas=trigger_providers, subscriptions=subscriptions ) + renewal = TriggerSubscriptionRenewalService( + subscriptions=subscriptions, + providers=subscription_provider, + oauth=trigger_oauth_provider, + clients=trigger_oauth_clients, + invalidate_cache=delete_cache_for_subscription, + invalidate_properties=invalidate_subscription_properties, + ) return WorkspaceServices( queries=WorkspaceQueryService(workspaces=workspaces, plans=DeploymentWorkspacePlanGateway()), management=WorkspaceService( @@ -204,13 +222,12 @@ def build_workspace_services( requests=TriggerHttpRequestCachingService, ), trigger_subscription_management=trigger_subscription_management, - trigger_subscription_renewal=TriggerSubscriptionRenewalService( + trigger_subscription_renewal=renewal, + trigger_subscription_refresh=TriggerSubscriptionRefreshService( subscriptions=subscriptions, - providers=subscription_provider, - oauth=trigger_oauth_provider, - clients=trigger_oauth_clients, - invalidate_cache=delete_cache_for_subscription, - invalidate_properties=invalidate_subscription_properties, + reservations=RedisTriggerSubscriptionRefreshReservations(redis=redis), + publish=enqueue_subscription_refresh, + renewal=renewal, ), trigger_builders=trigger_builders, trigger_endpoints=TriggerEndpointService( diff --git a/api/repositories/trigger/subscription_repository.py b/api/repositories/trigger/subscription_repository.py index da54ee5efc8f63..3a44237d7c6187 100644 --- a/api/repositories/trigger/subscription_repository.py +++ b/api/repositories/trigger/subscription_repository.py @@ -6,19 +6,22 @@ from contextlib import AbstractContextManager, contextmanager from copy import deepcopy from typing import Any, override +from uuid import uuid4 from psycopg2.errors import UniqueViolation from pymysql.err import IntegrityError as MySQLIntegrityError from redis.exceptions import LockError, RedisError -from sqlalchemy import delete, func, select +from sqlalchemy import and_, delete, func, or_, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, sessionmaker +from core.trigger.utils.locks import build_trigger_refresh_lock_key from extensions.ext_redis import RedisClientWrapper from models.enums import AppTriggerStatus from models.provider_ids import TriggerProviderID from models.trigger import AppTrigger, TriggerSubscription, WorkflowPluginTrigger from services.trigger.subscription_binding_service import TriggerSubscriber, TriggerSubscriptionBinding +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshJob, TriggerSubscriptionRefreshRef from services.trigger.subscription_service import ( TriggerSubscriptionChangedError, TriggerSubscriptionLimitError, @@ -58,6 +61,33 @@ class TriggerSubscriptionRepository(TriggerSubscriptionStore): def __init__(self, *, session_factory: sessionmaker[Session]) -> None: self._session_factory = session_factory + def list_due_for_refresh( + self, + *, + credential_due_at: int, + subscription_due_at: int, + after_id: str | None, + limit: int, + ) -> tuple[TriggerSubscriptionRefreshRef, ...]: + """Scan all workspaces for the scheduler, releasing SQL before reservation or publishing.""" + statement = select(TriggerSubscription.tenant_id, TriggerSubscription.id).where( + or_( + and_( + TriggerSubscription.credential_expires_at != -1, + TriggerSubscription.credential_expires_at <= credential_due_at, + ), + and_(TriggerSubscription.expires_at != -1, TriggerSubscription.expires_at <= subscription_due_at), + ) + ) + if after_id is not None: + statement = statement.where(TriggerSubscription.id > after_id) + with self._session_factory() as session: + rows = session.execute(statement.order_by(TriggerSubscription.id).limit(limit)).all() + return tuple( + TriggerSubscriptionRefreshRef(tenant_id=tenant_id, subscription_id=subscription_id) + for tenant_id, subscription_id in rows + ) + def sync_workflow_bindings( self, *, @@ -354,3 +384,38 @@ def _hold(self, key: str) -> Generator[None]: lock.release() except RedisError: logger.exception("Failed to release trigger subscription lock %s", key) + + +class RedisTriggerSubscriptionRefreshReservations: + _RELEASE = """ + if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) + end + return 0 + """ + + def __init__(self, *, redis: RedisClientWrapper) -> None: + self._redis = redis + + def reserve( + self, *, subscriptions: Sequence[TriggerSubscriptionRefreshRef], ttl: int + ) -> tuple[TriggerSubscriptionRefreshJob, ...]: + jobs = tuple( + TriggerSubscriptionRefreshJob(subscription=subscription, token=str(uuid4())) + for subscription in subscriptions + ) + with self._redis.pipeline(transaction=False) as pipe: + for job in jobs: + pipe.set(self._key(job), job.token, ex=ttl, nx=True) + acquired = pipe.execute() + return tuple(job for job, owned in zip(jobs, acquired, strict=True) if owned) + + def owns(self, job: TriggerSubscriptionRefreshJob) -> bool: + return self._redis.get(self._key(job)) in (job.token, job.token.encode()) + + def release(self, job: TriggerSubscriptionRefreshJob) -> None: + self._redis.eval(self._RELEASE, 1, self._key(job), job.token) + + @staticmethod + def _key(job: TriggerSubscriptionRefreshJob) -> str: + return build_trigger_refresh_lock_key(job.subscription.tenant_id, job.subscription.subscription_id) diff --git a/api/schedule/trigger_provider_refresh_task.py b/api/schedule/trigger_provider_refresh_task.py index b7b9d67da271c7..f9deaaab38459a 100644 --- a/api/schedule/trigger_provider_refresh_task.py +++ b/api/schedule/trigger_provider_refresh_task.py @@ -1,106 +1,16 @@ -import logging -import math -from collections.abc import Iterable, Sequence +from celery import shared_task -from celery import group -from sqlalchemy import ColumnElement, and_, func, or_, select -from sqlalchemy.engine.row import Row -from sqlalchemy.orm import Session - -import app from configs import dify_config -from core.trigger.utils.locks import build_trigger_refresh_lock_keys -from extensions.ext_database import db -from extensions.ext_redis import redis_client +from extensions.ext_application_services import application_services from libs.helper import current_timestamp -from models.trigger import TriggerSubscription -from tasks.trigger_subscription_refresh_tasks import trigger_subscription_refresh - -logger = logging.getLogger(__name__) - - -def _build_due_filter(now_ts: int): - """Build SQLAlchemy filter for due credential or subscription refresh.""" - credential_due: ColumnElement[bool] = and_( - TriggerSubscription.credential_expires_at != -1, - TriggerSubscription.credential_expires_at - <= now_ts + int(dify_config.TRIGGER_PROVIDER_CREDENTIAL_THRESHOLD_SECONDS), - ) - subscription_due: ColumnElement[bool] = and_( - TriggerSubscription.expires_at != -1, - TriggerSubscription.expires_at <= now_ts + int(dify_config.TRIGGER_PROVIDER_SUBSCRIPTION_THRESHOLD_SECONDS), - ) - return or_(credential_due, subscription_due) - - -def _acquire_locks(keys: Iterable[str], ttl_seconds: int) -> list[bool]: - """Attempt to acquire locks in a single pipelined round-trip. - - Returns a list of booleans indicating which locks were acquired. - """ - pipe = redis_client.pipeline(transaction=False) - for key in keys: - pipe.set(key, b"1", ex=ttl_seconds, nx=True) - results = pipe.execute() - return [bool(r) for r in results] -@app.celery.task(queue="trigger_refresh_publisher") +@shared_task(queue="trigger_refresh_publisher") def trigger_provider_refresh() -> None: - """ - Scan due trigger subscriptions and enqueue refresh tasks with in-flight locks. - """ - now: int = current_timestamp() - - batch_size: int = int(dify_config.TRIGGER_PROVIDER_REFRESH_BATCH_SIZE) - lock_ttl: int = max(300, int(dify_config.TRIGGER_PROVIDER_SUBSCRIPTION_THRESHOLD_SECONDS)) - - with Session(db.engine, expire_on_commit=False) as session: - filter: ColumnElement[bool] = _build_due_filter(now_ts=now) - total_due: int = int(session.scalar(statement=select(func.count()).where(filter)) or 0) - logger.info("Trigger refresh scan start: due=%d", total_due) - if total_due == 0: - return - - pages: int = math.ceil(total_due / batch_size) - for page in range(pages): - offset: int = page * batch_size - subscription_rows: Sequence[Row[tuple[str, str]]] = session.execute( - select(TriggerSubscription.tenant_id, TriggerSubscription.id) - .where(filter) - .order_by(TriggerSubscription.updated_at.asc()) - .offset(offset) - .limit(batch_size) - ).all() - if not subscription_rows: - logger.debug("Trigger refresh page %d/%d empty", page + 1, pages) - continue - - subscriptions: list[tuple[str, str]] = [ - (str(tenant_id), str(subscription_id)) for tenant_id, subscription_id in subscription_rows - ] - lock_keys: list[str] = build_trigger_refresh_lock_keys(subscriptions) - acquired: list[bool] = _acquire_locks(keys=lock_keys, ttl_seconds=lock_ttl) - - if not any(acquired): - continue - - jobs = [ - trigger_subscription_refresh.s(tenant_id=tenant_id, subscription_id=subscription_id) - for (tenant_id, subscription_id), is_locked in zip(subscriptions, acquired, strict=False) - if is_locked - ] - result = group(jobs).apply_async() - enqueued = len(jobs) - - logger.info( - "Trigger refresh page %d/%d: scanned=%d locks_acquired=%d enqueued=%d result=%s", - page + 1, - pages, - len(subscriptions), - sum(1 for x in acquired if x), - enqueued, - result, - ) - - logger.info("Trigger refresh scan done: due=%d", total_due) + """Scan due subscriptions in bounded batches and enqueue independently reserved jobs.""" + application_services().workspaces.trigger_subscription_refresh.enqueue_due( + now=current_timestamp(), + credential_threshold=dify_config.TRIGGER_PROVIDER_CREDENTIAL_THRESHOLD_SECONDS, + subscription_threshold=dify_config.TRIGGER_PROVIDER_SUBSCRIPTION_THRESHOLD_SECONDS, + batch_size=dify_config.TRIGGER_PROVIDER_REFRESH_BATCH_SIZE, + ) diff --git a/api/services/trigger/subscription_gateway.py b/api/services/trigger/subscription_gateway.py index 82b10c56fae00d..3c5ba2eb93d16e 100644 --- a/api/services/trigger/subscription_gateway.py +++ b/api/services/trigger/subscription_gateway.py @@ -11,6 +11,7 @@ from core.trigger.utils.encryption import TriggerProviderPropertiesCache from models.provider_ids import TriggerProviderID from services.trigger.subscription_query_service import TriggerSubscriptionSchemas +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshJob from services.trigger.subscription_service import ( TriggerSubscriptionInvalidError, TriggerSubscriptionProviderError, @@ -108,3 +109,13 @@ def invalidate_subscription_properties(*, tenant_id: str, provider_id: str, subs TriggerProviderPropertiesCache( tenant_id=tenant_id, provider_id=provider_id, subscription_id=subscription_id ).delete() + + +def enqueue_subscription_refresh(job: TriggerSubscriptionRefreshJob) -> None: + from tasks.trigger_subscription_refresh_tasks import trigger_subscription_refresh + + # Keep the two-argument task contract for older workers during rolling upgrades. + trigger_subscription_refresh.apply_async( + kwargs={"tenant_id": job.subscription.tenant_id, "subscription_id": job.subscription.subscription_id}, + headers={"reservation_token": job.token}, + ) diff --git a/api/services/trigger/subscription_refresh_service.py b/api/services/trigger/subscription_refresh_service.py new file mode 100644 index 00000000000000..efd1cbd7d01b00 --- /dev/null +++ b/api/services/trigger/subscription_refresh_service.py @@ -0,0 +1,111 @@ +"""Scan detached renewal candidates and coordinate their queued reservations.""" + +import logging +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True) +class TriggerSubscriptionRefreshRef: + tenant_id: str + subscription_id: str + + +@dataclass(frozen=True, slots=True) +class TriggerSubscriptionRefreshJob: + subscription: TriggerSubscriptionRefreshRef + token: str + + +class TriggerSubscriptionRefreshQuery(Protocol): + def list_due_for_refresh( + self, + *, + credential_due_at: int, + subscription_due_at: int, + after_id: str | None, + limit: int, + ) -> tuple[TriggerSubscriptionRefreshRef, ...]: + """None starts the scan; otherwise continue after the last immutable subscription ID.""" + ... + + +class TriggerSubscriptionRefreshReservations(Protocol): + def reserve( + self, *, subscriptions: Sequence[TriggerSubscriptionRefreshRef], ttl: int + ) -> tuple[TriggerSubscriptionRefreshJob, ...]: ... + + def owns(self, job: TriggerSubscriptionRefreshJob) -> bool: ... + + def release(self, job: TriggerSubscriptionRefreshJob) -> None: ... + + +class TriggerSubscriptionRefreshPublisher(Protocol): + def __call__(self, job: TriggerSubscriptionRefreshJob) -> None: ... + + +class TriggerSubscriptionRenewer(Protocol): + def refresh_if_due(self, *, tenant_id: str, subscription_id: str, now: int) -> None: ... + + +class TriggerSubscriptionRefreshService: + def __init__( + self, + *, + subscriptions: TriggerSubscriptionRefreshQuery, + reservations: TriggerSubscriptionRefreshReservations, + publish: TriggerSubscriptionRefreshPublisher, + renewal: TriggerSubscriptionRenewer, + ) -> None: + self._subscriptions = subscriptions + self._reservations = reservations + self._publish = publish + self._renewal = renewal + + def enqueue_due(self, *, now: int, credential_threshold: int, subscription_threshold: int, batch_size: int) -> int: + if batch_size <= 0: + raise ValueError("Trigger subscription refresh batch size must be positive.") + after_id = None + enqueued = 0 + while subscriptions := self._subscriptions.list_due_for_refresh( + credential_due_at=now + credential_threshold, + subscription_due_at=now + subscription_threshold, + after_id=after_id, + limit=batch_size, + ): + # Renewed rows can leave the due set while the scan runs; OFFSET would skip later rows. + after_id = subscriptions[-1].subscription_id + jobs = self._reservations.reserve(subscriptions=subscriptions, ttl=max(300, subscription_threshold)) + for index, job in enumerate(jobs): + try: + self._publish(job) + except Exception: + logger.exception("Failed to publish trigger renewal: %s", job.subscription) + # Keep already published reservations; failed and unsent jobs can retry next tick. + for unpublished in jobs[index:]: + self._release(unpublished) + raise + else: + enqueued += 1 + logger.info("Trigger refresh scan done: enqueued=%d", enqueued) + return enqueued + + def execute(self, *, job: TriggerSubscriptionRefreshJob, now: int) -> None: + if not self._reservations.owns(job): + logger.debug("Refresh reservation missing or replaced: %s", job.subscription) + return + try: + self._renewal.refresh_if_due( + tenant_id=job.subscription.tenant_id, subscription_id=job.subscription.subscription_id, now=now + ) + finally: + self._release(job) + + def _release(self, job: TriggerSubscriptionRefreshJob) -> None: + try: + self._reservations.release(job) + except Exception: + logger.warning("Failed to release trigger refresh reservation: %s", job.subscription, exc_info=True) diff --git a/api/tasks/trigger_subscription_refresh_tasks.py b/api/tasks/trigger_subscription_refresh_tasks.py index 333cbcbecc48de..aa330f9dc515f9 100644 --- a/api/tasks/trigger_subscription_refresh_tasks.py +++ b/api/tasks/trigger_subscription_refresh_tasks.py @@ -1,32 +1,20 @@ -import logging +from __future__ import annotations -from celery import shared_task +from celery import Task, shared_task -from core.trigger.utils.locks import build_trigger_refresh_lock_key from extensions.ext_application_services import application_services -from extensions.ext_redis import redis_client from libs.helper import current_timestamp +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshJob, TriggerSubscriptionRefreshRef -logger = logging.getLogger(__name__) - -@shared_task(queue="trigger_refresh_executor") -def trigger_subscription_refresh(tenant_id: str, subscription_id: str) -> None: - """Refresh a trigger subscription if needed, guarded by a Redis in-flight lock.""" - lock_key: str = build_trigger_refresh_lock_key(tenant_id, subscription_id) - if not redis_client.get(lock_key): - logger.debug("Refresh lock missing, skip: %s", lock_key) - return - - logger.info("Begin subscription refresh: tenant=%s id=%s", tenant_id, subscription_id) - try: - application_services().workspaces.trigger_subscription_renewal.refresh_if_due( - tenant_id=tenant_id, subscription_id=subscription_id, now=current_timestamp() - ) - finally: - try: - redis_client.delete(lock_key) - logger.debug("Lock released: %s", lock_key) - except Exception: - # Best-effort lock cleanup - logger.warning("Failed to release lock: %s", lock_key, exc_info=True) +@shared_task(queue="trigger_refresh_executor", bind=True) +def trigger_subscription_refresh(self: Task[[str, str], None], tenant_id: str, subscription_id: str) -> None: + """Read reservation ownership from headers; old publishers used the fixed token '1'.""" + reservation_token = (self.request.headers or {}).get("reservation_token", "1") + application_services().workspaces.trigger_subscription_refresh.execute( + job=TriggerSubscriptionRefreshJob( + subscription=TriggerSubscriptionRefreshRef(tenant_id=tenant_id, subscription_id=subscription_id), + token=reservation_token, + ), + now=current_timestamp(), + ) diff --git a/api/tests/test_containers_integration_tests/services/test_trigger_subscription_refresh.py b/api/tests/test_containers_integration_tests/services/test_trigger_subscription_refresh.py new file mode 100644 index 00000000000000..5acdc29f1e19dc --- /dev/null +++ b/api/tests/test_containers_integration_tests/services/test_trigger_subscription_refresh.py @@ -0,0 +1,275 @@ +"""SQLite, real Redis and the actual memory-broker publisher exercise renewal scheduling.""" + +import json +from collections.abc import Iterator +from contextvars import Context +from dataclasses import replace +from uuid import UUID, uuid4 + +import pytest +from celery import Celery +from celery.signals import before_task_publish +from kombu.exceptions import EncodeError +from kombu.serialization import register, unregister +from kombu.simple import SimpleQueue +from sqlalchemy import Engine, update +from sqlalchemy.exc import OperationalError +from sqlalchemy.pool import QueuePool + +from core.trigger.utils.locks import build_trigger_refresh_lock_key +from extensions.ext_application_services import application_services +from extensions.ext_redis import RedisClientWrapper +from models.base import TypeBase +from models.trigger import TriggerSubscription +from repositories.trigger.subscription_repository import RedisTriggerSubscriptionRefreshReservations +from schedule.trigger_provider_refresh_task import trigger_provider_refresh +from services.trigger.subscription_gateway import enqueue_subscription_refresh +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshJob, TriggerSubscriptionRefreshRef +from tasks.trigger_subscription_refresh_tasks import trigger_subscription_refresh +from tests.test_containers_integration_tests.services import test_trigger_subscription_builder +from tests.unit_tests.config_override import config_overrides_context +from tests.unit_tests.controllers.console.workspace import test_trigger_oauth_client +from tests.unit_tests.controllers.console.workspace.test_trigger_oauth_client import _Harness +from tests.unit_tests.repositories.trigger.test_subscription_repository import _persist, _snapshot + +application_redis = test_trigger_subscription_builder.application_redis +sqlite_engine = test_trigger_subscription_builder.sqlite_engine +sqlite_session_factory = test_trigger_subscription_builder.sqlite_session_factory +harness = test_trigger_oauth_client.harness + + +@pytest.fixture +def broker() -> Iterator[tuple[Celery, SimpleQueue]]: + previous_app = trigger_subscription_refresh.app + app = Celery("trigger-renewal-scan-test", broker="memory://", set_as_current=False) + app.conf.update(task_always_eager=False, task_serializer="json") + app.set_current() + try: + with app.connection_for_write() as connection: + with SimpleQueue(connection.channel(), "trigger_refresh_executor") as queue: + queue.clear() + try: + yield app, queue + finally: + queue.clear() + finally: + previous_app.set_current() + app.close() + + +def _ref(*, tenant_id: str) -> TriggerSubscriptionRefreshRef: + return TriggerSubscriptionRefreshRef(tenant_id=tenant_id, subscription_id=str(uuid4())) + + +def _key(ref: TriggerSubscriptionRefreshRef) -> str: + return build_trigger_refresh_lock_key(ref.tenant_id, ref.subscription_id) + + +def _scan(harness: _Harness, *, batch_size: int = 2) -> int: + with harness.app.app_context(): + return application_services().workspaces.trigger_subscription_refresh.enqueue_due( + now=100, credential_threshold=10, subscription_threshold=20, batch_size=batch_size + ) + + +def test_scheduler_releases_sql_before_reservation_and_publish_and_scans_past_renewed_rows( + harness: _Harness, + application_redis: RedisClientWrapper, + sqlite_engine: Engine, + broker: tuple[Celery, SimpleQueue], +) -> None: + snapshots = [ + replace(_snapshot(), id=str(UUID(int=index)), tenant_id=harness.tenant_id, name=f"due-{index}") + for index in range(1, 6) + ] + _persist(harness.session_factory, *snapshots) + reservations = RedisTriggerSubscriptionRefreshReservations(redis=application_redis) + locked = reservations.reserve( + subscriptions=[TriggerSubscriptionRefreshRef(tenant_id=harness.tenant_id, subscription_id=snapshots[2].id)], + ttl=300, + ) + assert len(locked) == 1 + pool = sqlite_engine.pool + assert isinstance(pool, QueuePool) + publications: list[str] = [] + + def on_publish( + sender: str, body: tuple[object, dict[str, str], object], headers: dict[str, object], **_kwargs: object + ) -> None: + assert sender == trigger_subscription_refresh.name + assert pool.checkedout() == 0 + job = body[1] + assert set(job) == {"tenant_id", "subscription_id"} + reservation_token = headers["reservation_token"] + assert isinstance(reservation_token, str) + assert ( + application_redis.get( + _key(TriggerSubscriptionRefreshRef(tenant_id=job["tenant_id"], subscription_id=job["subscription_id"])) + ) + == reservation_token.encode() + ) + publications.append(job["subscription_id"]) + # Simulate a fast worker removing already queued rows from the due set. + with harness.session_factory.begin() as session: + session.execute( + update(TriggerSubscription) + .where(TriggerSubscription.id == job["subscription_id"]) + .values(credential_expires_at=1000) + ) + + before_task_publish.connect(on_publish, weak=False, sender=trigger_subscription_refresh.name) + try: + assert _scan(harness) == 4 + finally: + before_task_publish.disconnect( # pyrefly: ignore[missing-attribute] + on_publish, sender=trigger_subscription_refresh.name + ) + assert publications == [snapshots[index].id for index in (0, 1, 3, 4)] + assert broker[1].qsize() == 4 + assert _scan(harness) == 0 + for snapshot in snapshots: + key = build_trigger_refresh_lock_key(snapshot.tenant_id, snapshot.id) + assert 0 < application_redis.ttl(key) <= 300 + application_redis.delete(key) + + +def test_failed_serialization_keeps_published_reservation_and_releases_failed_and_unsent_jobs( + harness: _Harness, application_redis: RedisClientWrapper, broker: tuple[Celery, SimpleQueue] +) -> None: + snapshots = [ + replace(_snapshot(), id=str(UUID(int=index)), tenant_id=harness.tenant_id, name=f"job-{index}") + for index in range(1, 4) + ] + _persist(harness.session_factory, *snapshots) + attempts = 0 + reject = True + + def encode(payload: object) -> str: + nonlocal attempts + attempts += 1 + if reject and attempts > 1: + raise ValueError("Serializer failed after publishing the first job") + return json.dumps(payload) + + serializer = "trigger-refresh-test" + register( + serializer, encode, json.loads, content_type="application/x-trigger-refresh-test", content_encoding="utf-8" + ) + broker[0].conf.task_serializer = serializer + keys = [build_trigger_refresh_lock_key(snapshot.tenant_id, snapshot.id) for snapshot in snapshots] + try: + with pytest.raises(EncodeError, match="Serializer failed"): + _scan(harness, batch_size=3) + assert [application_redis.exists(key) for key in keys] == [1, 0, 0] + assert broker[1].qsize() == 1 + + reject = False + assert _scan(harness, batch_size=3) == 2 + assert broker[1].qsize() == 3 + assert all(application_redis.exists(key) for key in keys) + finally: + unregister(serializer) + for key in keys: + application_redis.delete(key) + + +def test_reservation_scope_and_expired_owner_cannot_release_replacement( + harness: _Harness, application_redis: RedisClientWrapper +) -> None: + store = RedisTriggerSubscriptionRefreshReservations(redis=application_redis) + ref = _ref(tenant_id=harness.tenant_id) + other_tenant = replace(ref, tenant_id=str(uuid4())) + first = store.reserve(subscriptions=[ref, other_tenant], ttl=300) + assert len(first) == 2 + assert store.reserve(subscriptions=[ref, other_tenant], ttl=300) == () + application_redis.delete(_key(ref)) + (replacement,) = store.reserve(subscriptions=[ref], ttl=300) + + assert not store.owns(first[0]) + store.release(first[0]) + assert store.owns(replacement) + assert store.owns(first[1]) + store.release(replacement) + store.release(first[1]) + + +@pytest.mark.parametrize("publisher", ["legacy", "new"]) +def test_worker_accepts_legacy_and_new_reservations_and_releases_missing_subscription( + harness: _Harness, + application_redis: RedisClientWrapper, + broker: tuple[Celery, SimpleQueue], + publisher: str, +) -> None: + ref = _ref(tenant_id=harness.tenant_id) + token = "1" if publisher == "legacy" else "new-owner" + application_redis.setex(_key(ref), 300, token) + if publisher == "legacy": + broker[0].send_task( + trigger_subscription_refresh.name, + kwargs={"tenant_id": ref.tenant_id, "subscription_id": ref.subscription_id}, + queue="trigger_refresh_executor", + ) + else: + enqueue_subscription_refresh(TriggerSubscriptionRefreshJob(subscription=ref, token=token)) + message = broker[1].get(block=False) + + def run() -> None: + with harness.app.app_context(): + result = trigger_subscription_refresh.apply( + args=tuple(message.payload[0]), kwargs=message.payload[1], headers=message.headers, throw=True + ) + assert result.successful() + + Context().run(run) + message.ack() + assert application_redis.exists(_key(ref)) == 0 + + +@pytest.mark.usefixtures("broker") +@pytest.mark.parametrize("token", [None, "old-owner"], ids=["legacy-publisher", "new-publisher"]) +def test_stale_worker_does_not_use_or_delete_another_reservation( + harness: _Harness, application_redis: RedisClientWrapper, token: str | None +) -> None: + ref = _ref(tenant_id=harness.tenant_id) + application_redis.setex(_key(ref), 300, "replacement-owner") + with harness.app.app_context(): + result = trigger_subscription_refresh.apply( + kwargs={"tenant_id": ref.tenant_id, "subscription_id": ref.subscription_id}, + headers=None if token is None else {"reservation_token": token}, + throw=True, + ) + assert result.successful() + assert application_redis.get(_key(ref)) == b"replacement-owner" + application_redis.delete(_key(ref)) + + +def test_worker_releases_owned_reservation_when_repository_read_fails( + harness: _Harness, application_redis: RedisClientWrapper, sqlite_engine: Engine +) -> None: + ref = _ref(tenant_id=harness.tenant_id) + job = TriggerSubscriptionRefreshJob(subscription=ref, token="owned") + application_redis.setex(_key(ref), 300, job.token) + TypeBase.metadata.tables["trigger_subscriptions"].drop(sqlite_engine) + with harness.app.app_context(), pytest.raises(OperationalError): + application_services().workspaces.trigger_subscription_refresh.execute(job=job, now=100) + assert application_redis.exists(_key(ref)) == 0 + + +def test_scheduled_task_uses_injected_service_and_configured_thresholds( + harness: _Harness, application_redis: RedisClientWrapper, broker: tuple[Celery, SimpleQueue] +) -> None: + snapshot = replace(_snapshot(), tenant_id=harness.tenant_id) + _persist(harness.session_factory, snapshot) + with ( + harness.app.app_context(), + config_overrides_context( + TRIGGER_PROVIDER_REFRESH_BATCH_SIZE=1, + TRIGGER_PROVIDER_CREDENTIAL_THRESHOLD_SECONDS=10, + TRIGGER_PROVIDER_SUBSCRIPTION_THRESHOLD_SECONDS=400, + ), + ): + trigger_provider_refresh.run() + assert broker[1].qsize() == 1 + key = build_trigger_refresh_lock_key(snapshot.tenant_id, snapshot.id) + assert 300 < application_redis.ttl(key) <= 400 + application_redis.delete(key) diff --git a/api/tests/test_containers_integration_tests/services/test_trigger_subscription_renewal.py b/api/tests/test_containers_integration_tests/services/test_trigger_subscription_renewal.py index 6693272a8a515e..1d719dfefbd93f 100644 --- a/api/tests/test_containers_integration_tests/services/test_trigger_subscription_renewal.py +++ b/api/tests/test_containers_integration_tests/services/test_trigger_subscription_renewal.py @@ -397,7 +397,7 @@ def test_task_requires_existing_reservation_and_releases_it_after_processing( subscription_id = "missing-subscription" if job_state == "missing-subscription" else subscription.id key = build_trigger_refresh_lock_key(harness.tenant_id, subscription_id) if job_state != "no-lock": - redis_client.setex(key, 60, "reserved") + redis_client.setex(key, 60, "1") def run() -> None: with harness.app.app_context(): diff --git a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py index 28c6ec83fdd767..d667f382ad413c 100644 --- a/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/inner_api/workspace/test_workspace.py @@ -26,7 +26,11 @@ from repositories.credentials.query_repository import CredentialQueryRepository from repositories.trigger.oauth_client_repository import TriggerOAuthClientRepository from repositories.trigger.subscription_builder_repository import RedisTriggerSubscriptionBuilderRepository -from repositories.trigger.subscription_repository import RedisTriggerSubscriptionLocks, TriggerSubscriptionRepository +from repositories.trigger.subscription_repository import ( + RedisTriggerSubscriptionLocks, + RedisTriggerSubscriptionRefreshReservations, + TriggerSubscriptionRepository, +) from services.account_errors import AccountNotFoundError from services.errors.workspace import ( InvalidWorkspaceMemberRoleError, @@ -49,8 +53,13 @@ from services.trigger.oauth_service import TriggerOAuthService from services.trigger.provider_gateway import PluginTriggerProviderQuery from services.trigger.subscription_binding_service import TriggerSubscriptionBindingService -from services.trigger.subscription_gateway import PluginTriggerSubscriptionGateway, invalidate_subscription_properties +from services.trigger.subscription_gateway import ( + PluginTriggerSubscriptionGateway, + enqueue_subscription_refresh, + invalidate_subscription_properties, +) from services.trigger.subscription_query_service import TriggerSubscriptionQueryService +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshService from services.trigger.subscription_renewal_service import TriggerSubscriptionRenewalService from services.trigger.subscription_service import TriggerSubscriptionService from services.trigger.trigger_request_service import TriggerHttpRequestCachingService @@ -92,6 +101,14 @@ def provisioning(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessio schemas=PluginTriggerProviderQuery(), subscriptions=TriggerSubscriptionRepository(session_factory=sqlite_session_factory), ) + renewal = TriggerSubscriptionRenewalService( + subscriptions=TriggerSubscriptionRepository(session_factory=sqlite_session_factory), + providers=PluginTriggerSubscriptionGateway(), + oauth=oauth_provider, + clients=oauth_clients, + invalidate_cache=delete_cache_for_subscription, + invalidate_properties=invalidate_subscription_properties, + ) services = WorkspaceServices( queries=create_autospec(WorkspaceQueryService, instance=True), management=create_autospec(WorkspaceService, instance=True), @@ -130,13 +147,12 @@ def provisioning(monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessio publish=enqueue_triggered_workflows, ), trigger_subscription_management=subscription_management, - trigger_subscription_renewal=TriggerSubscriptionRenewalService( + trigger_subscription_renewal=renewal, + trigger_subscription_refresh=TriggerSubscriptionRefreshService( subscriptions=TriggerSubscriptionRepository(session_factory=sqlite_session_factory), - providers=PluginTriggerSubscriptionGateway(), - oauth=oauth_provider, - clients=oauth_clients, - invalidate_cache=delete_cache_for_subscription, - invalidate_properties=invalidate_subscription_properties, + reservations=RedisTriggerSubscriptionRefreshReservations(redis=redis), + publish=enqueue_subscription_refresh, + renewal=renewal, ), trigger_oauth_clients=oauth_clients, trigger_oauth=TriggerOAuthService( diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index 3a6828a49f59a5..1c32aa0d386ad3 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -1050,6 +1050,11 @@ def test_build_application_services_wires_credential_query( services.workspaces.trigger_subscription_renewal._providers is services.workspaces.trigger_subscription_management._providers ) + assert services.workspaces.trigger_subscription_refresh._renewal is services.workspaces.trigger_subscription_renewal + assert ( + services.workspaces.trigger_subscription_refresh._subscriptions + is services.workspaces.trigger_subscription_management._subscriptions + ) assert services.workspaces.trigger_subscription_renewal._clients is services.workspaces.trigger_oauth_clients assert services.workspaces.trigger_subscription_renewal._oauth is services.workspaces.trigger_oauth._providers diff --git a/api/tests/unit_tests/repositories/trigger/test_subscription_refresh_query.py b/api/tests/unit_tests/repositories/trigger/test_subscription_refresh_query.py new file mode 100644 index 00000000000000..e13a5b8df27cf6 --- /dev/null +++ b/api/tests/unit_tests/repositories/trigger/test_subscription_refresh_query.py @@ -0,0 +1,78 @@ +"""Due scans use real SQL and remain stable while workers renew earlier rows.""" + +from dataclasses import replace +from uuid import UUID + +import pytest +from sqlalchemy import Engine, update +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.pool import QueuePool + +from models.trigger import TriggerSubscription +from repositories.trigger.subscription_repository import TriggerSubscriptionRepository +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshRef +from tests.unit_tests.repositories.trigger.test_subscription_repository import _persist, _snapshot + + +@pytest.mark.parametrize( + ("credential_expiry", "subscription_expiry", "due"), + [ + (-1, -1, False), + (110, -1, True), + (111, -1, False), + (-1, 120, True), + (-1, 121, False), + (111, 120, True), + (0, -1, True), + (-1, 0, True), + ], +) +def test_due_thresholds_and_never_expiring_sentinels( + sqlite_session_factory: sessionmaker[Session], + sqlite_engine: Engine, + credential_expiry: int, + subscription_expiry: int, + due: bool, +) -> None: + snapshot = replace(_snapshot(), credential_expires_at=credential_expiry, expires_at=subscription_expiry) + _persist(sqlite_session_factory, snapshot) + repository = TriggerSubscriptionRepository(session_factory=sqlite_session_factory) + + result = repository.list_due_for_refresh(credential_due_at=110, subscription_due_at=120, after_id=None, limit=10) + + assert result == ( + (TriggerSubscriptionRefreshRef(tenant_id=snapshot.tenant_id, subscription_id=snapshot.id),) if due else () + ) + assert isinstance(sqlite_engine.pool, QueuePool) + assert sqlite_engine.pool.checkedout() == 0 + + +def test_cursor_keeps_later_candidates_when_earlier_rows_are_renewed( + sqlite_session_factory: sessionmaker[Session], +) -> None: + snapshots = [replace(_snapshot(), id=str(UUID(int=index)), credential_expires_at=10) for index in range(1, 6)] + _persist(sqlite_session_factory, *reversed(snapshots)) + repository = TriggerSubscriptionRepository(session_factory=sqlite_session_factory) + first = repository.list_due_for_refresh(credential_due_at=100, subscription_due_at=100, after_id=None, limit=2) + assert [item.subscription_id for item in first] == [snapshot.id for snapshot in snapshots[:2]] + with sqlite_session_factory.begin() as session: + session.execute( + update(TriggerSubscription) + .where(TriggerSubscription.id.in_([snapshot.id for snapshot in snapshots[:2]])) + .values(credential_expires_at=1000) + ) + + second = repository.list_due_for_refresh( + credential_due_at=100, subscription_due_at=100, after_id=first[-1].subscription_id, limit=2 + ) + final = repository.list_due_for_refresh( + credential_due_at=100, subscription_due_at=100, after_id=second[-1].subscription_id, limit=2 + ) + assert [item.subscription_id for item in (*first, *second, *final)] == [snapshot.id for snapshot in snapshots] + assert [item.tenant_id for item in (*first, *second, *final)] == [snapshot.tenant_id for snapshot in snapshots] + assert ( + repository.list_due_for_refresh( + credential_due_at=100, subscription_due_at=100, after_id=final[-1].subscription_id, limit=2 + ) + == () + ) diff --git a/api/tests/unit_tests/services/trigger/test_subscription_refresh_publisher.py b/api/tests/unit_tests/services/trigger/test_subscription_refresh_publisher.py new file mode 100644 index 00000000000000..594ddad4534c5d --- /dev/null +++ b/api/tests/unit_tests/services/trigger/test_subscription_refresh_publisher.py @@ -0,0 +1,68 @@ +"""The renewal publisher sends the real Celery task and its reservation owner.""" + +import pytest +from celery import Celery +from kombu.simple import SimpleQueue + +from extensions.ext_application_services import application_services +from services.trigger.subscription_gateway import enqueue_subscription_refresh +from services.trigger.subscription_refresh_service import TriggerSubscriptionRefreshJob, TriggerSubscriptionRefreshRef +from tasks.trigger_subscription_refresh_tasks import trigger_subscription_refresh +from tests.unit_tests.controllers.console.workspace import test_trigger_oauth_client +from tests.unit_tests.controllers.console.workspace.test_trigger_oauth_client import _Harness + +application_redis = test_trigger_oauth_client.application_redis +harness = test_trigger_oauth_client.harness + + +def test_publisher_preserves_legacy_worker_arguments_and_sends_reservation_in_headers() -> None: + previous_app = trigger_subscription_refresh.app + app = Celery("trigger-renewal-publisher-test", broker="memory://", set_as_current=False) + app.conf.update(task_always_eager=False, task_serializer="json") + app.set_current() + + # Exercise Celery's header/argument separation with the old consumer signature. + @app.task + def legacy_refresh(tenant_id: str, subscription_id: str) -> tuple[str, str]: + return tenant_id, subscription_id + + job = TriggerSubscriptionRefreshJob( + subscription=TriggerSubscriptionRefreshRef(tenant_id="tenant", subscription_id="subscription"), + token="reservation-owner", + ) + try: + with app.connection_for_write() as connection: + with SimpleQueue(connection.channel(), "trigger_refresh_executor") as queue: + queue.clear() + try: + enqueue_subscription_refresh(job) + message = queue.get(block=False) + assert message.headers["task"] == trigger_subscription_refresh.name + assert message.payload[1] == { + "tenant_id": "tenant", + "subscription_id": "subscription", + } + assert message.headers["reservation_token"] == "reservation-owner" + result = legacy_refresh.apply( + args=tuple(message.payload[0]), + kwargs=message.payload[1], + headers=message.headers, + throw=True, + ) + assert result.get() == ("tenant", "subscription") + message.ack() + assert queue.qsize() == 0 + finally: + queue.clear() + finally: + previous_app.set_current() + app.close() + + +@pytest.mark.parametrize("batch_size", [0, -1]) +def test_invalid_batch_size_stops_before_database_or_external_io(harness: _Harness, batch_size: int) -> None: + with harness.app.app_context(), pytest.raises(ValueError, match="batch size must be positive"): + application_services().workspaces.trigger_subscription_refresh.enqueue_due( + now=100, credential_threshold=10, subscription_threshold=20, batch_size=batch_size + ) + assert not harness.metadata.sessions