Race Condition in Worker._process_task_completion leading to a crash

[alonkukl]

Description

There is a race condition between _process_task_completion and _handle_cancel methods in the Worker class, where concurrent modification of self._tasks can lead to an exception. The race condition occurs when one thread is processing task completion while another thread handles a cancellation request.

Affected Methods

Method 1: _process_task_completion

def _process_task_completion(self, task: RuntimeTask, result: Any) -> None:
    assert task is self._active_task
    packaged_result = RuntimeResult(task.return_address, result, self._id)
    if task.return_address not in self._tasks:  # Check existence
        return
    # ... other operations ...
    self._tasks.pop(task.return_address)  # Attempt to remove

Method 2: _handle_cancel
This method runs in a separate thread and modifies the same self._tasks dictionary, causing the race condition.

Race Condition Scenario

1. Thread 1 (in _process_task_completion): Checks if task.return_address not in self._tasks - condition is false (task exists)
2. Thread 2 (in _handle_cancel): Removes the task from self._tasks
3. Thread 1 continues: Attempts self._tasks.pop(task.return_address)
   Raises KeyError because the task was already removed by the cancel handler