Repository: dask/distributed

Issue:
P2P shuffle restart may deadlock if no running workers exist
An edge case exists where attempting to restart a P2P shuffle may deadlock. 

In this scenario, there are some non-running workers holding inputs to the shuffle, but no more running workers when the shuffle is restarted. This will cause a `processing` `shuffle-transfer` to transition `processing -> released -> waiting -> waiting -> no-worker -> released`. The task will then remain stuck in `released`, causing a deadlock.

There are two possible ways of fixing this: 
1. Carefully design the restarting mechanism to avoid this transition chain.
2. Recommend `waiting` when transitioning a task `no-worker -> released` if it has tasks or clients waiting on its results (similar to the `processing -> released` and `waiting -> released`.

**Reproducer:**

```python
@gen_cluster(client=True, nthreads=[])
async def test_restart_deadlock(c, s):
    async with Worker(s.address) as a:
        async with Nanny(s.address) as b:
            with dask.annotate(workers=[a.address]):
                df = dask.datasets.timeseries(
                    start="2000-01-01",
                    end="2000-03-01",
                    dtypes={"x": float, "y": float},
                    freq="10 s",
                )
            out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
            fut = c.compute(out.x.size)
            await wait_until_worker_has_tasks(
                "shuffle-transfer", b.worker_address, 1, s
            )
            a.status = Status.paused
            while len(s.running) > 1:
                await asyncio.sleep(0.01) 
            await b.process.process.kill()

            while s.running:
                await asyncio.sleep(0.01)


            a.status = Status.running
            await fut
```


Discussion:

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2707,14 +2707,10 @@ def transition_no_worker_released(self, key: str, stimulus_id: str) -> RecsMsgs:
             assert not ts.waiting_on
 
         self.unrunnable.remove(ts)
-        ts.state = "released"
-
-        for dts in ts.dependencies:
-            dts.waiters.discard(ts)
-
-        ts.waiters.clear()
 
-        return {}, {}, {}
+        recommendations: Recs = {}
+        self._propagate_released(ts, recommendations)
+        return recommendations, {}, {}
 
     def transition_waiting_queued(self, key: str, stimulus_id: str) -> RecsMsgs:
         ts = self.tasks[key]


Test patch:
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -15,6 +15,8 @@
 
 import pytest
 
+from distributed.worker import Status
+
 pd = pytest.importorskip("pandas")
 dd = pytest.importorskip("dask.dataframe")
 
@@ -421,6 +423,47 @@ async def test_crashed_worker_during_transfer(c, s, a):
         await check_scheduler_cleanup(s)
 
 
+@gen_cluster(
+    client=True,
+    nthreads=[],
+    # Effectively disable the memory monitor to be able to manually control
+    # the worker status
+    config={"distributed.worker.memory.monitor-interval": "60s"},
+)
+async def test_restarting_does_not_deadlock(c, s):
+    """Regression test for https://github.com/dask/distributed/issues/8088"""
+    async with Worker(s.address) as a:
+        async with Nanny(s.address) as b:
+            # Ensure that a holds the input tasks to the shuffle
+            with dask.annotate(workers=[a.address]):
+                df = dask.datasets.timeseries(
+                    start="2000-01-01",
+                    end="2000-03-01",
+                    dtypes={"x": float, "y": float},
+                    freq="10 s",
+                )
+            out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
+            fut = c.compute(out.x.size)
+            await wait_until_worker_has_tasks(
+                "shuffle-transfer", b.worker_address, 1, s
+            )
+            a.status = Status.paused
+            while len(s.running) > 1:
+                await asyncio.sleep(0.01)
+            b.close_gracefully()
+            await b.process.process.kill()
+
+            while s.running:
+                await asyncio.sleep(0.01)
+
+            a.status = Status.running
+
+            while not s.running:
+                await asyncio.sleep(0.01)
+            pass
+            await fut
+
+
 @gen_cluster(client=True, nthreads=[("", 1)] * 2)
 async def test_closed_input_only_worker_during_transfer(c, s, a, b):
     def mock_get_worker_for_range_sharding(