Refactor scheduler code