Scheduler worker reconnect drops messages

[gjoseph92]

When a worker connects, it sends a list of data it already has. On the scheduler side, we were handling each item and transitioning it to memory. But on each iteration, we overwrite the list of messages to send to the client and worker. So if a worker reconnected with >1 key in memory, we'd drop critical instructions to send to clients and the worker.

Progress towards #6228, #5480.

Pulled out from #6272/#6329

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

cc @crusaderky @fjetter @mrocklin

[gjoseph92]

This is @fjetter's key fix. Notice we were overwriting client_msgs, worker_msgs on each iteration.

[gjoseph92]

This is the test that probably should have caught the regression, but it didn't because Scheduler.add_worker only cares about nbytes. By just adding data, but not corresponding TaskStates, nbytes was empty in the worker's reconnection message:

distributed/distributed/worker.py

Lines 1110 to 1117 in e23f0ea

	nbytes={
	ts.key: ts.get_nbytes()
	for ts in self.tasks.values()
	# Only if the task is in memory this is a sensible
	# result since otherwise it simply submits the
	# default value
	if ts.state == "memory"
	},

The keys field of this message is unused by the scheduler and should probably be removed. It should be redundant to nbytes anyway.

[gjoseph92]

The previous test wasn't really testing much besides this. Whether x-5 would come from normal compute on worker a, or from memory on the new worker, was a race condition, and was untested anyway.

[gjoseph92]

Note that this updated test fails on main (it times out waiting here, because the client message saying l1 and l2 are done is overwritten and never sent).

[gjoseph92]

My updated test is failing locally as well. Need to look into this more next week.

[github-actions]

Unit Test Results

       15 files  ±0         15 suites  ±0   6h 53m 45s ⏱️ + 3m 6s
  2 775 tests ±0    2 694 ✔️ ±0    78 💤  - 1  3 ❌ +1 
20 587 runs  ±0  19 674 ✔️  - 5  906 💤 ±0  7 ❌ +5 

For more details on these failures, see this check.

Results for commit e23f0ea. ± Comparison against base commit 50d2911.

[fjetter]

This should no longer be relevant after #6361