enable generic key cache stacking

dnstapir · Oct 17, 2024 · e5f680e · e5f680e
1 parent 40c5888
commit e5f680e
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 22 deletions.
diff --git a/aggrec/key_cache.py b/aggrec/key_cache.py
@@ -31,7 +31,7 @@ class MemoryKeyCache(KeyCache):
     def __init__(self, size: int, ttl: int):
         super().__init__()
         self.cache = ExpiringDict(max_len=size, max_age_seconds=ttl)
-        self.logger.info("Using memory cache size=%d ttl=%d", size, ttl)
+        self.logger.info("Configured memory key cache size=%d ttl=%d", size, ttl)
 
     def get(self, key: str) -> bytes | None:
         res = self.cache.get(key)
@@ -48,7 +48,7 @@ def __init__(self, redis_client: redis.Redis, ttl: int):
         super().__init__()
         self.redis_client = redis_client
         self.ttl = ttl
-        self.logger.info("Using Redis cache ttl=%d", ttl)
+        self.logger.info("Configured Redis key cache ttl=%d", ttl)
 
     def get(self, key: str) -> bytes | None:
         res = self.redis_client.get(name=key)
@@ -61,16 +61,16 @@ def set(self, key: str, value: bytes) -> None:
         self.redis_client.set(name=key, value=value, exat=expires_at)
 
 
-class MemoryRedisKeyCache(KeyCache):
-    def __init__(self, size: int, ttl: int, redis_client: redis.Redis):
-        super().__init__()
-        self.memory_cache = MemoryKeyCache(size, ttl)
-        self.redis_cache = RedisKeyCache(redis_client, ttl)
-        self.logger.info("Using memory+redis cache size=%d ttl=%d", size, ttl)
+class CombinedKeyCache(KeyCache):
+    def __init__(self, caches: list[KeyCache]):
+        self.caches = caches
 
     def get(self, key: str) -> bytes | None:
-        return self.memory_cache.get(key) or self.redis_cache.get(key)
+        for cache in self.caches:
+            if res := cache.get(key):
+                return res
+        return None
 
     def set(self, key: str, value: bytes) -> None:
-        self.memory_cache.set(key, value)
-        self.redis_cache.set(key, value)
+        for cache in self.caches:
+            cache.set(key, value)
diff --git a/aggrec/server.py b/aggrec/server.py
@@ -14,7 +14,7 @@
 import aggrec.extras
 
 from . import OPENAPI_METADATA, __verbose_version__
-from .key_cache import KeyCache, MemoryKeyCache, MemoryRedisKeyCache, RedisKeyCache
+from .key_cache import CombinedKeyCache, KeyCache, MemoryKeyCache, RedisKeyCache
 from .logging import JsonFormatter  # noqa
 from .settings import Settings
 from .telemetry import configure_opentelemetry
@@ -72,19 +72,17 @@ def __init__(self, settings: Settings):
         )
         self.key_cache: KeyCache | None = None
         if self.settings.key_cache:
+            memory_key_cache = MemoryKeyCache(size=self.settings.key_cache.size, ttl=self.settings.key_cache.ttl)
             if redis_settings := self.settings.key_cache.redis:
                 redis_client = redis.StrictRedis(host=redis_settings.host, port=redis_settings.port)
                 self.logger.debug("Using REDIS at %s:%d", redis_settings.host, redis_settings.port)
+                redis_key_cache = RedisKeyCache(redis_client=redis_client, ttl=self.settings.key_cache.ttl)
                 if self.settings.key_cache.size:
-                    self.key_cache = MemoryRedisKeyCache(
-                        redis_client=redis_client,
-                        ttl=self.settings.key_cache.ttl,
-                        size=self.settings.key_cache.size,
-                    )
+                    self.key_cache = CombinedKeyCache([memory_key_cache, redis_key_cache])
                 else:
-                    self.key_cache = RedisKeyCache(redis_client=redis_client, ttl=self.settings.key_cache.ttl)
+                    self.key_cache = redis_key_cache
             elif self.settings.key_cache.size:
-                self.key_cache = MemoryKeyCache(size=self.settings.key_cache.size, ttl=self.settings.key_cache.ttl)
+                self.key_cache = memory_key_cache
 
     @staticmethod
     def connect_mongodb(settings: Settings):

diff --git a/tests/test_key_cache.py b/tests/test_key_cache.py
@@ -2,7 +2,7 @@
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import ed25519
 
-from aggrec.key_cache import KeyCache, MemoryKeyCache, MemoryRedisKeyCache, RedisKeyCache
+from aggrec.key_cache import CombinedKeyCache, KeyCache, MemoryKeyCache, RedisKeyCache
 
 
 def _test_key_cache(key_cache: KeyCache):
@@ -32,7 +32,10 @@ def test_memory_cache():
     _test_key_cache(key_cache=key_cache)
 
 
-def test_memory_redis_cache():
+def test_memory_stack():
     redis_client = fakeredis.FakeRedis()
-    key_cache = MemoryRedisKeyCache(size=100, ttl=60, redis_client=redis_client)
+    memory_key_cache = MemoryKeyCache(size=100, ttl=60)
+    redis_key_cache = RedisKeyCache(redis_client=redis_client, ttl=60)
+
+    key_cache = CombinedKeyCache([memory_key_cache, redis_key_cache])
     _test_key_cache(key_cache=key_cache)