Repository: dask/distributed

Issue:
Regression: frequent deadlocks when gather_dep fails to contact peer
`test_worker_metrics.py::test_gather_dep_network_error` has started being heavily flaky recently.
The failure has nothing to do with metrics; a `Worker.gather_dep` method that fails to open a new RPC channel to its peer seems now to be left dangling forever. 

```python
@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    config={"distributed.comm.timeouts.connect": "500ms"},
)
async def test_gather_dep_network_error(c, s, a):
    x = c.submit(inc, 1, key="x")
    await wait(x)
    async with BlockedGatherDep(s.address) as b:
        y = c.submit(inc, x, key="y", workers=[b.address])
        await b.in_gather_dep.wait()   # <-- Before b tries to connect to a
        await a.close()
        b.block_gather_dep.set()  # <-- b will now attempt and gracefully fail RPC call to a.get_data()
        await wait(y)  # <--- Times out here
```

This may be (happy to be proven wrong) a major regression that could deadlock entire clusters. I suggest treating it as a blocker for the next release.

This also shows we are missing a dedicated unit test; this regression should not have been fortuitously caught by a test that is about metrics.

The [test report](https://dask.github.io/distributed/test_short_report.html) seems to strongly correlate it to #7969:

![image](https://github.com/dask/distributed/assets/6213168/baaf3c21-08d2-4d88-9b95-a4a99cfcb881)

CC @graingert @jrbourbeau @fjetter @hendrikmakait 


Discussion:

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2081,7 +2081,10 @@ async def gather_dep(
                 stimulus_id=f"gather-dep-success-{time()}",
             )
 
-        except OSError:
+        # Note: CancelledError and asyncio.TimeoutError are rare conditions
+        # that can be raised by the network stack.
+        # See https://github.com/dask/distributed/issues/8006
+        except (OSError, asyncio.CancelledError, asyncio.TimeoutError):
             logger.exception("Worker stream died during communication: %s", worker)
             self.state.log.append(
                 ("gather-dep-failed", worker, to_gather, stimulus_id, time())
@@ -2094,6 +2097,8 @@ async def gather_dep(
 
         except Exception as e:
             # e.g. data failed to deserialize
+            # FIXME this will deadlock the cluster
+            #       https://github.com/dask/distributed/issues/6705
             logger.exception(e)
             self.state.log.append(
                 ("gather-dep-failed", worker, to_gather, stimulus_id, time())


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -3351,6 +3351,78 @@ async def test_gather_dep_no_longer_in_flight_tasks(c, s, a):
         assert not any("missing-dep" in msg for msg in f2_story)
 
 
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_gather_dep_cancelled_error(c, s, a):
+    """Something somewhere in the networking stack raises CancelledError while
+    gather_dep is running
+
+    See Also
+    --------
+    test_get_data_cancelled_error
+    https://github.com/dask/distributed/issues/8006
+    """
+    async with BlockedGetData(s.address) as b:
+        x = c.submit(inc, 1, key="x", workers=[b.address])
+        y = c.submit(inc, x, key="y", workers=[a.address])
+        await b.in_get_data.wait()
+        tasks = {
+            task for task in asyncio.all_tasks() if "gather_dep" in task.get_name()
+        }
+        assert tasks
+        # There should be only one task but cope with finding more just in case a
+        # previous test didn't properly clean up
+        for task in tasks:
+            task.cancel()
+
+        b.block_get_data.set()
+        assert await y == 3
+
+    assert_story(
+        a.state.story("x"),
+        [
+            ("x", "fetch", "flight", "flight", {}),
+            ("x", "flight", "missing", "missing", {}),
+            ("x", "missing", "fetch", "fetch", {}),
+            ("x", "fetch", "flight", "flight", {}),
+            ("x", "flight", "memory", "memory", {"y": "ready"}),
+        ],
+    )
+
+
+@gen_cluster(client=True, nthreads=[("", 1)], timeout=5)
+async def test_get_data_cancelled_error(c, s, a):
+    """Something somewhere in the networking stack raises CancelledError while
+    get_data is running
+
+    See Also
+    --------
+    test_gather_dep_cancelled_error
+    https://github.com/dask/distributed/issues/8006
+    """
+
+    class FlakyInboundRPC(Worker):
+        flake = 0
+
+        def handle_comm(self, comm):
+            if self.flake:
+                self.flake -= 1
+
+                async def write(*args, **kwargs):
+                    raise asyncio.CancelledError()
+
+                comm.write = write
+
+            return super().handle_comm(comm)
+
+    async with FlakyInboundRPC(s.address) as b:
+        x = c.submit(inc, 1, key="x", workers=[b.address])
+        await wait(x)
+        b.flake = 2
+        y = c.submit(inc, x, key="y", workers=[a.address])
+        assert await y == 3
+        assert b.flake == 0
+
+
 @gen_cluster(client=True, nthreads=[("", 1)])
 async def test_Worker__to_dict(c, s, a):
     x = c.submit(inc, 1, key="x")