Fix wait statement hang in some cases introduced by previous commit