diff --git a/CHANGES.rst b/CHANGES.rst index b0ec203aac7..bc8c24141fb 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -13,6 +13,8 @@ Changes - Allow to disable redirect url re-quoting #1474 +- Cancel websocket heartbeat on close #1793 + - Dropped "%O" in access logger #1673 diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py index e16c7c7cf95..acf19094988 100644 --- a/aiohttp/client_ws.py +++ b/aiohttp/client_ws.py @@ -64,10 +64,11 @@ def _send_heartbeat(self): self._pong_not_received, self._pong_heartbeat, self._loop) def _pong_not_received(self): - self._closed = True - self._close_code = 1006 - self._exception = asyncio.TimeoutError() - self._response.close() + if not self._closed: + self._closed = True + self._close_code = 1006 + self._exception = asyncio.TimeoutError() + self._response.close() @property def closed(self): @@ -174,14 +175,14 @@ def receive(self, timeout=None): return WS_CLOSED_MESSAGE try: + self._waiting = create_future(self._loop) try: - self._waiting = create_future(self._loop) with Timeout( timeout or self._receive_timeout, loop=self._loop): msg = yield from self._reader.read() - finally: self._reset_heartbeat() + finally: waiter = self._waiting self._waiting = None waiter.set_result(True) diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py index ad58bc73406..fde267898cf 100644 --- a/aiohttp/web_ws.py +++ b/aiohttp/web_ws.py @@ -80,11 +80,10 @@ def _send_heartbeat(self): self._pong_not_received, self._pong_heartbeat, self._loop) def _pong_not_received(self): - self._closed = True - self._close_code = 1006 - self._exception = asyncio.TimeoutError() - - if self._req is not None: + if self._req is not None and self._req.transport is not None: + self._closed = True + self._close_code = 1006 + self._exception = asyncio.TimeoutError() self._req.transport.close() @asyncio.coroutine @@ -202,6 +201,8 @@ def close(self, *, code=1000, message=b''): if self._writer is None: raise RuntimeError('Call .prepare() first') + self._cancel_heartbeat() + # we need to break `receive()` cycle first, # `close()` may be called from different task if self._waiting is not None and not self._closed: @@ -209,7 +210,6 @@ def close(self, *, code=1000, message=b''): yield from self._waiting if not self._closed: - self._cancel_heartbeat() self._closed = True try: self._writer.close(code, message) @@ -270,8 +270,8 @@ def receive(self, timeout=None): with Timeout( timeout or self._receive_timeout, loop=self._loop): msg = yield from self._reader.read() - finally: self._reset_heartbeat() + finally: waiter = self._waiting self._waiting = None waiter.set_result(True)