UN-2901 [FIX] Prevent race conditions in distributed file processing causing status corruption and premature cleanup (#1609)

This commit implements multiple defensive fixes to prevent race conditions when duplicate
workers process the same file execution, which was causing:
1. COMPLETED/ERROR status being overwritten with PENDING→EXECUTING by late workers
2. Premature chord callback cleanup (FileExpired errors) while workers still processing
3. Missing FINALIZATION stage detection during grace period
4. Incorrect COMPLETED:SUCCESS status for failed executions

Fixes implemented:
- Pre-creation: Check FileExecutionStatusTracker (Redis) + DB before PENDING update
- Grace period: Extended duplicate detection to include FINALIZATION stage
- Destination: Wait for DESTINATION_PROCESSING lock release (prevents chord cleanup)
- Terminal states: Handle both COMPLETED and ERROR in processor.py
- Status tracking: Set COMPLETED:FAILED for failed executions (was SUCCESS)
- Refactoring: Consolidated completed_ttl to method initialization

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude <noreply@anthropic.com>

Author: muhammad-ali-e

unstract/tool-sandbox/src/unstract/tool_sandbox/helper.py

@@ -166,20 +166,27 @@ def poll_tool_status(
                     current_tracker_data = tracker.get_data(
                         self.execution_id, file_execution_id
                     )
+                    # Check if already at FINALIZATION or COMPLETED (tool execution done)
                     if (
                         current_tracker_data
                         and current_tracker_data.stage_status.stage
-                        == FileExecutionStage.COMPLETED
+                        in [
+                            FileExecutionStage.FINALIZATION,
+                            FileExecutionStage.COMPLETED,
+                        ]
                     ):
+                        stage = current_tracker_data.stage_status.stage
+                        status = current_tracker_data.stage_status.status
                         logger.warning(
-                            f"File execution COMPLETED by another worker during NOT_FOUND grace period - "
-                            f"duplicate run detected after {not_found_duration:.1f}s for execution_id: {self.execution_id}, file_execution_id: {file_execution_id}"
+                            f"File execution already at {stage}:{status} by another worker "
+                            f"during NOT_FOUND grace period - duplicate run detected after {not_found_duration:.1f}s "
+                            f"for execution_id: {self.execution_id}, file_execution_id: {file_execution_id}"
                         )
                         # Break with ERROR to stop further processing (destination, finalization)
-                        # This is not a real error - the file was successfully processed by another worker
+                        # This is not a real error - the file was already processed by another worker
                         response = self._create_run_response(
                             status=RunnerContainerRunStatus.ERROR,
-                            error="File already completed by another worker (duplicate run avoided)",
+                            error=f"File already processed by another worker (stage: {stage}:{status}, duplicate run avoided)",
                         )
                         break
 

workers/file_processing/tasks.py

@@ -1154,16 +1154,19 @@ def _pre_create_file_executions(
                 force_create=not use_file_history,  # Force create when file history is disabled
             )
 
-            # EARLY EXIT: Skip if file already COMPLETED in this execution
-            # This catches race conditions where Worker 1 completes before Worker 2's pre-create runs
+            # EARLY EXIT: Skip if file already in terminal state (COMPLETED/ERROR) in this execution
+            # This catches race conditions where Worker 1 completes/fails before Worker 2's pre-create runs
             # Only checks current execution_id (respects file_history cleanup for different executions)
-            if (
-                hasattr(workflow_file_execution, "status")
-                and workflow_file_execution.status == ExecutionStatus.COMPLETED.value
-            ):
+            if hasattr(
+                workflow_file_execution, "status"
+            ) and workflow_file_execution.status in [
+                ExecutionStatus.COMPLETED.value,
+                ExecutionStatus.ERROR.value,
+            ]:
                 logger.info(
-                    f"File '{file_name}' already COMPLETED in execution {execution_id} "
-                    f"(file_execution_id: {workflow_file_execution.id}). Skipping duplicate task creation."
+                    f"File '{file_name}' already in terminal state: {workflow_file_execution.status} "
+                    f"in execution {execution_id} (file_execution_id: {workflow_file_execution.id}). "
+                    f"Skipping duplicate task creation."
                 )
                 file_identifier = f"{file_hash.provider_file_uuid}:{file_hash.file_path}"
                 skipped_already_completed.append(file_identifier)
@@ -1219,16 +1222,65 @@ def _pre_create_file_executions(
 
             else:
                 # NORMAL FILE: Set initial status to PENDING for processing
-                try:
-                    api_client.file_client.update_file_execution_status(
-                        file_execution_id=workflow_file_execution.id,
-                        status=ExecutionStatus.PENDING.value,
+                # But preserve EXECUTING status for crash recovery (don't overwrite active processing)
+
+                # RACE CONDITION FIX: Check FileExecutionStatusTracker (Redis, real-time) before updating
+                # This prevents late-arriving workers from overwriting COMPLETED/ERROR status with PENDING
+                from unstract.core.file_execution_tracker import (
+                    FileExecutionStage,
+                    FileExecutionStatusTracker,
+                )
+
+                tracker = FileExecutionStatusTracker()
+                current_tracker_data = tracker.get_data(
+                    execution_id, workflow_file_execution.id
+                )
+
+                # Primary check: FileExecutionStatusTracker (real-time, Redis-based)
+                if current_tracker_data and current_tracker_data.stage_status.stage in [
+                    FileExecutionStage.FINALIZATION,
+                    FileExecutionStage.COMPLETED,
+                ]:
+                    logger.info(
+                        f"File '{file_name}' already at {current_tracker_data.stage_status.stage} "
+                        f"stage by another worker - skipping PENDING update "
+                        f"(file_execution_id: {workflow_file_execution.id})"
                     )
-                except Exception as status_error:
-                    logger.warning(
-                        f"Failed to set initial PENDING status for file {file_name}: {status_error}"
+                    # Skip to next file without updating status or adding to pre_created_data
+                    continue
+
+                # Fallback check: DB status (in case Redis down or tracker expired)
+                current_status = getattr(workflow_file_execution, "status", None)
+                if current_status in [
+                    ExecutionStatus.COMPLETED.value,
+                    ExecutionStatus.ERROR.value,
+                ]:
+                    logger.info(
+                        f"File '{file_name}' already in terminal state {current_status} "
+                        f"(DB status) - skipping PENDING update "
+                        f"(file_execution_id: {workflow_file_execution.id})"
+                    )
+                    # Skip to next file
+                    continue
+
+                # Preserve EXECUTING for crash recovery (don't reset to PENDING)
+                if current_status != ExecutionStatus.EXECUTING.value:
+                    try:
+                        api_client.file_client.update_file_execution_status(
+                            file_execution_id=workflow_file_execution.id,
+                            status=ExecutionStatus.PENDING.value,
+                        )
+                        logger.info(f"Set initial PENDING status for file {file_name}")
+                    except Exception as status_error:
+                        logger.warning(
+                            f"Failed to set initial PENDING status for file {file_name}: {status_error}"
+                        )
+                        # Don't fail the entire creation if status update fails
+                else:
+                    logger.info(
+                        f"File '{file_name}' already EXECUTING - preserving status for crash recovery/resume "
+                        f"(file_execution_id: {workflow_file_execution.id})"
                     )
-                    # Don't fail the entire creation if status update fails
 
                 # Add to pre_created_data for processing
                 pre_created_data[file_name] = PreCreatedFileData(

workers/shared/processing/files/processor.py

@@ -182,23 +182,34 @@ def validate_workflow_file_execution(
         )
 
         # RACE CONDITION FIX: Fetch fresh status from DB instead of using cached object
-        # This prevents late-arriving workers from checking stale data and overwriting COMPLETED status
+        # This prevents late-arriving workers from checking stale data and overwriting terminal status
         workflow_file_execution = context.api_client.get_workflow_file_execution(
             context.workflow_file_execution_id
         )
 
-        # Check if file execution is already completed
-        if workflow_file_execution.status == ExecutionStatus.COMPLETED.value:
+        # Check if file execution is already in terminal state (COMPLETED or ERROR)
+        if workflow_file_execution.status in [
+            ExecutionStatus.COMPLETED.value,
+            ExecutionStatus.ERROR.value,
+        ]:
             logger.info(
-                f"File already completed. Skipping execution for execution_id: {context.execution_id}, "
+                f"File already in terminal state: {workflow_file_execution.status}. "
+                f"Skipping execution for execution_id: {context.execution_id}, "
                 f"file_execution_id: {workflow_file_execution.id}"
             )
 
+            # Return appropriate result based on terminal status
             return FileProcessingResult(
                 file_name=context.file_name,
                 file_execution_id=workflow_file_execution.id,
-                success=True,
-                error=None,
+                success=(
+                    workflow_file_execution.status == ExecutionStatus.COMPLETED.value
+                ),
+                error=(
+                    getattr(workflow_file_execution, "execution_error", None)
+                    if workflow_file_execution.status == ExecutionStatus.ERROR.value
+                    else None
+                ),
                 result=getattr(workflow_file_execution, "result", None),
                 metadata=getattr(workflow_file_execution, "metadata", None) or {},
             )

workers/shared/workflow/destination_connector.py

@@ -330,36 +330,6 @@ def _check_and_acquire_destination_lock(
         try:
             tracker = FileExecutionStatusTracker()
 
-            # Check if destination already processed or in progress
-            existing_data = tracker.get_data(
-                exec_ctx.execution_id, exec_ctx.file_execution_id
-            )
-
-            if existing_data:
-                current_stage = existing_data.stage_status.stage
-                # Check for DESTINATION_PROCESSING, FINALIZATION, or COMPLETED
-                if current_stage in [
-                    FileExecutionStage.DESTINATION_PROCESSING,
-                    FileExecutionStage.FINALIZATION,
-                    FileExecutionStage.COMPLETED,
-                ]:
-                    # Enhanced debug log with full context for internal debugging
-                    logger.warning(
-                        f"DUPLICATE DETECTION: File '{file_ctx.file_name}' destination already at stage "
-                        f"{current_stage.value}. execution_id={exec_ctx.execution_id}, "
-                        f"file_execution_id={exec_ctx.file_execution_id}, worker_pid={os.getpid()}. "
-                        f"Skipping duplicate processing - another worker completed this already."
-                    )
-                    # UI log removed - duplication is internal detail, not user-facing error
-                    return False  # Duplicate detected
-
-            # Atomically acquire lock using Redis SET NX
-            # Use helper method for consistent lock key format
-            lock_key = tracker.get_destination_lock_key(
-                exec_ctx.execution_id, exec_ctx.file_execution_id
-            )
-            lock_token = str(uuid.uuid4())  # Unique token for debugging
-
             # Get TTL values from environment
             LOCK_TTL = int(
                 os.environ.get("DESTINATION_PROCESSING_LOCK_TTL_IN_SECOND", 120)
@@ -368,27 +338,65 @@ def _check_and_acquire_destination_lock(
                 os.environ.get("DESTINATION_PROCESSING_STAGE_TTL_IN_SECOND", 600)
             )
 
+            # Get lock key for atomic operations
+            lock_key = tracker.get_destination_lock_key(
+                exec_ctx.execution_id, exec_ctx.file_execution_id
+            )
+            lock_token = str(uuid.uuid4())  # Unique token for debugging
+
+            # STEP 1: Try to acquire lock atomically (SOURCE OF TRUTH)
+            # This is atomic - if lock exists, another worker is processing
             logger.info(
                 f"Attempting to acquire destination lock for file '{file_ctx.file_name}' "
                 f"(lock_key={lock_key}, lock_ttl={LOCK_TTL}s, stage_ttl={STAGE_TTL}s)"
             )
 
-            # Use Redis SET NX (Set if Not eXists) for true atomic lock acquisition
-            # Returns True if key was set (lock acquired), False if key already exists
             lock_acquired = tracker.redis_client.set(
                 lock_key, lock_token, nx=True, ex=LOCK_TTL
             )
 
+            # STEP 2: If lock acquisition failed, another worker is processing - WAIT
             if not lock_acquired:
-                # Enhanced debug log with full context for internal debugging
-                logger.warning(
-                    f"DUPLICATE DETECTION: Lock acquisition failed for file '{file_ctx.file_name}'. "
-                    f"execution_id={exec_ctx.execution_id}, file_execution_id={exec_ctx.file_execution_id}, "
-                    f"lock_key={lock_key}, worker_pid={os.getpid()}. "
-                    f"Another worker holds the lock - skipping duplicate processing."
+                logger.info(
+                    f"Lock already held by another worker for file '{file_ctx.file_name}'. "
+                    f"Waiting for lock release to prevent premature chord cleanup (max {LOCK_TTL}s)..."
                 )
-                # UI log removed - duplication is internal detail, not user-facing error
-                return False
+
+                wait_start = time.time()
+                max_wait = min(LOCK_TTL, 120)  # Wait up to lock TTL or 120s
+
+                # Poll until lock released or timeout
+                while time.time() - wait_start < max_wait:
+                    # Check if lock released
+                    if not tracker.redis_client.exists(lock_key):
+                        wait_duration = time.time() - wait_start
+                        # Lock released BEFORE timeout → Other worker finished (success or error) → Skip
+                        logger.info(
+                            f"Lock released for '{file_ctx.file_name}' after {wait_duration:.1f}s - "
+                            f"other worker completed processing. Skipping as duplicate."
+                        )
+                        return False  # Skip immediately
+
+                    time.sleep(2)  # Poll every 2 seconds
+
+                # STEP 3: After wait, try to acquire lock again
+                lock_acquired = tracker.redis_client.set(
+                    lock_key, lock_token, nx=True, ex=LOCK_TTL
+                )
+
+                if not lock_acquired:
+                    # Still can't acquire lock (timeout or another worker grabbed it)
+                    if tracker.redis_client.exists(lock_key):
+                        logger.warning(
+                            f"Lock still held after {max_wait}s timeout for '{file_ctx.file_name}'. "
+                            f"Another worker may have acquired it. Skipping as duplicate."
+                        )
+                    else:
+                        logger.warning(
+                            f"Failed to acquire lock for '{file_ctx.file_name}' after wait. "
+                            f"Another worker may have grabbed it first. Skipping as duplicate."
+                        )
+                    return False  # Skip
 
             # Lock acquired successfully - now set DESTINATION_PROCESSING stage
             logger.info(

workers/shared/workflow/execution/service.py

@@ -94,6 +94,11 @@ def execute_workflow_for_file(
         destination_start_time = start_time
         destination_end_time = start_time
 
+        # Configure TTL for COMPLETED stage (shorter TTL to optimize Redis memory)
+        completed_ttl = int(
+            os.environ.get("FILE_EXECUTION_TRACKER_COMPLETED_TTL_IN_SECOND", 300)
+        )
+
         try:
             logger.info(f"Executing workflow {workflow_id} for file {file_name}")
 
@@ -409,13 +414,6 @@ def execute_workflow_for_file(
                         ),
                     )
                     logger.info(f"Marked finalization as successful for {file_name}")
-
-                    # Use shorter TTL for COMPLETED stage to optimize Redis memory
-                    completed_ttl = int(
-                        os.environ.get(
-                            "FILE_EXECUTION_TRACKER_COMPLETED_TTL_IN_SECOND", 300
-                        )
-                    )
                     tracker.update_stage_status(
                         execution_id=execution_id,
                         file_execution_id=workflow_file_execution_id,
@@ -442,6 +440,17 @@ def execute_workflow_for_file(
                             error=error_msg,
                         ),
                     )
+                    # Set COMPLETED:FAILED to mark file as done (prevents duplicate processing)
+                    tracker.update_stage_status(
+                        execution_id=execution_id,
+                        file_execution_id=workflow_file_execution_id,
+                        stage_status=FileExecutionStageData(
+                            stage=FileExecutionStage.COMPLETED,
+                            status=FileExecutionStageStatus.FAILED,
+                            error=error_msg,
+                        ),
+                        ttl_in_second=completed_ttl,
+                    )
                     logger.info(f"Tracked failed execution for {file_name}: {error_msg}")
                 else:
                     # Duplicate skip - don't update any stages to avoid interfering with active worker