From e53edb0fc2cdac12ea32d5e06e8e5c89e39e41cf Mon Sep 17 00:00:00 2001 From: Yeuoly <45712896+Yeuoly@users.noreply.github.com> Date: Wed, 27 Aug 2025 17:35:04 +0800 Subject: [PATCH] refactor: optimize TenantDailyRateLimiter to use UTC internally with timezone-aware error messages (#24632) Co-authored-by: Claude --- api/services/async_workflow_service.py | 11 +-- api/services/workflow/queue_dispatcher.py | 10 +- api/services/workflow/rate_limiter.py | 106 +++++++++------------- 3 files changed, 50 insertions(+), 77 deletions(-) diff --git a/api/services/async_workflow_service.py b/api/services/async_workflow_service.py index 3b20ec500a..58b75fc07f 100644 --- a/api/services/async_workflow_service.py +++ b/api/services/async_workflow_service.py @@ -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) diff --git a/api/services/workflow/queue_dispatcher.py b/api/services/workflow/queue_dispatcher.py index f1e5db9073..a3d71bbdd8 100644 --- a/api/services/workflow/queue_dispatcher.py +++ b/api/services/workflow/queue_dispatcher.py @@ -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() ) diff --git a/api/services/workflow/rate_limiter.py b/api/services/workflow/rate_limiter.py index 49f936d253..6f83ed8ad4 100644 --- a/api/services/workflow/rate_limiter.py +++ b/api/services/workflow/rate_limiter.py @@ -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)