Fix code review issues