Make ShuffleSchedulerExtension.remove_worker more robust

[hendrikmakait]

If processing the transactions causes a task to get released, this removes the shuffle from self.states. Therefore, we must process them outside of the loop.

Under normal circumstances, remove_worker does not cause tasks to transition to released, so writing a test for this is hard. However, it can happen if you really screw with your cluster:

Traceback (most recent call last):
  File "/opt/coiled/env/lib/python3.9/site-packages/distributed/scheduler.py", line 5057, in remove_worker
    result = plugin.remove_worker(scheduler=self, worker=address)
  File "/opt/coiled/env/lib/python3.9/site-packages/distributed/shuffle/_scheduler_extension.py", line 297, in remove_worker
    for shuffle_id, shuffle in self.states.items():
RuntimeError: dictionary changed size during iteration

• Tests added / passed
• Passes pre-commit run --all-files

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       20 files  +       2         20 suites  +2   12h 26m 51s ⏱️ + 2h 15m 18s
  3 681 tests ±       0    3 568 ✔️ +       7     108 💤  -   10  5 ❌ +3 
35 600 runs  +3 644  33 831 ✔️ +3 505  1 764 💤 +136  5 ❌ +3 

For more details on these failures, see this check.

Results for commit 1088a9d. ± Comparison against base commit 4d8dbad.