Repository: dask/distributed

Issue:
P2P hash joins deadlock if workers are replaced
When removing shuffle annotations in #8329, we also removed our mechanism for aligning output worker restrictions between shuffles in a hash join. It appears that there has never been a dedicated test for this, which is why we missed it. Now, if a worker A gets replaced by a worker B, the following may happen:

When worker A gets removed, both shuffles get restarted. Shuffle 1 gets directly re-initialized with N-1 workers and range-shards based on N-1. Worker B joins the cluster, which now has N workers again. Eventually, all transfer tasks for shuffle 1 have been executed, so shuffle 2 gets initialized with N workers and range-shards based on N. There is now a mismatch in output worker restrictions between shuffle 1 and shuffle 2 that will cause the output tasks to be stuck in an infinite rescheduling loop.


Discussion:
I suggest reverting https://github.com/dask/distributed/pull/8329 until we have a better strategy
Generally +1 on reverting. I'm currently looking into this and might have something useful by EOD. Let's hold off until then and revert tomorrow morning in time for the release in case this takes longer. 
For this scenario, it is sufficient if a worker is added before the second shuffle is initialized, there's no requirement for a worker to leave.

Resolving patch:
diff --git a/distributed/shuffle/_scheduler_plugin.py b/distributed/shuffle/_scheduler_plugin.py
--- a/distributed/shuffle/_scheduler_plugin.py
+++ b/distributed/shuffle/_scheduler_plugin.py
@@ -171,6 +171,11 @@ def _raise_if_task_not_processing(self, key: Key) -> None:
     def _calculate_worker_for(self, spec: ShuffleSpec) -> dict[Any, str]:
         """Pin the outputs of a P2P shuffle to specific workers.
 
+        The P2P implementation of a hash join combines the loading of shuffled output
+        partitions for the left and right side with the actual merge operation into a
+        single output task. As a consequence, we need to make sure that shuffles with
+        shared output tasks align on the output mapping.
+
         Parameters
         ----------
         id: ID of the shuffle to pin
@@ -181,7 +186,7 @@ def _calculate_worker_for(self, spec: ShuffleSpec) -> dict[Any, str]:
             This function assumes that the barrier task and the output tasks share
             the same worker restrictions.
         """
-        mapping = {}
+        existing: dict[Any, str] = {}
         shuffle_id = spec.id
         barrier = self.scheduler.tasks[barrier_key(shuffle_id)]
 
@@ -190,10 +195,45 @@ def _calculate_worker_for(self, spec: ShuffleSpec) -> dict[Any, str]:
         else:
             workers = list(self.scheduler.workers)
 
+        # Check if this shuffle shares output tasks with a different shuffle that has
+        # already been initialized and needs to be taken into account when
+        # mapping output partitions to workers.
+        # Naively, you could delete this whole paragraph and just call
+        # spec.pick_worker; it would return two identical sets of results on both calls
+        # of this method... until the set of available workers changes between the two
+        # calls, which would cause misaligned shuffle outputs and a deadlock.
+        seen = {barrier}
+        for dependent in barrier.dependents:
+            for possible_barrier in dependent.dependencies:
+                if possible_barrier in seen:
+                    continue
+                seen.add(possible_barrier)
+                if not (other_barrier_key := id_from_key(possible_barrier.key)):
+                    continue
+                if not (shuffle := self.active_shuffles.get(other_barrier_key)):
+                    continue
+                current_worker_for = shuffle.run_spec.worker_for
+                # This is a fail-safe for future three-ways merges. At the moment there
+                # should only ever be at most one other shuffle that shares output
+                # tasks, so existing will always be empty.
+                if existing:  # pragma: nocover
+                    for shared_key in existing.keys() & current_worker_for.keys():
+                        if existing[shared_key] != current_worker_for[shared_key]:
+                            raise RuntimeError(
+                                f"Failed to initialize shuffle {spec.id} because "
+                                "it cannot align output partition mappings between "
+                                f"existing shuffles {seen}. "
+                                f"Mismatch encountered for output partition {shared_key!r}: "
+                                f"{existing[shared_key]} != {current_worker_for[shared_key]}."
+                            )
+                existing.update(current_worker_for)
+
+        worker_for = {}
         for partition in spec.output_partitions:
-            worker = spec.pick_worker(partition, workers)
-            mapping[partition] = worker
-        return mapping
+            if (worker := existing.get(partition, None)) is None:
+                worker = spec.pick_worker(partition, workers)
+            worker_for[partition] = worker
+        return worker_for
 
     def _ensure_output_tasks_are_non_rootish(self, spec: ShuffleSpec) -> None:
         """Output tasks are created without worker restrictions and run once with the


Test patch:
diff --git a/distributed/shuffle/tests/test_merge.py b/distributed/shuffle/tests/test_merge.py
--- a/distributed/shuffle/tests/test_merge.py
+++ b/distributed/shuffle/tests/test_merge.py
@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+import asyncio
 import contextlib
+from typing import Any
 from unittest import mock
 
 import pytest
 
-from distributed.shuffle._core import id_from_key
+from dask.typing import Key
+
+from distributed import Worker
+from distributed.shuffle._core import ShuffleId, ShuffleSpec, id_from_key
 from distributed.shuffle._merge import hash_join
+from distributed.shuffle._worker_plugin import ShuffleRun, _ShuffleRunManager
 from distributed.utils_test import gen_cluster
 
 dd = pytest.importorskip("dask.dataframe")
@@ -478,3 +484,51 @@ async def test_index_merge_p2p(c, s, a, b, how):
         ),
         pdf_right.merge(pdf_left, how=how, right_index=True, left_on="a"),
     )
+
+
+class LimitedGetOrCreateShuffleRunManager(_ShuffleRunManager):
+    seen: set[ShuffleId]
+    block_get_or_create: asyncio.Event
+    blocking_get_or_create: asyncio.Event
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.seen = set()
+        self.limit = 1
+        self.blocking_get_or_create = asyncio.Event()
+        self.block_get_or_create = asyncio.Event()
+
+    async def get_or_create(self, spec: ShuffleSpec, key: Key) -> ShuffleRun:
+        if len(self.seen) >= self.limit and spec.id not in self.seen:
+            self.blocking_get_or_create.set()
+            await self.block_get_or_create.wait()
+        self.seen.add(spec.id)
+        return await super().get_or_create(spec, key)
+
+
+@mock.patch(
+    "distributed.shuffle._worker_plugin._ShuffleRunManager",
+    LimitedGetOrCreateShuffleRunManager,
+)
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_merge_does_not_deadlock_if_worker_joins(c, s, a):
+    """Regression test for https://github.com/dask/distributed/issues/8411"""
+    pdf1 = pd.DataFrame({"a": range(100), "b": range(0, 200, 2)})
+    pdf2 = pd.DataFrame({"x": range(200), "y": [1, 2, 3, 4] * 50})
+    df1 = dd.from_pandas(pdf1, npartitions=10)
+    df2 = dd.from_pandas(pdf2, npartitions=20)
+
+    run_manager_A = a.plugins["shuffle"].shuffle_runs
+
+    joined = dd.merge(df1, df2, left_on="a", right_on="x", shuffle="p2p")
+    result = c.compute(joined)
+
+    await run_manager_A.blocking_get_or_create.wait()
+
+    async with Worker(s.address) as b:
+        run_manager_A.block_get_or_create.set()
+        run_manager_B = b.plugins["shuffle"].shuffle_runs
+        run_manager_B.block_get_or_create.set()
+        result = await result
+    expected = pd.merge(pdf1, pdf2, left_on="a", right_on="x")
+    assert_eq(result, expected, check_index=False)