fix(prompt-cache): avoid redundant refresh races

Author: hassiebp

langfuse/_client/client.py

@@ -3482,8 +3482,9 @@ def refresh_task() -> None:
                         fetch_timeout_seconds=fetch_timeout_seconds,
                     )
 
-                self._resources.prompt_cache.add_refresh_prompt_task(
+                self._resources.prompt_cache.add_refresh_prompt_task_if_current(
                     cache_key,
+                    cached_prompt,
                     refresh_task,
                 )
                 langfuse_logger.debug(

langfuse/_utils/prompt_cache.py

@@ -4,7 +4,7 @@
 import os
 from datetime import datetime
 from queue import Queue
-from threading import Thread
+from threading import RLock, Thread
 from typing import Callable, Dict, List, Optional, Set
 
 from langfuse._client.environment_variables import (
@@ -77,12 +77,14 @@ class PromptCacheTaskManager(object):
     _threads: int
     _queue: Queue
     _processing_keys: Set[str]
+    _lock: RLock
 
     def __init__(self, threads: int = 1):
         self._queue = Queue()
         self._consumers = []
         self._threads = threads
         self._processing_keys = set()
+        self._lock = RLock()
 
         for i in range(self._threads):
             consumer = PromptCacheRefreshConsumer(self._queue, i)
@@ -92,16 +94,20 @@ def __init__(self, threads: int = 1):
         atexit.register(self.shutdown)
 
     def add_task(self, key: str, task: Callable[[], None]) -> None:
-        if key not in self._processing_keys:
-            logger.debug(f"Adding prompt cache refresh task for key: {key}")
-            self._processing_keys.add(key)
-            wrapped_task = self._wrap_task(key, task)
-            self._queue.put((wrapped_task))
-        else:
-            logger.debug(f"Prompt cache refresh task already submitted for key: {key}")
+        with self._lock:
+            if key not in self._processing_keys:
+                logger.debug(f"Adding prompt cache refresh task for key: {key}")
+                self._processing_keys.add(key)
+                wrapped_task = self._wrap_task(key, task)
+                self._queue.put((wrapped_task))
+            else:
+                logger.debug(
+                    f"Prompt cache refresh task already submitted for key: {key}"
+                )
 
     def active_tasks(self) -> int:
-        return len(self._processing_keys)
+        with self._lock:
+            return len(self._processing_keys)
 
     def wait_for_idle(self) -> None:
         self._queue.join()
@@ -112,7 +118,8 @@ def wrapped() -> None:
             try:
                 task()
             finally:
-                self._processing_keys.remove(key)
+                with self._lock:
+                    self._processing_keys.remove(key)
                 logger.debug(f"Refreshed prompt cache for key: {key}")
 
         return wrapped
@@ -139,6 +146,7 @@ def shutdown(self) -> None:
 
 class PromptCache:
     _cache: Dict[str, PromptCacheItem]
+    _lock: RLock
 
     _task_manager: PromptCacheTaskManager
     """Task manager for refreshing cache"""
@@ -147,34 +155,60 @@ def __init__(
         self, max_prompt_refresh_workers: int = DEFAULT_PROMPT_CACHE_REFRESH_WORKERS
     ):
         self._cache = {}
+        self._lock = RLock()
         self._task_manager = PromptCacheTaskManager(threads=max_prompt_refresh_workers)
         logger.debug("Prompt cache initialized.")
 
     def get(self, key: str) -> Optional[PromptCacheItem]:
-        return self._cache.get(key, None)
+        with self._lock:
+            return self._cache.get(key, None)
 
     def set(self, key: str, value: PromptClient, ttl_seconds: Optional[int]) -> None:
         if ttl_seconds is None:
             ttl_seconds = DEFAULT_PROMPT_CACHE_TTL_SECONDS
 
-        self._cache[key] = PromptCacheItem(value, ttl_seconds)
+        with self._lock:
+            self._cache[key] = PromptCacheItem(value, ttl_seconds)
 
     def delete(self, key: str) -> None:
-        self._cache.pop(key, None)
+        with self._lock:
+            self._cache.pop(key, None)
 
     def invalidate(self, prompt_name: str) -> None:
         """Invalidate all cached prompts with the given prompt name."""
-        for key in list(self._cache):
-            if key.startswith(prompt_name):
-                del self._cache[key]
+        with self._lock:
+            for key in list(self._cache):
+                if key.startswith(prompt_name):
+                    del self._cache[key]
 
     def add_refresh_prompt_task(self, key: str, fetch_func: Callable[[], None]) -> None:
         logger.debug(f"Submitting refresh task for key: {key}")
         self._task_manager.add_task(key, fetch_func)
 
+    def add_refresh_prompt_task_if_current(
+        self,
+        key: str,
+        expected_item: PromptCacheItem,
+        fetch_func: Callable[[], None],
+    ) -> None:
+        with self._lock:
+            current_item = self._cache.get(key)
+            if (
+                current_item is not None
+                and current_item is not expected_item
+                and not current_item.is_expired()
+            ):
+                logger.debug(
+                    f"Skipping refresh task for key: {key} because cache is already fresh."
+                )
+                return
+
+        self.add_refresh_prompt_task(key, fetch_func)
+
     def clear(self) -> None:
         """Clear the entire prompt cache, removing all cached prompts."""
-        self._cache.clear()
+        with self._lock:
+            self._cache.clear()
 
     @staticmethod
     def generate_cache_key(

tests/unit/test_prompt.py

@@ -492,6 +492,58 @@ def test_get_stale_prompt_when_expired_cache_default_ttl(mock_time, langfuse: La
     assert updated_result == TextPromptClient(updated_prompt)
 
 
+@patch.object(PromptCacheItem, "get_epoch_seconds")
+def test_skip_redundant_refresh_when_cache_already_updated(
+    mock_time, langfuse: Langfuse
+) -> None:
+    prompt_name = "test_skip_redundant_refresh_when_cache_already_updated"
+    cache_key = PromptCache.generate_cache_key(prompt_name, version=None, label=None)
+
+    mock_time.return_value = 0
+
+    initial_prompt = Prompt_Text(
+        name=prompt_name,
+        version=1,
+        prompt="Make me laugh",
+        labels=[],
+        type="text",
+        config={},
+        tags=[],
+    )
+    updated_prompt = Prompt_Text(
+        name=prompt_name,
+        version=2,
+        prompt="Make me laugh",
+        labels=[],
+        type="text",
+        config={},
+        tags=[],
+    )
+
+    stale_result = TextPromptClient(initial_prompt)
+    fresh_result = TextPromptClient(updated_prompt)
+
+    langfuse._resources.prompt_cache.set(cache_key, stale_result, None)
+    stale_item = langfuse._resources.prompt_cache.get(cache_key)
+    assert stale_item is not None
+
+    mock_time.return_value = DEFAULT_PROMPT_CACHE_TTL_SECONDS + 1
+    assert stale_item.is_expired()
+
+    langfuse._resources.prompt_cache.set(cache_key, fresh_result, None)
+
+    add_task_mock = Mock()
+    langfuse._resources.prompt_cache._task_manager.add_task = add_task_mock
+
+    langfuse._resources.prompt_cache.add_refresh_prompt_task_if_current(
+        cache_key,
+        stale_item,
+        Mock(),
+    )
+
+    add_task_mock.assert_not_called()
+
+
 @patch.object(PromptCacheItem, "get_epoch_seconds")
 def test_get_fresh_prompt_when_expired_cache_default_ttl(mock_time, langfuse: Langfuse):
     mock_time.return_value = 0