Why not leave it the way it was? This change doesn't have anything to do with the DAGScheduler cleanup, and it risks introducing a bug. We can investigate it in a later patch.