Cleaner implementation of UpdateShardStates + more tests

(Since we shouldn't assume pool implementation allows duplicate registration)