More robust handling of a closing connection (Fixes #164)

miguelgrinberg · Mar 2, 2020 · bdd584b · bdd584b
1 parent 658e811
commit bdd584b
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 6 deletions.
diff --git a/engineio/asyncio_socket.py b/engineio/asyncio_socket.py
@@ -20,11 +20,16 @@ async def poll(self):
             raise exceptions.QueueEmpty()
         if packets == [None]:
             return []
-        try:
-            packets.append(self.queue.get_nowait())
-            self.queue.task_done()
-        except asyncio.QueueEmpty:
-            pass
+        while True:
+            try:
+                pkt = self.queue.get_nowait()
+                self.queue.task_done()
+                if pkt is None:
+                    self.queue.put_nowait(None)
+                    break
+                packets.append(pkt)
+            except asyncio.QueueEmpty:
+                break
         return packets
 
     async def receive(self, pkt):

diff --git a/engineio/socket.py b/engineio/socket.py
@@ -36,8 +36,12 @@ def poll(self):
             return []
         while True:
             try:
-                packets.append(self.queue.get(block=False))
+                pkt = self.queue.get(block=False)
                 self.queue.task_done()
+                if pkt is None:
+                    self.queue.put(None)
+                    break
+                packets.append(pkt)
             except queue_empty:
                 break
         return packets

diff --git a/tests/asyncio/test_asyncio_socket.py b/tests/asyncio/test_asyncio_socket.py
@@ -94,6 +94,15 @@ def test_poll_none(self):
         _run(s.queue.put(None))
         self.assertEqual(_run(s.poll()), [])
 
+    def test_poll_none_after_packet(self):
+        mock_server = self._get_mock_server()
+        s = asyncio_socket.AsyncSocket(mock_server, 'sid')
+        pkt = packet.Packet(packet.MESSAGE, data='hello')
+        _run(s.send(pkt))
+        _run(s.queue.put(None))
+        self.assertEqual(_run(s.poll()), [pkt])
+        self.assertEqual(_run(s.poll()), [])
+
     def test_ping_pong(self):
         mock_server = self._get_mock_server()
         s = asyncio_socket.AsyncSocket(mock_server, 'sid')

diff --git a/tests/common/test_socket.py b/tests/common/test_socket.py
@@ -77,6 +77,21 @@ def test_poll(self):
         s.send(pkt2)
         self.assertEqual(s.poll(), [pkt1, pkt2])
 
+    def test_poll_none(self):
+        mock_server = self._get_mock_server()
+        s = socket.Socket(mock_server, 'sid')
+        s.queue.put(None)
+        self.assertEqual(s.poll(), [])
+
+    def test_poll_none_after_packet(self):
+        mock_server = self._get_mock_server()
+        s = socket.Socket(mock_server, 'sid')
+        pkt = packet.Packet(packet.MESSAGE, data='hello')
+        s.send(pkt)
+        s.queue.put(None)
+        self.assertEqual(s.poll(), [pkt])
+        self.assertEqual(s.poll(), [])
+
     def test_ping_pong(self):
         mock_server = self._get_mock_server()
         s = socket.Socket(mock_server, 'sid')