From 632eef48b2dbf9a3588bf0bcfc263f828c2528a8 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 17 Aug 2023 17:59:00 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(manager.py):=20change=20thre?=
 =?UTF-8?q?ading.Lock=20to=20threading.RLock=20to=20allow=20reentrant=20lo?=
 =?UTF-8?q?cking=20and=20prevent=20deadlocks=20=F0=9F=94=80=20merge(manage?=
 =?UTF-8?q?r.py):=20refactor=20=5Fget=20method=20to=20extract=20logic=20in?=
 =?UTF-8?q?to=20=5Fget=5Fwithout=5Flock=20method=20for=20code=20reuse=20an?=
 =?UTF-8?q?d=20readability=20=F0=9F=94=80=20merge(manager.py):=20refactor?=
 =?UTF-8?q?=20set=20method=20to=20use=20=5Fget=5Fwithout=5Flock=20method?=
 =?UTF-8?q?=20for=20code=20reuse=20and=20readability=20=F0=9F=94=80=20merg?=
 =?UTF-8?q?e(manager.py):=20add=20lock=20acquisition=20in=20delete=20metho?=
 =?UTF-8?q?d=20to=20ensure=20thread=20safety=20=F0=9F=94=80=20merge(utils.?=
 =?UTF-8?q?py):=20add=20update=5Fcache=20method=20to=20Memoize=20class=20f?=
 =?UTF-8?q?or=20updating=20cache=20values?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../langflow/services/cache/manager.py        | 37 +++++++++++--------
 src/backend/langflow/services/cache/utils.py  |  5 +++
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/src/backend/langflow/services/cache/manager.py b/src/backend/langflow/services/cache/manager.py
index ec9909c75..a3f04518d 100644
--- a/src/backend/langflow/services/cache/manager.py
+++ b/src/backend/langflow/services/cache/manager.py
@@ -43,7 +43,7 @@ class InMemoryCache(BaseCacheManager):
             expiration_time (int, optional): Time in seconds after which a cached item expires. Default is 1 hour.
         """
         self._cache = OrderedDict()
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
         self.max_size = max_size
         self.expiration_time = expiration_time
 
@@ -58,18 +58,23 @@ class InMemoryCache(BaseCacheManager):
             The value associated with the key, or None if the key is not found or the item has expired.
         """
         with self._lock:
-            if key in self._cache:
-                item = self._cache.pop(key)
-                if (
-                    self.expiration_time is None
-                    or time.time() - item["time"] < self.expiration_time
-                ):
-                    # Move the key to the end to make it recently used
-                    self._cache[key] = item
-                    return item["value"]
-                else:
-                    self.delete(key)
-            return None
+            return self._get_without_lock(key)
+
+    def _get_without_lock(self, key):
+        """
+        Retrieve an item from the cache without acquiring the lock.
+        """
+        if item := self._cache.get(key):
+            if (
+                self.expiration_time is None
+                or time.time() - item["time"] < self.expiration_time
+            ):
+                # Move the key to the end to make it recently used
+                self._cache.move_to_end(key)
+                return item["value"]
+            else:
+                self.delete(key)
+        return None
 
     def set(self, key, value):
         """
@@ -100,7 +105,7 @@ class InMemoryCache(BaseCacheManager):
             value: The value to insert or update.
         """
         with self._lock:
-            existing_value = self.get(key)
+            existing_value = self._get_without_lock(key)
             if (
                 existing_value is not None
                 and isinstance(existing_value, dict)
@@ -135,8 +140,8 @@ class InMemoryCache(BaseCacheManager):
         Args:
             key: The key of the item to remove.
         """
-        # with self._lock:
-        self._cache.pop(key, None)
+        with self._lock:
+            self._cache.pop(key, None)
 
     def clear(self):
         """
diff --git a/src/backend/langflow/services/cache/utils.py b/src/backend/langflow/services/cache/utils.py
index 13e24dd6d..f73b8f023 100644
--- a/src/backend/langflow/services/cache/utils.py
+++ b/src/backend/langflow/services/cache/utils.py
@@ -212,6 +212,10 @@ class Memoize:
         cache_manager = self.get_cache_manager()
         return cache_manager.get(session_id)
 
+    def update_cache(self, session_id, value):
+        cache_manager = self.get_cache_manager()
+        cache_manager.set(session_id, value)
+
     def __call__(self, func: Callable[..., Any]):
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
@@ -228,5 +232,6 @@ class Memoize:
             return result
 
         wrapper.clear_cache = self.clear_cache
+        wrapper.update_cache = self.update_cache
         wrapper.get_result_by_session_id = self.get_result_by_session_id
         return wrapper