Refactor the scheduler code