Merge pull request #3 from adeptvin1/feature/fix-group-status-update

fix race condition

Author: adeptvin1

client/api/endpoints.py

@@ -41,6 +41,8 @@ def default(self, obj):
 
 # Сет для отслеживания активных экспериментов
 active_experiments: Set[str] = set()
+# Блокировка для предотвращения race conditions при работе с active_experiments
+active_experiments_lock = asyncio.Lock()
 
 
 # Функция для преобразования документа MongoDB в JSON-сериализуемый формат
@@ -396,21 +398,52 @@ async def check_and_update_group_status(experiment_id: str, db):
 
 async def run_experiment(experiment_id: str):
     """Фоновая задача для выполнения эксперимента с динамической нагрузкой."""
-    if experiment_id in active_experiments:
-        logger.warning(f"Experiment {experiment_id} is already running")
-        return
-
-    active_experiments.add(experiment_id)
+    # Используем блокировку для предотвращения race conditions
+    async with active_experiments_lock:
+        if experiment_id in active_experiments:
+            logger.warning(f"Experiment {experiment_id} is already running")
+            return
+        
+        # Проверяем статус эксперимента в БД перед запуском
+        try:
+            db = await get_database()
+            experiment = await db.experiments.find_one({"_id": ObjectId(experiment_id)})
+            if not experiment:
+                logger.error(f"Experiment {experiment_id} not found")
+                return
+            
+            # Если эксперимент уже завершен или провалился, не запускаем его снова
+            exp_state = experiment.get("state")
+            if exp_state in [ExperimentState.COMPLETED, ExperimentState.FAILED]:
+                logger.info(f"Experiment {experiment_id} is already {exp_state}, skipping")
+                return
+            
+            # Если эксперимент уже выполняется (по статусу в БД), не запускаем его снова
+            if exp_state == ExperimentState.RUNNING:
+                logger.warning(f"Experiment {experiment_id} is already running (status in DB)")
+                return
+        
+        except Exception as e:
+            logger.error(f"Error checking experiment {experiment_id} status: {str(e)}")
+            return
+        
+        # Добавляем в активные эксперименты только после всех проверок
+        active_experiments.add(experiment_id)
+    
+    # Инициализация переменных вне блока блокировки
     semaphore = asyncio.Semaphore(MAX_CONCURRENT_TASKS)
     save_semaphore = asyncio.Semaphore(1)  # Семафор для сохранения результатов
-    db = None  # Инициализируем для доступа в finally блоке
 
     try:
+        # Получаем эксперимент еще раз для работы
         db = await get_database()
-        experiment = await with_retry(get_experiment_by_id, experiment_id, db)
+        experiment = await with_retry(_get_experiment_from_db, experiment_id, db)
 
         if not experiment:
             logger.error(f"Experiment {experiment_id} not found")
+            # Удаляем из активных, если эксперимент не найден
+            async with active_experiments_lock:
+                active_experiments.discard(experiment_id)
             return
 
         settings = experiment["settings"]
@@ -598,7 +631,9 @@ async def run_experiment(experiment_id: str):
                 # Логируем ошибку, но не прерываем выполнение
                 logger.error(f"Failed to check group status in finally block: {str(group_check_error)}")
 
-        active_experiments.remove(experiment_id)
+        # Удаляем эксперимент из активных с блокировкой
+        async with active_experiments_lock:
+            active_experiments.discard(experiment_id)  # Используем discard вместо remove для безопасности
 
 
 @router.get("/experiment_stats")
@@ -852,20 +887,73 @@ async def manage_group(
     if not group:
         raise HTTPException(status_code=404, detail="Group not found")
 
+    current_group_state = group.get("state", ExperimentState.PENDING)
+    
     new_state = ""
     if state == "start":
+        # Проверяем, не завершена ли уже группа
+        if current_group_state == ExperimentState.COMPLETED:
+            logger.warning(f"Group {group_id} is already completed, cannot start")
+            raise HTTPException(
+                status_code=400,
+                detail="Group is already completed and cannot be started again"
+            )
+        
+        # Проверяем, не запущена ли уже группа
+        if current_group_state == ExperimentState.RUNNING:
+            logger.warning(f"Group {group_id} is already running")
+            # Возвращаем текущий статус без повторного запуска
+            return {"_id": group_id, "state": current_group_state, "message": "Group is already running"}
+        
         new_state = ExperimentState.RUNNING
-        # Запускаем все эксперименты в группе
-        for exp_id in group["experiment_ids"]:
-            asyncio.create_task(run_experiment(exp_id))
+        
+        # Запускаем только те эксперименты, которые еще не завершены и не запущены
+        experiment_ids = group.get("experiment_ids", [])
+        started_count = 0
+        skipped_count = 0
+        
+        for exp_id in experiment_ids:
+            try:
+                # Проверяем статус эксперимента перед запуском
+                experiment = await db.experiments.find_one({"_id": ObjectId(exp_id)})
+                if not experiment:
+                    logger.warning(f"Experiment {exp_id} not found, skipping")
+                    skipped_count += 1
+                    continue
+                
+                exp_state = experiment.get("state")
+                
+                # Пропускаем уже завершенные или провалившиеся эксперименты
+                if exp_state in [ExperimentState.COMPLETED, ExperimentState.FAILED]:
+                    logger.debug(f"Experiment {exp_id} is already {exp_state}, skipping")
+                    skipped_count += 1
+                    continue
+                
+                # Проверяем, не запущен ли уже эксперимент
+                async with active_experiments_lock:
+                    if exp_id in active_experiments:
+                        logger.debug(f"Experiment {exp_id} is already running, skipping")
+                        skipped_count += 1
+                        continue
+                
+                # Запускаем эксперимент
+                asyncio.create_task(run_experiment(exp_id))
+                started_count += 1
+                
+            except Exception as e:
+                logger.error(f"Error starting experiment {exp_id}: {str(e)}")
+                skipped_count += 1
+        
+        logger.info(f"Group {group_id}: started {started_count} experiments, skipped {skipped_count}")
+        
     elif state == "pause":
         new_state = ExperimentState.PAUSED
     elif state == "stop":
         new_state = ExperimentState.COMPLETED
 
     await db.groups.update_one(
         {"_id": ObjectId(group_id)},
-        {"$set": {"state": new_state}}
+        {"$set": {"state": new_state, "updated_at": datetime.now()}}
     )
 
     return {"_id": group_id, "state": new_state}