refactor and fix timing issue