admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: trigger billing (#28335)

Browse Source 
Signed-off-by: lyzno1 <yuanyouhuilyz@gmail.com>
Co-authored-by: lyzno1 <yuanyouhuilyz@gmail.com>
Co-authored-by: lyzno1 <92089059+lyzno1@users.noreply.github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Maries
2025-11-20 10:15:23 +08:00
committed by GitHub
parent c0b7ffd5d0
commit a1b735a4c0
61 changed files with 1475 additions and 465 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
api/services/workflow/queue_dispatcher.py
View File

@@ -2,16 +2,14 @@
Queue dispatcher system for async workflow execution.
Implements an ABC-based pattern for handling different subscription tiers
with appropriate queue routing and rate limiting.
with appropriate queue routing and priority assignment.
"""
from abc import ABC, abstractmethod
from enum import StrEnum
from configs import dify_config
from extensions.ext_redis import redis_client
from services.billing_service import BillingService
from services.workflow.rate_limiter import TenantDailyRateLimiter
class QueuePriority(StrEnum):
@@ -25,50 +23,16 @@ class QueuePriority(StrEnum):
class BaseQueueDispatcher(ABC):
"""Abstract base class for queue dispatchers"""
def __init__(self):
self.rate_limiter = TenantDailyRateLimiter(redis_client)
@abstractmethod
def get_queue_name(self) -> str:
"""Get the queue name for this dispatcher"""
pass
@abstractmethod
def get_daily_limit(self) -> int:
"""Get daily execution limit"""
pass
@abstractmethod
def get_priority(self) -> int:
"""Get task priority level"""
pass
def check_daily_quota(self, tenant_id: str) -> bool:
"""
Check if tenant has remaining daily quota
Args:
tenant_id: The tenant identifier
Returns:
True if quota available, False otherwise
"""
# Check without consuming
remaining = self.rate_limiter.get_remaining_quota(tenant_id=tenant_id, max_daily_limit=self.get_daily_limit())
return remaining > 0
def consume_quota(self, tenant_id: str) -> bool:
"""
Consume one execution from daily quota
Args:
tenant_id: The tenant identifier
Returns:
True if quota consumed successfully, False if limit reached
"""
return self.rate_limiter.check_and_consume(tenant_id=tenant_id, max_daily_limit=self.get_daily_limit())
class ProfessionalQueueDispatcher(BaseQueueDispatcher):
"""Dispatcher for professional tier"""
@@ -76,9 +40,6 @@ class ProfessionalQueueDispatcher(BaseQueueDispatcher):
def get_queue_name(self) -> str:
return QueuePriority.PROFESSIONAL
def get_daily_limit(self) -> int:
return int(1e9)
def get_priority(self) -> int:
return 100
@@ -89,9 +50,6 @@ class TeamQueueDispatcher(BaseQueueDispatcher):
def get_queue_name(self) -> str:
return QueuePriority.TEAM
def get_daily_limit(self) -> int:
return int(1e9)
def get_priority(self) -> int:
return 50
@@ -102,9 +60,6 @@ class SandboxQueueDispatcher(BaseQueueDispatcher):
def get_queue_name(self) -> str:
return QueuePriority.SANDBOX
def get_daily_limit(self) -> int:
return dify_config.APP_DAILY_RATE_LIMIT
def get_priority(self) -> int:
return 10

183
api/services/workflow/rate_limiter.py
View File

@@ -1,183 +0,0 @@
"""
Day-based rate limiter for workflow executions.
Implements UTC-based daily quotas that reset at midnight UTC for consistent rate limiting.
"""
from datetime import UTC, datetime, time, timedelta
from typing import Union
import pytz
from redis import Redis
from sqlalchemy import select
from extensions.ext_database import db
from extensions.ext_redis import RedisClientWrapper
from models.account import Account, TenantAccountJoin, TenantAccountRole
class TenantDailyRateLimiter:
"""
Day-based rate limiter that resets at midnight UTC
This class provides Redis-based rate limiting with the following features:
- Daily quotas that reset at midnight UTC for consistency
- Atomic check-and-consume operations
- Automatic cleanup of stale counters
- Timezone-aware error messages for better UX
"""
def __init__(self, redis_client: Union[Redis, RedisClientWrapper]):
self.redis = redis_client
def get_tenant_owner_timezone(self, tenant_id: str) -> str:
"""
Get timezone of tenant owner
Args:
tenant_id: The tenant identifier
Returns:
Timezone string (e.g., 'America/New_York', 'UTC')
"""
# Query to get tenant owner's timezone using scalar and select
owner = db.session.scalar(
select(Account)
.join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id)
.where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.role == TenantAccountRole.OWNER)
)
if not owner:
return "UTC"
return owner.timezone or "UTC"
def _get_day_key(self, tenant_id: str) -> str:
"""
Get Redis key for current UTC day
Args:
tenant_id: The tenant identifier
Returns:
Redis key for the current UTC day
"""
utc_now = datetime.now(UTC)
date_str = utc_now.strftime("%Y-%m-%d")
return f"workflow:daily_limit:{tenant_id}:{date_str}"
def _get_ttl_seconds(self) -> int:
"""
Calculate seconds until UTC midnight
Returns:
Number of seconds until UTC midnight
"""
utc_now = datetime.now(UTC)
# Get next midnight in UTC
next_midnight = datetime.combine(utc_now.date() + timedelta(days=1), time.min)
next_midnight = next_midnight.replace(tzinfo=UTC)
return int((next_midnight - utc_now).total_seconds())
def check_and_consume(self, tenant_id: str, max_daily_limit: int) -> bool:
"""
Check if quota available and consume one execution
Args:
tenant_id: The tenant identifier
max_daily_limit: Maximum daily limit
Returns:
True if quota consumed successfully, False if limit reached
"""
key = self._get_day_key(tenant_id)
ttl = self._get_ttl_seconds()
# Check current usage
current = self.redis.get(key)
if current is None:
# First execution of the day - set to 1
self.redis.setex(key, ttl, 1)
return True
current_count = int(current)
if current_count < max_daily_limit:
# Within limit, increment
new_count = self.redis.incr(key)
# Update TTL
self.redis.expire(key, ttl)
# Double-check in case of race condition
if new_count <= max_daily_limit:
return True
else:
# Race condition occurred, decrement back
self.redis.decr(key)
return False
else:
# Limit exceeded
return False
def get_remaining_quota(self, tenant_id: str, max_daily_limit: int) -> int:
"""
Get remaining quota for the day
Args:
tenant_id: The tenant identifier
max_daily_limit: Maximum daily limit
Returns:
Number of remaining executions for the day
"""
key = self._get_day_key(tenant_id)
used = int(self.redis.get(key) or 0)
return max(0, max_daily_limit - used)
def get_current_usage(self, tenant_id: str) -> int:
"""
Get current usage for the day
Args:
tenant_id: The tenant identifier
Returns:
Number of executions used today
"""
key = self._get_day_key(tenant_id)
return int(self.redis.get(key) or 0)
def reset_quota(self, tenant_id: str) -> bool:
"""
Reset quota for testing purposes
Args:
tenant_id: The tenant identifier
Returns:
True if key was deleted, False if key didn't exist
"""
key = self._get_day_key(tenant_id)
return bool(self.redis.delete(key))
def get_quota_reset_time(self, tenant_id: str, timezone_str: str) -> datetime:
"""
Get the time when quota will reset (next UTC midnight in tenant's timezone)
Args:
tenant_id: The tenant identifier
timezone_str: Tenant's timezone for display purposes
Returns:
Datetime when quota resets (next UTC midnight in tenant's timezone)
"""
tz = pytz.timezone(timezone_str)
utc_now = datetime.now(UTC)
# Get next midnight in UTC, then convert to tenant's timezone
next_utc_midnight = datetime.combine(utc_now.date() + timedelta(days=1), time.min)
next_utc_midnight = pytz.UTC.localize(next_utc_midnight)
return next_utc_midnight.astimezone(tz)