Address code review comments / fix minor issues