opensource
/
dify
mirror of https://github.com/langgenius/dify.git
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor: optimize TenantDailyRateLimiter to use UTC internally with timezone-aware error messages (#24632) 

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Yeuoly 2025-08-27 17:35:04 +08:00 committed by GitHub
parent 17908fbf6b
commit e53edb0fc2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 50 additions and 77 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
api/services/async_workflow_service.py
View File

@ -94,8 +94,7 @@ class AsyncWorkflowService:
# 3. Get dispatcher based on tenant subscription
dispatcher = dispatcher_manager.get_dispatcher(trigger_data.tenant_id)
# 4. Get tenant owner timezone for rate limiting
tenant_owner_tz = rate_limiter._get_tenant_owner_timezone(trigger_data.tenant_id)
# 4. Rate limiting check will be done without timezone first
# 5. Determine user role and ID
if isinstance(user, Account):
@ -125,16 +124,16 @@ class AsyncWorkflowService:
session.commit()
# 7. Check and consume daily quota
if not dispatcher.consume_quota(trigger_data.tenant_id, tenant_owner_tz):
if not dispatcher.consume_quota(trigger_data.tenant_id):
# Update trigger log status
trigger_log.status = WorkflowTriggerStatus.RATE_LIMITED
trigger_log.error = f"Daily limit reached for {dispatcher.get_queue_name()}"
trigger_log_repo.update(trigger_log)
session.commit()
remaining = rate_limiter.get_remaining_quota(
trigger_data.tenant_id, dispatcher.get_daily_limit(), tenant_owner_tz
)
tenant_owner_tz = rate_limiter._get_tenant_owner_timezone(trigger_data.tenant_id)
remaining = rate_limiter.get_remaining_quota(trigger_data.tenant_id, dispatcher.get_daily_limit())
reset_time = rate_limiter.get_quota_reset_time(trigger_data.tenant_id, tenant_owner_tz)

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

@ -44,36 +44,34 @@ class BaseQueueDispatcher(ABC):
"""Get task priority level"""
pass
def check_daily_quota(self, tenant_id: str, tenant_owner_tz: str) -> bool:
def check_daily_quota(self, tenant_id: str) -> bool:
"""
Check if tenant has remaining daily quota
Args:
tenant_id: The tenant identifier
tenant_owner_tz: Tenant owner's timezone
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(), timezone_str=tenant_owner_tz
tenant_id=tenant_id, max_daily_limit=self.get_daily_limit()
)
return remaining > 0
def consume_quota(self, tenant_id: str, tenant_owner_tz: str) -> bool:
def consume_quota(self, tenant_id: str) -> bool:
"""
Consume one execution from daily quota
Args:
tenant_id: The tenant identifier
tenant_owner_tz: Tenant owner's timezone
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(), timezone_str=tenant_owner_tz
tenant_id=tenant_id, max_daily_limit=self.get_daily_limit()
)

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

@ -1,11 +1,11 @@
"""
Day-based rate limiter for workflow executions.
Implements timezone-aware daily quotas that reset at midnight in the tenant owner's timezone.
Implements UTC-based daily quotas that reset at midnight UTC for consistent rate limiting.
"""
from datetime import datetime, time, timedelta
from typing import Optional, Union
from typing import Union
import pytz
from redis import Redis
@ -18,13 +18,13 @@ from models.account import Account, TenantAccountJoin, TenantAccountRole
class TenantDailyRateLimiter:
"""
Day-based rate limiter that resets at midnight in tenant owner's timezone
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 in tenant owner's timezone
- Daily quotas that reset at midnight UTC for consistency
- Atomic check-and-consume operations
- Automatic cleanup of stale counters
- Support for timezone changes without duplicate limits
- Timezone-aware error messages for better UX
"""
def __init__(self, redis_client: Union[Redis, RedisClientWrapper]):
@ -52,57 +52,47 @@ class TenantDailyRateLimiter:
return owner.timezone or "UTC"
def _get_day_key(self, tenant_id: str, timezone_str: str) -> str:
def _get_day_key(self, tenant_id: str) -> str:
"""
Get Redis key for current day in tenant's timezone
Get Redis key for current UTC day
Args:
tenant_id: The tenant identifier
timezone_str: Timezone string
Returns:
Redis key for the current day
Redis key for the current UTC day
"""
tz = pytz.timezone(timezone_str)
now = datetime.now(tz)
date_str = now.strftime("%Y-%m-%d")
return f"workflow:daily_limit:{tenant_id}:{date_str}:{timezone_str}"
utc_now = datetime.utcnow()
date_str = utc_now.strftime("%Y-%m-%d")
return f"workflow:daily_limit:{tenant_id}:{date_str}"
def _get_ttl_seconds(self, timezone_str: str) -> int:
def _get_ttl_seconds(self) -> int:
"""
Calculate seconds until midnight in given timezone
Args:
timezone_str: Timezone string
Calculate seconds until UTC midnight
Returns:
Number of seconds until midnight
Number of seconds until UTC midnight
"""
tz = pytz.timezone(timezone_str)
now = datetime.now(tz)
utc_now = datetime.utcnow()
# Get next midnight in UTC
next_midnight = datetime.combine(utc_now.date() + timedelta(days=1), time.min)
return int((next_midnight - utc_now).total_seconds())
# Get next midnight in the timezone
midnight = tz.localize(datetime.combine(now.date() + timedelta(days=1), time.min))
return int((midnight - now).total_seconds())
def check_and_consume(self, tenant_id: str, max_daily_limit: int, timezone_str: Optional[str] = None) -> bool:
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
timezone_str: Optional timezone string (will be fetched if not provided)
Returns:
True if quota consumed successfully, False if limit reached
"""
if not timezone_str:
timezone_str = self._get_tenant_owner_timezone(tenant_id)
key = self._get_day_key(tenant_id, timezone_str)
ttl = self._get_ttl_seconds(timezone_str)
key = self._get_day_key(tenant_id)
ttl = self._get_ttl_seconds()
# Check current usage
current = self.redis.get(key)
@ -116,7 +106,7 @@ class TenantDailyRateLimiter:
if current_count < max_daily_limit:
# Within limit, increment
new_count = self.redis.incr(key)
# Update TTL in case timezone changed
# Update TTL
self.redis.expire(key, ttl)
# Double-check in case of race condition
@ -130,77 +120,63 @@ class TenantDailyRateLimiter:
# Limit exceeded
return False
def get_remaining_quota(self, tenant_id: str, max_daily_limit: int, timezone_str: Optional[str] = None) -> int:
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
timezone_str: Optional timezone string (will be fetched if not provided)
Returns:
Number of remaining executions for the day
"""
if not timezone_str:
timezone_str = self._get_tenant_owner_timezone(tenant_id)
key = self._get_day_key(tenant_id, timezone_str)
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, timezone_str: Optional[str] = None) -> int:
def get_current_usage(self, tenant_id: str) -> int:
"""
Get current usage for the day
Args:
tenant_id: The tenant identifier
timezone_str: Optional timezone string (will be fetched if not provided)
Returns:
Number of executions used today
"""
if not timezone_str:
timezone_str = self._get_tenant_owner_timezone(tenant_id)
key = self._get_day_key(tenant_id, timezone_str)
key = self._get_day_key(tenant_id)
return int(self.redis.get(key) or 0)
def reset_quota(self, tenant_id: str, timezone_str: Optional[str] = None) -> bool:
def reset_quota(self, tenant_id: str) -> bool:
"""
Reset quota for testing purposes
Args:
tenant_id: The tenant identifier
timezone_str: Optional timezone string (will be fetched if not provided)
Returns:
True if key was deleted, False if key didn't exist
"""
if not timezone_str:
timezone_str = self._get_tenant_owner_timezone(tenant_id)
key = self._get_day_key(tenant_id, timezone_str)
key = self._get_day_key(tenant_id)
return bool(self.redis.delete(key))
def get_quota_reset_time(self, tenant_id: str, timezone_str: Optional[str] = None) -> datetime:
def get_quota_reset_time(self, tenant_id: str, timezone_str: str) -> datetime:
"""
Get the time when quota will reset (midnight in tenant's timezone)
Get the time when quota will reset (next UTC midnight in tenant's timezone)
Args:
tenant_id: The tenant identifier
timezone_str: Optional timezone string (will be fetched if not provided)
timezone_str: Tenant's timezone for display purposes
Returns:
Datetime when quota resets
Datetime when quota resets (next UTC midnight in tenant's timezone)
"""
if not timezone_str:
timezone_str = self._get_tenant_owner_timezone(tenant_id)
tz = pytz.timezone(timezone_str)
now = datetime.now(tz)
# Get next midnight in the timezone
midnight = tz.localize(datetime.combine(now.date() + timedelta(days=1), time.min))
return midnight
utc_now = datetime.utcnow()
# 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)