Cancel websocket heartbeat on close #1793

aio-libs · Apr 8, 2017 · 9b7af75 · 9b7af75
1 parent aed0bfe
commit 9b7af75
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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,14 +201,15 @@ 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:
             self._reader.feed_data(WS_CLOSING_MESSAGE, 0)
             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)