Fix endless loop when disconnecting on multi-server deployments (Fixes #441)

Author: miguelgrinberg

socketio/asyncio_pubsub_manager.py

@@ -70,17 +70,13 @@ async def emit(self, event, data, namespace=None, room=None, skip_sid=None,
                              'host_id': self.host_id})
 
     async def can_disconnect(self, sid, namespace):
-        await self._publish({'method': 'disconnect', 'sid': sid,
-                             'namespace': namespace or '/'})
-
-    async def disconnect(self, sid, namespace=None):
-        """Disconnect a client."""
-        # this is a bit weird, the can_disconnect call on pubsub managers just
-        # issues a disconnect request to the message queue and returns None,
-        # indicating that the client cannot disconnect immediately. The
-        # server(s) listening on the queue will get this request and carry out
-        # the disconnect appropriately.
-        await self.can_disconnect(sid, namespace)
+        if self.is_connected(sid, namespace):
+            # client is in this server, so we can disconnect directly
+            return super().can_disconnect(sid, namespace)
+        else:
+            # client is in another server, so we post request to the queue
+            await self._publish({'method': 'disconnect', 'sid': sid,
+                                'namespace': namespace or '/'})
 
     async def close_room(self, room, namespace=None):
         await self._publish({'method': 'close_room', 'room': room,

socketio/asyncio_server.py

@@ -326,8 +326,11 @@ async def disconnect(self, sid, namespace=None, ignore_queue=False):
         Note: this method is a coroutine.
         """
         namespace = namespace or '/'
-        if (ignore_queue and self.manager.is_connected(sid, namespace)) or \
-                await self.manager.can_disconnect(sid, namespace):
+        if ignore_queue:
+            do_it = self.manager.is_connected(sid, namespace)
+        else:
+            do_it = await self.manager.can_disconnect(sid, namespace)
+        if do_it:
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
             self.manager.pre_disconnect(sid, namespace=namespace)
             await self._send_packet(sid, packet.Packet(packet.DISCONNECT,
@@ -440,9 +443,11 @@ async def _handle_disconnect(self, sid, namespace):
             namespace_list = [namespace]
         for n in namespace_list:
             if n != '/' and self.manager.is_connected(sid, n):
+                self.manager.pre_disconnect(sid, namespace=namespace)
                 await self._trigger_event('disconnect', n, sid)
                 self.manager.disconnect(sid, n)
         if namespace == '/' and self.manager.is_connected(sid, namespace):
+            self.manager.pre_disconnect(sid, namespace=namespace)
             await self._trigger_event('disconnect', '/', sid)
             self.manager.disconnect(sid, '/')
 

socketio/pubsub_manager.py

@@ -68,17 +68,13 @@ def emit(self, event, data, namespace=None, room=None, skip_sid=None,
                        'host_id': self.host_id})
 
     def can_disconnect(self, sid, namespace):
-        self._publish({'method': 'disconnect', 'sid': sid,
-                       'namespace': namespace or '/'})
-
-    def disconnect(self, sid, namespace=None):
-        """Disconnect a client."""
-        # this is a bit weird, the can_disconnect call on pubsub managers just
-        # issues a disconnect request to the message queue and returns None,
-        # indicating that the client cannot disconnect immediately. The
-        # server(s) listening on the queue will get this request and carry out
-        # the disconnect appropriately.
-        self.can_disconnect(sid, namespace)
+        if self.is_connected(sid, namespace):
+            # client is in this server, so we can disconnect directly
+            return super().can_disconnect(sid, namespace)
+        else:
+            # client is in another server, so we post request to the queue
+            self._publish({'method': 'disconnect', 'sid': sid,
+                           'namespace': namespace or '/'})
 
     def close_room(self, room, namespace=None):
         self._publish({'method': 'close_room', 'room': room,

socketio/server.py

@@ -517,8 +517,11 @@ def disconnect(self, sid, namespace=None, ignore_queue=False):
                              its default value of ``False``.
         """
         namespace = namespace or '/'
-        if (ignore_queue and self.manager.is_connected(sid, namespace)) or \
-                self.manager.can_disconnect(sid, namespace):
+        if ignore_queue:
+            do_it = self.manager.is_connected(sid, namespace)
+        else:
+            do_it = self.manager.can_disconnect(sid, namespace)
+        if do_it:
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
             self.manager.pre_disconnect(sid, namespace=namespace)
             self._send_packet(sid, packet.Packet(packet.DISCONNECT,
@@ -649,9 +652,11 @@ def _handle_disconnect(self, sid, namespace):
             namespace_list = [namespace]
         for n in namespace_list:
             if n != '/' and self.manager.is_connected(sid, n):
+                self.manager.pre_disconnect(sid, namespace=namespace)
                 self._trigger_event('disconnect', n, sid)
                 self.manager.disconnect(sid, n)
         if namespace == '/' and self.manager.is_connected(sid, namespace):
+            self.manager.pre_disconnect(sid, namespace=namespace)
             self._trigger_event('disconnect', '/', sid)
             self.manager.disconnect(sid, '/')
 

tests/asyncio/test_asyncio_pubsub_manager.py

@@ -116,8 +116,10 @@ def test_emit_with_ignore_queue(self):
         self.pm.server._emit_internal.mock.assert_called_once_with(
             '123', 'foo', 'bar', '/', None)
 
-    def test_disconnect(self):
-        _run(self.pm.disconnect('123', '/foo'))
+    def test_can_disconnect(self):
+        self.pm.connect('123', '/')
+        self.assertTrue(_run(self.pm.can_disconnect('123', '/')))
+        _run(self.pm.can_disconnect('123', '/foo'))
         self.pm._publish.mock.assert_called_once_with(
             {'method': 'disconnect', 'sid': '123', 'namespace': '/foo'})
 

tests/common/test_pubsub_manager.py

@@ -112,8 +112,10 @@ def test_emit_with_ignore_queue(self):
         self.pm.server._emit_internal.assert_called_once_with('123', 'foo',
                                                               'bar', '/', None)
 
-    def test_disconnect(self):
-        self.pm.disconnect('123', '/foo')
+    def test_can_disconnect(self):
+        self.pm.connect('123', '/')
+        self.assertTrue(self.pm.can_disconnect('123', '/'))
+        self.pm.can_disconnect('123', '/foo')
         self.pm._publish.assert_called_once_with(
             {'method': 'disconnect', 'sid': '123', 'namespace': '/foo'})