Refactor logging code