Address code review feedback
CC @lizlux