Created by: zhaify

Refactor the code from PR 16865

please help to review @luotao1 @wzzju @wojtuss @bingyanghuang