Improved locking and logging when cleaning old cached sessions.

This commit is contained in:
Marc-Andre Ferland 2022-10-15 04:39:45 -04:00
parent d3b28c42e6
commit 8fdb1e7ec9

View File

@ -70,7 +70,7 @@ class ImageRequest(BaseModel):
class TaskCache(): class TaskCache():
def __init__(self): def __init__(self):
self._base = dict() self._base = dict()
self._lock: threading.Lock = threading.Lock() self._lock: threading.Lock = threading.RLock()
def _get_ttl_time(self, ttl: int) -> int: def _get_ttl_time(self, ttl: int) -> int:
return int(time.time()) + ttl return int(time.time()) + ttl
def _is_expired(self, timestamp: int) -> bool: def _is_expired(self, timestamp: int) -> bool:
@ -78,10 +78,16 @@ class TaskCache():
def clean(self) -> None: def clean(self) -> None:
self._lock.acquire() self._lock.acquire()
try: try:
# Create a list of expired keys to delete
to_delete = []
for key in self._base: for key in self._base:
ttl, _ = self._base[key] ttl, _ = self._base[key]
if self._is_expired(ttl): if self._is_expired(ttl):
to_delete.append(key)
# Remove Items
for key in to_delete:
del self._base[key] del self._base[key]
print(f'Session {key} expired. Data removed.')
finally: finally:
self._lock.release() self._lock.release()
def clear(self) -> None: def clear(self) -> None:
@ -126,6 +132,7 @@ class TaskCache():
try: try:
ttl, value = self._base.get(key, (None, None)) ttl, value = self._base.get(key, (None, None))
if ttl is not None and self._is_expired(ttl): if ttl is not None and self._is_expired(ttl):
print(f'Session {key} expired. Discarding data.')
self.delete(key) self.delete(key)
return None return None
return value return value