Merge remote-tracking branch 'origin/main' into feat/collaboration

api/core/app/apps/base_app_queue_manager.py

@@ -1,10 +1,12 @@
 import logging
 import queue
+import threading
 import time
 from abc import abstractmethod
 from enum import IntEnum, auto
 from typing import Any
 
+from cachetools import TTLCache, cachedmethod
 from redis.exceptions import RedisError
 from sqlalchemy.orm import DeclarativeMeta
 
@@ -45,6 +47,8 @@ class AppQueueManager:
         q: queue.Queue[WorkflowQueueMessage | MessageQueueMessage | None] = queue.Queue()
 
         self._q = q
+        self._stopped_cache: TTLCache[tuple, bool] = TTLCache(maxsize=1, ttl=1)
+        self._cache_lock = threading.Lock()
 
     def listen(self):
         """
@@ -157,6 +161,7 @@ class AppQueueManager:
         stopped_cache_key = cls._generate_stopped_cache_key(task_id)
         redis_client.setex(stopped_cache_key, 600, 1)
 
+    @cachedmethod(lambda self: self._stopped_cache, lock=lambda self: self._cache_lock)
     def _is_stopped(self) -> bool:
         """
         Check if task is stopped

api/core/rag/embedding/cached_embedding.py

@@ -43,8 +43,7 @@ class CacheEmbedding(Embeddings):
             else:
                 embedding_queue_indices.append(i)
 
-        # release database connection, because embedding may take a long time
-        db.session.close()
+        # NOTE: avoid closing the shared scoped session here; downstream code may still have pending work
 
         if embedding_queue_indices:
             embedding_queue_texts = [texts[i] for i in embedding_queue_indices]