refactor(api): inject trigger subscription renewal scheduling

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

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(

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)

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,
+    )

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},
+    )