Test deadlock

Author: crusaderky

distributed/tests/test_worker_state_machine.py

@@ -5,7 +5,13 @@
 
 from distributed.protocol.serialize import Serialize
 from distributed.utils import recursive_to_dict
-from distributed.utils_test import _LockedCommPool, assert_story, gen_cluster, inc
+from distributed.utils_test import (
+    _LockedCommPool,
+    assert_story,
+    freeze_data_fetching,
+    gen_cluster,
+    inc,
+)
 from distributed.worker_state_machine import (
     ExecuteFailureEvent,
     ExecuteSuccessEvent,
@@ -348,3 +354,31 @@ async def test_in_memory_while_in_flight(c, s, a, b):
     # Let the comm from b to a return the result
     event.set()
     assert await y == 4  # Data in flight from b has been discarded
+
+
+@gen_cluster(client=True)
+async def test_forget_data_needed(c, s, a, b):
+    """
+    1. A task transitions to fetch and is added to data_needed
+    2. _ensure_communicating runs, but the network is saturated so the task is not
+       popped from data_needed
+    3. Task is forgotten
+    4. Task is recreated from scratch and transitioned to fetch again
+    5. BUG: at the moment of writing this test, data_needed.push silently did nothing,
+       because it still contained the forgotten task, which is a different TaskState
+       instance which will be no longer updated.
+    6. _ensure_communicating runs. It pops the forgotten task and discards it.
+    7. We now have a task stuck in fetch state.
+    """
+    x = c.submit(inc, 1, key="x", workers=[a.address])
+    with freeze_data_fetching(b):
+        y = c.submit(inc, x, key="y", workers=[b.address])
+        await wait_for_state("x", "fetch", b)
+        x.release()
+        y.release()
+        while s.tasks or a.tasks or b.tasks:
+            await asyncio.sleep(0.01)
+
+    x = c.submit(inc, 2, key="x", workers=[a.address])
+    y = c.submit(inc, x, key="y", workers=[b.address])
+    assert await y == 4

distributed/utils_test.py

@@ -2261,3 +2261,32 @@ def wait_for_log_line(
         if match in line:
             return line
         i += 1
+
+
+@contextmanager
+def freeze_data_fetching(w: Worker, *, jump_start: bool = False):
+    """Prevent any task from transitioning from fetch to flight on the worker while
+    inside the context, simulating a situation where the worker's network comms are
+    saturated.
+    This is not the same as setting the worker to Status=paused, which would also
+    inform the Scheduler and prevent further tasks to be enqueued on the worker.
+    Parameters
+    ----------
+    w: Worker
+        The Worker on which tasks will not transition from fetch to flight
+    jump_start: bool
+        If False, tasks will remain in fetch state after exiting the context, until
+        something else triggers ensure_communicating.
+        If True, trigger ensure_communicating on exit; this simulates e.g. an unrelated
+        worker moving out of in_flight_workers.
+    """
+    old_out_connections = w.total_out_connections
+    old_comm_threshold = w.comm_threshold_bytes
+    w.total_out_connections = 0
+    w.comm_threshold_bytes = 0
+    yield
+    w.total_out_connections = old_out_connections
+    w.comm_threshold_bytes = old_comm_threshold
+    if jump_start:
+        w.status = Status.paused
+        w.status = Status.running