Added eviction policy configuration

Author: vladvildanov

redis/cache.py

@@ -1,19 +1,17 @@
-from typing import Callable, TypeVar, Any, NoReturn, List, Union
-from typing import Optional
+from abc import ABC, abstractmethod
+from typing import TypeVar
 from enum import Enum
 
-from cachetools import TTLCache, Cache, LRUCache
-from cachetools.keys import hashkey
-
-from redis.typing import ResponseT
+from cachetools import LRUCache, LFUCache, RRCache, Cache, TTLCache
 
 T = TypeVar('T')
 
 
 class EvictionPolicy(Enum):
-    LRU = "lru"
-    LFU = "lfu"
-    RANDOM = "random"
+    LRU = "LRU"
+    LFU = "LFU"
+    RANDOM = "RANDOM"
+    TTL = "TTL"
 
 
 class CacheConfiguration:
@@ -97,18 +95,61 @@ class CacheConfiguration:
     ]
 
     def __init__(self, **kwargs):
-        self._max_size = kwargs.get("cache_size", 10000)
-        self._ttl = kwargs.get("cache_ttl", 0)
-        self._eviction_policy = kwargs.get("eviction_policy", self.DEFAULT_EVICTION_POLICY)
+        self._max_size = kwargs.get("cache_size", None)
+        self._ttl = kwargs.get("cache_ttl", None)
+        self._eviction_policy = kwargs.get("cache_eviction", None)
+        if self._max_size is None:
+            self._max_size = 10000
+        if self._ttl is None:
+            self._ttl = 0
+        if self._eviction_policy is None:
+            self._eviction_policy = EvictionPolicy.LRU
+
+        if self._eviction_policy not in EvictionPolicy:
+            raise ValueError(f"Invalid eviction_policy {self._eviction_policy}")
 
     def get_ttl(self) -> int:
         return self._ttl
 
+    def get_max_size(self) -> int:
+        return self._max_size
+
     def get_eviction_policy(self) -> EvictionPolicy:
         return self._eviction_policy
 
     def is_exceeds_max_size(self, count: int) -> bool:
         return count > self._max_size
 
     def is_allowed_to_cache(self, command: str) -> bool:
-        return command in self.DEFAULT_ALLOW_LIST
+        return command in self.DEFAULT_ALLOW_LIST
+
+
+class CacheClass(Enum):
+    LRU = LRUCache
+    LFU = LFUCache
+    RANDOM = RRCache
+    TTL = TTLCache
+
+
+class CacheFactoryInterface(ABC):
+    @abstractmethod
+    def get_cache(self) -> Cache:
+        pass
+
+
+class CacheFactory(CacheFactoryInterface):
+    def __init__(self, conf: CacheConfiguration):
+        self._conf = conf
+
+    def get_cache(self) -> Cache:
+        eviction_policy = self._conf.get_eviction_policy()
+        cache_class = self._get_cache_class(eviction_policy).value
+
+        if eviction_policy == EvictionPolicy.TTL:
+            return cache_class(self._conf.get_max_size(), self._conf.get_ttl())
+
+        return cache_class(self._conf.get_max_size())
+
+    def _get_cache_class(self, eviction_policy: EvictionPolicy) -> CacheClass:
+        return CacheClass[eviction_policy.value]
+

redis/client.py

@@ -15,6 +15,7 @@
     _RedisCallbacksRESP3,
     bool_ok,
 )
+from redis.cache import EvictionPolicy
 from redis.commands import (
     CoreCommands,
     RedisModuleCommands,
@@ -216,6 +217,7 @@ def __init__(
         protocol: Optional[int] = 2,
         use_cache: bool = False,
         cache: Optional[Cache] = None,
+        cache_eviction: Optional[EvictionPolicy] = None,
         cache_size: int = 128,
         cache_ttl: int = 300,
     ) -> None:
@@ -315,6 +317,7 @@ def __init__(
                         {
                             "use_cache": use_cache,
                             "cache": cache,
+                            "cache_eviction": cache_eviction,
                             "cache_size": cache_size,
                             "cache_ttl": cache_ttl,
                         }

redis/cluster.py

@@ -11,6 +11,7 @@
 from redis._parsers import CommandsParser, Encoder
 from redis._parsers.helpers import parse_scan
 from redis.backoff import default_backoff
+from redis.cache import EvictionPolicy
 from redis.client import CaseInsensitiveDict, PubSub, Redis
 from redis.commands import READ_COMMANDS, RedisClusterCommands
 from redis.commands.helpers import list_or_args
@@ -505,6 +506,7 @@ def __init__(
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
         use_cache: bool = False,
         cache: Optional[Cache] = None,
+        cache_eviction: Optional[EvictionPolicy] = None,
         cache_size: int = 128,
         cache_ttl: int = 300,
         **kwargs,
@@ -648,6 +650,7 @@ def __init__(
             address_remap=address_remap,
             use_cache=use_cache,
             cache=cache,
+            cache_eviction=cache_eviction,
             cache_size=cache_size,
             cache_ttl=cache_ttl,
             **kwargs,
@@ -1331,6 +1334,7 @@ def __init__(
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
         use_cache: bool = False,
         cache: Optional[Cache] = None,
+        cache_eviction: Optional[EvictionPolicy] = None,
         cache_size: int = 128,
         cache_ttl: int = 300,
         **kwargs,
@@ -1347,6 +1351,7 @@ def __init__(
         self.address_remap = address_remap
         self.use_cache = use_cache
         self.cache = cache
+        self.cache_eviction = cache_eviction
         self.cache_size = cache_size
         self.cache_ttl = cache_ttl
         self._moved_exception = None
@@ -1494,6 +1499,7 @@ def create_redis_node(self, host, port, **kwargs):
             kwargs.update({"port": port})
             kwargs.update({"use_cache": self.use_cache})
             kwargs.update({"cache": self.cache})
+            kwargs.update({"cache_eviction": self.cache_eviction})
             kwargs.update({"cache_size": self.cache_size})
             kwargs.update({"cache_ttl": self.cache_ttl})
             r = Redis(connection_pool=self.connection_pool_class(**kwargs))
@@ -1503,6 +1509,7 @@ def create_redis_node(self, host, port, **kwargs):
                 port=port,
                 use_cache=self.use_cache,
                 cache=self.cache,
+                cache_eviction=self.cache_eviction,
                 cache_size=self.cache_size,
                 cache_ttl=self.cache_ttl,
                 **kwargs,
@@ -1682,9 +1689,6 @@ def remap_host_port(self, host: str, port: int) -> Tuple[str, int]:
             return self.address_remap((host, port))
         return host, port
 
-    def get_cache(self) -> Optional[Cache]:
-        return self.connection_pool.cache
-
 
 class ClusterPubSub(PubSub):
     """

redis/connection.py

@@ -15,7 +15,7 @@
 from apscheduler.schedulers.background import BackgroundScheduler
 from cachetools import TTLCache, Cache, LRUCache
 from cachetools.keys import hashkey
-from redis.cache import CacheConfiguration
+from redis.cache import CacheConfiguration, CacheFactory
 
 from ._parsers import Encoder, _HiredisParser, _RESP2Parser, _RESP3Parser
 from .backoff import NoBackoff
@@ -1246,6 +1246,7 @@ def __init__(
         self.max_connections = max_connections
         self.cache = None
         self._cache_conf = None
+        self._cache_factory = None
         self.cache_lock = None
         self.scheduler = None
 
@@ -1260,13 +1261,15 @@ def __init__(
             if cache is not None:
                 self.cache = cache
             else:
-                self.cache = TTLCache(self.connection_kwargs["cache_size"], self.connection_kwargs["cache_ttl"])
+                cache_factory = CacheFactory(self._cache_conf)
+                self.cache = cache_factory.get_cache()
 
             self.scheduler = BackgroundScheduler()
             self.scheduler.add_job(self._perform_health_check, "interval", seconds=2, id="cache_health_check")
             self.scheduler.start()
 
         connection_kwargs.pop("use_cache", None)
+        connection_kwargs.pop("cache_eviction", None)
         connection_kwargs.pop("cache_size", None)
         connection_kwargs.pop("cache_ttl", None)
         connection_kwargs.pop("cache", None)