From c9d7e3b5cb488528c4ed400471e4489aca8ed5c0 Mon Sep 17 00:00:00 2001 From: chenjh3 Date: Thu, 13 Feb 2025 13:36:22 +0800 Subject: [PATCH 1/7] feat: When using redis cluster for caching, add global_prefix to all keys cached by redis --- api/.env.example | 1 + api/configs/middleware/cache/redis_config.py | 5 ++ .../app/features/rate_limiting/rate_limit.py | 6 ++- api/extensions/ext_redis.py | 52 +++++++++++++++++-- 4 files changed, 57 insertions(+), 7 deletions(-) diff --git a/api/.env.example b/api/.env.example index 95da531a1d..41744504d7 100644 --- a/api/.env.example +++ b/api/.env.example @@ -36,6 +36,7 @@ REDIS_USERNAME= REDIS_PASSWORD=difyai123456 REDIS_USE_SSL=false REDIS_DB=0 +REDIS_KEY_PREFIX= # redis Sentinel configuration. REDIS_USE_SENTINEL=false diff --git a/api/configs/middleware/cache/redis_config.py b/api/configs/middleware/cache/redis_config.py index 2e98c31ec3..55f436f8bf 100644 --- a/api/configs/middleware/cache/redis_config.py +++ b/api/configs/middleware/cache/redis_config.py @@ -34,6 +34,11 @@ class RedisConfig(BaseSettings): default=0, ) + REDIS_KEY_PREFIX: Optional[str] = Field( + description="Redis global string key prefix (if required)", + default=None, + ) + REDIS_USE_SSL: bool = Field( description="Enable SSL/TLS for the Redis connection", default=False, diff --git a/api/core/app/features/rate_limiting/rate_limit.py b/api/core/app/features/rate_limiting/rate_limit.py index dcc2b4e55f..df63db63cc 100644 --- a/api/core/app/features/rate_limiting/rate_limit.py +++ b/api/core/app/features/rate_limiting/rate_limit.py @@ -2,6 +2,7 @@ import logging import time import uuid from collections.abc import Generator, Mapping +from configs import dify_config from datetime import timedelta from typing import Any, Optional, Union @@ -12,8 +13,9 @@ logger = logging.getLogger(__name__) class RateLimit: - _MAX_ACTIVE_REQUESTS_KEY = "dify:rate_limit:{}:max_active_requests" - _ACTIVE_REQUESTS_KEY = "dify:rate_limit:{}:active_requests" + _KEY_PREFIX = dify_config.REDIS_KEY_PREFIX if dify_config.REDIS_KEY_PREFIX is not None else "dify" + _MAX_ACTIVE_REQUESTS_KEY = _KEY_PREFIX + ":rate_limit:{}:max_active_requests" + _ACTIVE_REQUESTS_KEY = _KEY_PREFIX + ":rate_limit:{}:active_requests" _UNLIMITED_REQUEST_ID = "unlimited_request_id" _REQUEST_MAX_ALIVE_TIME = 10 * 60 # 10 minutes _ACTIVE_REQUESTS_COUNT_FLUSH_INTERVAL = 5 * 60 # recalculate request_count from request_detail every 5 minutes diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index da41805707..f8d4c21826 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -1,4 +1,5 @@ -from typing import Any, Union +from typing import Any, Union, Callable, TypeVar, Generic +import functools import redis from redis.cluster import ClusterNode, RedisCluster @@ -8,6 +9,42 @@ from redis.sentinel import Sentinel from configs import dify_config from dify_app import DifyApp +T = TypeVar('T') + + +class KeyPrefixMethodProxy(Generic[T]): + """ + KeyPrefixMethodProxy is a generic class used to proxy method calls. + It can preprocess parameters before calling a method, such as prefixing specific keys. + Key features include: + If a method's argument contains'key', prefix it. + If the method's positional parameter contains a key of type string, prefix it. + """ + def __init__(self, client: T, prefix: str): + self._client = client + self._prefix = prefix + + def __getattr__(self, item): + attr = getattr(self._client, item) + if callable(attr): + return self._wrap_method(attr) + return attr + + def _wrap_method(self, method: Callable) -> Callable: + @functools.wraps(method) + def wrapper(*args, **kwargs): + if 'key' in kwargs: + if isinstance(kwargs['key'], str): + kwargs['key'] = self._add_prefix(kwargs['key']) + elif args: + if args and isinstance(args[0], str): + args = (self._add_prefix(args[0]),) + args[1:] + return method(*args, **kwargs) + return wrapper + + def _add_prefix(self, key: str) -> str: + return f"{self._prefix}:{key}" if self._prefix else key + class RedisClientWrapper: """ @@ -32,14 +69,19 @@ class RedisClientWrapper: def __init__(self): self._client = None + self._prefix = None - def initialize(self, client): + def initialize(self, client, prefix=None): if self._client is None: self._client = client + if prefix is not None: + self._prefix = prefix def __getattr__(self, item): if self._client is None: raise RuntimeError("Redis client is not initialized. Call init_app first.") + if self._prefix is not None: + return getattr(KeyPrefixMethodProxy(self._client, self._prefix), item) return getattr(self._client, item) @@ -75,7 +117,7 @@ def init_app(app: DifyApp): }, ) master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params) - redis_client.initialize(master) + redis_client.initialize(master, prefix=dify_config.REDIS_KEY_PREFIX) elif dify_config.REDIS_USE_CLUSTERS: assert dify_config.REDIS_CLUSTERS is not None, "REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True" nodes = [ @@ -83,7 +125,7 @@ def init_app(app: DifyApp): for node in dify_config.REDIS_CLUSTERS.split(",") ] # FIXME: mypy error here, try to figure out how to fix it - redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD)) # type: ignore + redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), prefix=dify_config.REDIS_KEY_PREFIX) # type: ignore else: redis_params.update( { @@ -93,6 +135,6 @@ def init_app(app: DifyApp): } ) pool = redis.ConnectionPool(**redis_params) - redis_client.initialize(redis.Redis(connection_pool=pool)) + redis_client.initialize(redis.Redis(connection_pool=pool), prefix=dify_config.REDIS_KEY_PREFIX) app.extensions["redis"] = redis_client From 5fb7beae13b56fbb9be7ef08fb7e168ee3c8c985 Mon Sep 17 00:00:00 2001 From: chenjh3 Date: Thu, 13 Feb 2025 14:00:06 +0800 Subject: [PATCH 2/7] fix hint --- api/core/app/features/rate_limiting/rate_limit.py | 2 +- api/extensions/ext_redis.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/api/core/app/features/rate_limiting/rate_limit.py b/api/core/app/features/rate_limiting/rate_limit.py index df63db63cc..f53fa4506f 100644 --- a/api/core/app/features/rate_limiting/rate_limit.py +++ b/api/core/app/features/rate_limiting/rate_limit.py @@ -2,10 +2,10 @@ import logging import time import uuid from collections.abc import Generator, Mapping -from configs import dify_config from datetime import timedelta from typing import Any, Optional, Union +from configs import dify_config from core.errors.error import AppInvokeQuotaExceededError from extensions.ext_redis import redis_client diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index f8d4c21826..89342f6c33 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -1,5 +1,5 @@ -from typing import Any, Union, Callable, TypeVar, Generic import functools +from typing import Any, Callable, Generic, TypeVar, Union import redis from redis.cluster import ClusterNode, RedisCluster @@ -125,7 +125,8 @@ def init_app(app: DifyApp): for node in dify_config.REDIS_CLUSTERS.split(",") ] # FIXME: mypy error here, try to figure out how to fix it - redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), prefix=dify_config.REDIS_KEY_PREFIX) # type: ignore + redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), + prefix=dify_config.REDIS_KEY_PREFIX) # type: ignore else: redis_params.update( { From 25ec569a851d82b7c9cb5555f5c890bd64d009c0 Mon Sep 17 00:00:00 2001 From: chenjh3 Date: Thu, 13 Feb 2025 14:23:34 +0800 Subject: [PATCH 3/7] fix hint --- docker-legacy/docker-compose.yaml | 2 ++ docker/docker-compose.yaml | 1 + 2 files changed, 3 insertions(+) diff --git a/docker-legacy/docker-compose.yaml b/docker-legacy/docker-compose.yaml index 0a071e80b3..95c3b6c0dc 100644 --- a/docker-legacy/docker-compose.yaml +++ b/docker-legacy/docker-compose.yaml @@ -59,6 +59,7 @@ services: REDIS_USE_SSL: 'false' # use redis db 0 for redis cache REDIS_DB: 0 + REDIS_KEY_PREFIX: '' # The configurations of celery broker. # Use redis as the broker, and redis db 1 for celery broker. CELERY_BROKER_URL: redis://:difyai123456@redis:6379/1 @@ -254,6 +255,7 @@ services: REDIS_USERNAME: '' REDIS_PASSWORD: difyai123456 REDIS_DB: 0 + REDIS_KEY_PREFIX: '' REDIS_USE_SSL: 'false' # The configurations of celery broker. CELERY_BROKER_URL: redis://:difyai123456@redis:6379/1 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 38e4783806..a3db2d733c 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -62,6 +62,7 @@ x-shared-env: &shared-api-worker-env REDIS_PASSWORD: ${REDIS_PASSWORD:-difyai123456} REDIS_USE_SSL: ${REDIS_USE_SSL:-false} REDIS_DB: ${REDIS_DB:-0} + REDIS_KEY_PREFIX: ${REDIS_KEY_PREFIX:-} REDIS_USE_SENTINEL: ${REDIS_USE_SENTINEL:-false} REDIS_SENTINELS: ${REDIS_SENTINELS:-} REDIS_SENTINEL_SERVICE_NAME: ${REDIS_SENTINEL_SERVICE_NAME:-} From ee29dbea5718a102fbfc0c0d1ce92e05acde253f Mon Sep 17 00:00:00 2001 From: chenjh3 Date: Thu, 13 Feb 2025 14:23:48 +0800 Subject: [PATCH 4/7] fix hint --- docker/.env.example | 1 + 1 file changed, 1 insertion(+) diff --git a/docker/.env.example b/docker/.env.example index 3bc79059dc..5d84b176ca 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -235,6 +235,7 @@ REDIS_USERNAME= REDIS_PASSWORD=difyai123456 REDIS_USE_SSL=false REDIS_DB=0 +REDIS_KEY_PREFIX= # Whether to use Redis Sentinel mode. # If set to true, the application will automatically discover and connect to the master node through Sentinel. From e4da9823f8f25033318990af161e6e746ae50181 Mon Sep 17 00:00:00 2001 From: chenjh3 Date: Thu, 13 Feb 2025 14:52:39 +0800 Subject: [PATCH 5/7] fix hint --- api/extensions/ext_redis.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index 89342f6c33..dedafddeed 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -9,7 +9,7 @@ from redis.sentinel import Sentinel from configs import dify_config from dify_app import DifyApp -T = TypeVar('T') +T = TypeVar("T") class KeyPrefixMethodProxy(Generic[T]): @@ -20,6 +20,7 @@ class KeyPrefixMethodProxy(Generic[T]): If a method's argument contains'key', prefix it. If the method's positional parameter contains a key of type string, prefix it. """ + def __init__(self, client: T, prefix: str): self._client = client self._prefix = prefix @@ -33,13 +34,14 @@ class KeyPrefixMethodProxy(Generic[T]): def _wrap_method(self, method: Callable) -> Callable: @functools.wraps(method) def wrapper(*args, **kwargs): - if 'key' in kwargs: - if isinstance(kwargs['key'], str): - kwargs['key'] = self._add_prefix(kwargs['key']) + if "key" in kwargs: + if isinstance(kwargs["key"], str): + kwargs["key"] = self._add_prefix(kwargs["key"]) elif args: if args and isinstance(args[0], str): args = (self._add_prefix(args[0]),) + args[1:] return method(*args, **kwargs) + return wrapper def _add_prefix(self, key: str) -> str: @@ -125,8 +127,10 @@ def init_app(app: DifyApp): for node in dify_config.REDIS_CLUSTERS.split(",") ] # FIXME: mypy error here, try to figure out how to fix it - redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), - prefix=dify_config.REDIS_KEY_PREFIX) # type: ignore + redis_client.initialize( + RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), + prefix=dify_config.REDIS_KEY_PREFIX, + ) # type: ignore else: redis_params.update( { From efe0ed5d523da12b0306156caf5542facf411d31 Mon Sep 17 00:00:00 2001 From: chenjh3 Date: Tue, 18 Feb 2025 10:48:57 +0800 Subject: [PATCH 6/7] fix hint --- api/extensions/ext_redis.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index dedafddeed..c62bf3195f 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -128,9 +128,9 @@ def init_app(app: DifyApp): ] # FIXME: mypy error here, try to figure out how to fix it redis_client.initialize( - RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), + RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD), # type: ignore prefix=dify_config.REDIS_KEY_PREFIX, - ) # type: ignore + ) else: redis_params.update( { From b4c26c86c37b4f11a5ac25e66f36d0d6591bfa0d Mon Sep 17 00:00:00 2001 From: jiandanfeng Date: Tue, 18 Feb 2025 03:44:10 +0000 Subject: [PATCH 7/7] fix conflicts --- docker-legacy/docker-compose.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/docker-legacy/docker-compose.yaml b/docker-legacy/docker-compose.yaml index 95c3b6c0dc..0a071e80b3 100644 --- a/docker-legacy/docker-compose.yaml +++ b/docker-legacy/docker-compose.yaml @@ -59,7 +59,6 @@ services: REDIS_USE_SSL: 'false' # use redis db 0 for redis cache REDIS_DB: 0 - REDIS_KEY_PREFIX: '' # The configurations of celery broker. # Use redis as the broker, and redis db 1 for celery broker. CELERY_BROKER_URL: redis://:difyai123456@redis:6379/1 @@ -255,7 +254,6 @@ services: REDIS_USERNAME: '' REDIS_PASSWORD: difyai123456 REDIS_DB: 0 - REDIS_KEY_PREFIX: '' REDIS_USE_SSL: 'false' # The configurations of celery broker. CELERY_BROKER_URL: redis://:difyai123456@redis:6379/1