Cleaner implementation of UpdateShardStates + more tests

(Since we shouldn't assume pool implementation allows duplicate registration)
why is this cleaner than other way around?