♻️ Performance optimization of the knowledge base file list query interface (#3025)

* ♻️ Performance optimization of the knowledge base file list query interface

* ♻️ Performance optimization of the knowledge base file list query interface: ut

Author: WMC001

backend/apps/data_process_app.py

@@ -204,9 +204,14 @@ async def get_index_tasks(index_name: str):
 
     Returns tasks that are being processed or waiting to be processed
     """
+    import time
+    start = time.time()
     try:
-        return await service.get_index_tasks(index_name)
+        result = await service.get_index_tasks(index_name)
+        logger.info(f"[get_index_tasks] index={index_name}, tasks={len(result)}, duration={time.time()-start:.3f}s")
+        return result
     except Exception as e:
+        logger.error(f"[get_index_tasks] error: {e}")
         raise HTTPException(
             status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=str(e))
 

backend/services/data_process_service.py

@@ -54,7 +54,7 @@ def __init__(self):
 
         self._inspector = None
         self._inspector_last_time = 0
-        self._inspector_ttl = 60  # Inspector cache time in seconds
+        self._inspector_ttl = 300  # 5 minutes - inspector is expensive to create (ping all workers)
         self._inspector_lock = None
         self._inspector_lock = threading.Lock()
 
@@ -105,7 +105,7 @@ async def stop(self):
         logger.info("Data processing service stopped")
 
     def _get_celery_inspector(self):
-        """Get Celery inspector"""
+        """Get Celery inspector (cached for performance)"""
         with self._inspector_lock:
             now = time.time()
             if self._inspector and now - self._inspector_last_time < self._inspector_ttl:
@@ -117,9 +117,9 @@ def _get_celery_inspector(self):
                     f"Celery broker URL is not configured properly, reconfiguring to {celery_app.conf.broker_url}")
             try:
                 inspector = celery_app.control.inspect()
-                inspector.ping()
                 self._inspector = inspector
                 self._inspector_last_time = now
+                self._inspector_init_time = now
                 return inspector
             except Exception as e:
                 self._inspector = None
@@ -142,11 +142,9 @@ async def get_all_tasks(self, filter: bool = True) -> List[Dict[str, Any]]:
         all_tasks = []
         try:
             start_time = time.time()
-            logger.debug(
-                "Getting inspector to check for active and reserved tasks (concurrent)")
+            inspector_start = time.time()
             inspector = self._get_celery_inspector()
-            logger.debug(
-                f"⏰ Inspector initialization took {time.time() - start_time}s")
+            inspector_duration = time.time() - inspector_start
 
             # Collect task IDs from different sources and keep runtime metadata
             task_ids = set()
@@ -171,18 +169,37 @@ def _normalize_runtime_meta(task: Dict[str, Any]) -> Dict[str, Any]:
                     'original_filename': kwargs.get('original_filename', ''),
                 }
 
+            celery_start = time.time()
+
+            # Use short timeout for inspector since workers can respond in ~0.1s
+            # Default 1s timeout is unnecessary and causes delay
+            short_timeout = 0.2
+
             def get_active():
-                return inspector.active()
+                t = time.time()
+                # Create fresh inspector with short timeout for each call
+                short_inspector = celery_app.control.inspect(timeout=short_timeout)
+                result = short_inspector.active()
+                elapsed = time.time() - t
+                logger.info(f"[get_all_tasks] inspector.active() took {elapsed:.3f}s")
+                return result if result else {}
 
             def get_reserved():
-                return inspector.reserved()
+                t = time.time()
+                short_inspector = celery_app.control.inspect(timeout=short_timeout)
+                result = short_inspector.reserved()
+                elapsed = time.time() - t
+                logger.info(f"[get_all_tasks] inspector.reserved() took {elapsed:.3f}s")
+                return result if result else {}
+
             with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
                 future_active = executor.submit(get_active)
                 future_reserved = executor.submit(get_reserved)
-                active_tasks_dict = future_active.result()
-                reserved_tasks_dict = future_reserved.result()
-            logger.debug(
-                f"⏰ Get active and reserved tasks (concurrent) took {time.time() - start_time}s")
+                active_tasks_dict = future_active.result(timeout=short_timeout + 0.5)
+                reserved_tasks_dict = future_reserved.result(timeout=short_timeout + 0.5)
+            celery_duration = time.time() - celery_start
+            if celery_duration > 0.5:
+                logger.warning(f"[get_all_tasks] Inspector took {celery_duration:.3f}s (expected <0.5s)")
             if active_tasks_dict:
                 for worker, tasks in active_tasks_dict.items():
                     for task in tasks:
@@ -199,23 +216,17 @@ def get_reserved():
                             # Keep active metadata if already present
                             runtime_task_meta.setdefault(task_id, _normalize_runtime_meta(task))
 
-            # Currently, we don't have scheduled tasks, so skip getting scheduled tasks here
-            start_time = time.time()
-            logger.debug("Getting task IDs from Redis backend")
-            # Also get task IDs from Redis backend (covers completed/failed tasks within expiry)
+            # Get task IDs from Redis backend (covers completed/failed tasks within expiry)
             try:
                 redis_task_ids = get_all_task_ids_from_redis(self.redis_client)
-                logger.debug(
-                    f"⏰ Get Redis task IDs took {time.time() - start_time}s")
                 for task_id in redis_task_ids:
-                    # Add to the set, duplicates will be handled
                     task_ids.add(task_id)
             except Exception as redis_error:
                 logger.warning(
                     f"Failed to query Redis for stored task IDs: {str(redis_error)}")
-            logger.debug(
-                f"Total unique task IDs collected (inspector + Redis): {len(task_ids)}")
+
             task_id_list = list(task_ids)
+            # Batch fetch all task info
             tasks = [get_task_info(task_id) for task_id in task_id_list]
             all_task_infos = await asyncio.gather(*tasks, return_exceptions=True)
             for idx, task_info in enumerate(all_task_infos):
@@ -243,7 +254,6 @@ def get_reserved():
                     if not task_info.get('index_name'):
                         continue
                 all_tasks.append(task_info)
-            logger.debug(f"Retrieved {len(all_tasks)} tasks.")
         except Exception as e:
             logger.error(f"Error retrieving all tasks: {str(e)}")
             all_tasks = []

backend/services/redis_service.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import re
-from typing import Dict, Any, Optional, Tuple, Set
+from typing import Dict, Any, Optional, Tuple, Set, List
 
 import redis
 
@@ -24,8 +24,8 @@ def client(self) -> redis.Redis:
             if not REDIS_URL:
                 raise ValueError("REDIS_URL environment variable is not set")
             self._client = redis.from_url(
-                REDIS_URL, 
-                socket_timeout=5, 
+                REDIS_URL,
+                socket_timeout=5,
                 socket_connect_timeout=5,
                 decode_responses=True
             )
@@ -654,13 +654,13 @@ def save_error_info(self, task_id: str, error_reason: str, ttl_days: int = 30) -
             if not error_reason:
                 logger.error(f"Cannot save error info for task {task_id}: error_reason is empty")
                 return False
-            
+
             ttl_seconds = ttl_days * 24 * 60 * 60
             reason_key = f"error:reason:{task_id}"
 
             # Save error reason
             result = self.client.setex(reason_key, ttl_seconds, error_reason)
-            
+
             if result:
                 logger.info(f"Successfully saved error info to Redis for task {task_id}, key: {reason_key}")
                 # Verify the save by reading it back
@@ -695,13 +695,13 @@ def save_progress_info(self, task_id: str, processed_chunks: int, total_chunks:
             if not task_id:
                 logger.error("Cannot save progress info: task_id is empty")
                 return False
-            
+
             progress_key = f"progress:{task_id}"
             progress_data = {
                 'processed_chunks': processed_chunks,
                 'total_chunks': total_chunks
             }
-            
+
             ttl_seconds = ttl_hours * 3600
             progress_json = json.dumps(progress_data)
             self.client.setex(
@@ -874,6 +874,79 @@ def get_error_info(self, task_id: str) -> Optional[str]:
                 f"Failed to get error info for task {task_id}: {str(e)}")
             return None
 
+    def batch_get_progress_info(self, task_ids: List[str]) -> Dict[str, Optional[Dict[str, int]]]:
+        """
+        Batch get progress information for multiple tasks in a single Redis call.
+
+        Args:
+            task_ids: List of Celery task IDs
+
+        Returns:
+            Dict mapping task_id to progress info dict, or None if not found
+        """
+        if not task_ids:
+            return {}
+
+        try:
+            # Build list of keys
+            progress_keys = [f"progress:{tid}" for tid in task_ids]
+            # Use pipeline for batch operation
+            pipe = self.client.pipeline()
+            for key in progress_keys:
+                pipe.get(key)
+            results = pipe.execute()
+
+            # Build result dict
+            result = {}
+            for i, task_id in enumerate(task_ids):
+                progress_data = results[i]
+                if progress_data:
+                    try:
+                        if isinstance(progress_data, bytes):
+                            progress_data = progress_data.decode('utf-8')
+                        result[task_id] = json.loads(progress_data)
+                    except (json.JSONDecodeError, TypeError):
+                        result[task_id] = None
+                else:
+                    result[task_id] = None
+            return result
+        except Exception as e:
+            logger.warning(f"Failed to batch get progress info: {str(e)}")
+            return {tid: None for tid in task_ids}
+
+    def batch_get_error_info(self, task_ids: List[str]) -> Dict[str, Optional[str]]:
+        """
+        Batch get error information for multiple tasks in a single Redis call.
+
+        Args:
+            task_ids: List of Celery task IDs
+
+        Returns:
+            Dict mapping task_id to error reason string, or None if not found
+        """
+        if not task_ids:
+            return {}
+
+        try:
+            # Build list of keys
+            error_keys = [f"error:reason:{tid}" for tid in task_ids]
+            # Use pipeline for batch operation
+            pipe = self.client.pipeline()
+            for key in error_keys:
+                pipe.get(key)
+            results = pipe.execute()
+
+            # Build result dict
+            result = {}
+            for i, task_id in enumerate(task_ids):
+                reason = results[i]
+                # With decode_responses=True, reason is already a string
+                result[task_id] = reason if reason else None
+            return result
+        except Exception as e:
+            logger.warning(f"Failed to batch get error info: {str(e)}")
+            return {tid: None for tid in task_ids}
+
 # Global Redis service instance
 _redis_service = None