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
@@ -4804,7 +4804,7 @@ def stimulus_task_finished(self, key, worker, stimulus_id, run_id, **kwargs):
 
         ws: WorkerState = self.workers[worker]
         ts: TaskState = self.tasks.get(key)
-        if ts is None or ts.state in ("released", "queued"):
+        if ts is None or ts.state in ("released", "queued", "no-worker"):
             logger.debug(
                 "Received already computed task, worker: %s, state: %s"
                 ", key: %s, who_has: %s",


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
@@ -448,19 +448,15 @@ async def test_restarting_does_not_deadlock(c, s):
                 "shuffle-transfer", b.worker_address, 1, s
             )
             a.status = Status.paused
-            while len(s.running) > 1:
-                await asyncio.sleep(0.01)
+            await async_poll_for(lambda: len(s.running) == 1, timeout=5)
             b.close_gracefully()
             await b.process.process.kill()
 
-            while s.running:
-                await asyncio.sleep(0.01)
+            await async_poll_for(lambda: not s.running, timeout=5)
 
             a.status = Status.running
 
-            while not s.running:
-                await asyncio.sleep(0.01)
-            pass
+            await async_poll_for(lambda: s.running, timeout=5)
             await fut