Refactor connection closing a bit