Better connection loss handling