Avoid race condition in lazy dispatch registration

[jrbourbeau]

This makes it so we delete the toplevel entry from the _lazy lookup dictionary only after we're registered the corresponding dispatch method

Closes #9181, closes dask/distributed#7061

[jrbourbeau]

@zklaus I'm curious if this fixes the issue you were encountering

[zklaus]

I'll give it a try and get back to you.

[jakirkham]

Wonder if we could actually simplify this a bit

Suggested change

	try:
	register = self._lazy.pop(toplevel)
	register = self._lazy[toplevel]
	except KeyError:
	pass
	else:
	register()
	self._lazy.pop(toplevel, None)
	register = self._lazy.get(toplevel)
	if register is not None:
	register()
	self._lazy.pop(toplevel, None)

[zklaus]

lgtm

[jrbourbeau]

Though I agree this suggestion is equivalent to the current state of this PR, I actually find the try/except to be a bit more clear as catching a KeyError is more explicit in saying "this thing doesn't have this key" than get + is not None (as there could be a key with a value of None). Regardless, this isn't a big deal either way. If this is a blocking comment, then I'm happy to update accordingly.

[jakirkham]

There are other ways to do this if we don't want to rely on None (though I don't know how None would work as an intentional value given we need to call it like a function later, which would break).

Suggested change

	try:
	register = self._lazy.pop(toplevel)
	register = self._lazy[toplevel]
	except KeyError:
	pass
	else:
	register()
	self._lazy.pop(toplevel, None)
	missing = object()
	register = self._lazy.get(toplevel, missing)
	if register is not missing:
	register()
	self._lazy.pop(toplevel, None)

[jrbourbeau]

Sure, but I still find the try/except KeyError: more clear over defining a separate missing object. Again, this is a subjective point and probably not a big deal either way. Is this update needed to get this PR in? If so, happy to update

[jakirkham]

No I don't think it matters. Have just been bitten by unexpected else behavior in try clauses. This combined with the race condition was what motivated the attempt to simplify. That said, given the existing code works for users, would prefer getting the fix in.

[jrbourbeau]

Thanks for trying this out @zklaus and for reviewing @jakirkham

[jakirkham]

Thanks for working on this James! 🙏