Fix Orphaned tasks stuck in CeleryExecutor as running (apache#16550)

(cherry picked from commit 90f0088) (cherry picked from commit 83fb4bf) (cherry picked from commit d44c223)
astronomer · Jun 23, 2021 · 9c76e8f · 9c76e8f
1 parent 55b836b
commit 9c76e8f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/airflow/executors/celery_executor.py b/airflow/executors/celery_executor.py
@@ -377,9 +377,7 @@ def _check_for_stalled_adopted_tasks(self):
                 "\n\t".join([repr(x) for x in timedout_keys]),
             )
             for key in timedout_keys:
-                self.event_buffer[key] = (State.FAILED, None)
-                del self.tasks[key]
-                del self.adopted_task_timeouts[key]
+                self.change_state(key, State.FAILED)
 
     def debug_dump(self) -> None:
         """Called in response to SIGUSR2 by the scheduler"""

diff --git a/tests/executors/test_celery_executor.py b/tests/executors/test_celery_executor.py
@@ -385,11 +385,14 @@ def test_check_for_stalled_adopted_tasks(self):
             key_1: queued_dttm + executor.task_adoption_timeout,
             key_2: queued_dttm + executor.task_adoption_timeout,
         }
+        executor.running = {key_1, key_2}
         executor.tasks = {key_1: AsyncResult("231"), key_2: AsyncResult("232")}
         executor.sync()
-        self.assertEqual(executor.event_buffer, {key_1: (State.FAILED, None), key_2: (State.FAILED, None)})
-        self.assertEqual(executor.tasks, {})
-        self.assertEqual(executor.adopted_task_timeouts, {})
+
+        assert executor.event_buffer == {key_1: (State.FAILED, None), key_2: (State.FAILED, None)}
+        assert executor.tasks == {}
+        assert executor.running == set()
+        assert executor.adopted_task_timeouts == {}
 
 
 def test_operation_timeout_config():