Try to fix dropped data in P2P shuffle

[gjoseph92]

Attempt with @graingert at fixing #6277; doesn't work.

Also fixes a couple other small design issues (ShuffleExtension peeking into implementation details of Shuffle during shuffle_receive/shuffle_inputs_done.)

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

[graingert]

I think it's nicer to combine consecutive sync calls into one async def:

async def close():
    await shuffle.close()
    await self.worker.scheduler.shuffle_register_complete(id=shuffle_id, worker=self.worker.address)

sync(self.worker.loop, close)

[gjoseph92]

Agreed, I was just lazy here

[github-actions]

Unit Test Results

       16 files  ±0         16 suites  ±0   7h 56m 45s ⏱️ + 35m 11s
  2 756 tests ±0    2 673 ✔️  -   1       78 💤 ±0    4 ❌ ±0  1 🔥 +1 
22 010 runs  ±0  20 975 ✔️  - 13  1 022 💤 +4  12 ❌ +8  1 🔥 +1 

For more details on these failures and errors, see this check.

Results for commit cdbddc4. ± Comparison against base commit 7bd6442.

[graingert]

see #6279

[fjetter]

Next time something like this happens, please provide a bit of information why it doesn't work. A reproducer of the issue would also be welcome.

[hendrikmakait]

This should be superseded by #7486 and previous changes.