Simplify state management in P2P's `ShuffleSchedulerPlugin`

As mentioned in https://github.com/dask/distributed/pull/7970#discussion_r1266814524, keeping the internal state of the plugin is currently done manually. This is prone to errors, increases the cognitive load during maintenance, and makes the code generally harder to understand. We should refactor it so that it's easier to understand and maintain invariants.