put clients in a pre-disconnect state while their disconnect handler runs

This avoids potential endless recursion.
See miguelgrinberg/Flask-SocketIO#312

Author: miguelgrinberg

socketio/base_manager.py

@@ -16,6 +16,7 @@ def __init__(self):
         self.server = None
         self.rooms = {}
         self.callbacks = {}
+        self.pending_disconnect = {}
 
     def initialize(self, server):
         self.server = server
@@ -35,11 +36,26 @@ def connect(self, sid, namespace):
         self.enter_room(sid, namespace, sid)
 
     def is_connected(self, sid, namespace):
+        if namespace in self.pending_disconnect and \
+                sid in self.pending_disconnect[namespace]:
+            # the client is in the process of being disconnected
+            return False
         try:
             return self.rooms[namespace][None][sid]
         except KeyError:
             pass
 
+    def pre_disconnect(self, sid, namespace):
+        """Put the client in the to-be-disconnected list.
+
+        This allows the client data structures to be present while the
+        disconnect handler is invoked, but still recognize the fact that the
+        client is soon going away.
+        """
+        if namespace not in self.pending_disconnect:
+            self.pending_disconnect[namespace] = []
+        self.pending_disconnect[namespace].append(sid)
+
     def disconnect(self, sid, namespace):
         """Register a client disconnect from a namespace."""
         rooms = []
@@ -52,6 +68,11 @@ def disconnect(self, sid, namespace):
             del self.callbacks[sid][namespace]
             if len(self.callbacks[sid]) == 0:
                 del self.callbacks[sid]
+        if namespace in self.pending_disconnect and \
+                sid in self.pending_disconnect[namespace]:
+            self.pending_disconnect[namespace].remove(sid)
+            if len(self.pending_disconnect[namespace]) == 0:
+                del self.pending_disconnect[namespace]
 
     def enter_room(self, sid, namespace, room):
         """Add a client to a room."""

socketio/server.py

@@ -298,6 +298,7 @@ def disconnect(self, sid, namespace=None):
         namespace = namespace or '/'
         if self.manager.is_connected(sid, namespace=namespace):
             self.logger.info('Disconnecting %s [%s]', sid, namespace)
+            self.manager.pre_disconnect(sid, namespace=namespace)
             self._send_packet(sid, packet.Packet(packet.DISCONNECT,
                                                  namespace=namespace))
             self._trigger_event('disconnect', namespace, sid)

tests/test_base_manager.py

@@ -24,6 +24,20 @@ def test_connect(self):
         self.assertEqual(self.bm.rooms['/foo'], {None: {'123': True},
                                                  '123': {'123': True}})
 
+    def test_pre_disconnect(self):
+        self.bm.connect('123', '/foo')
+        self.bm.connect('456', '/foo')
+        self.bm.pre_disconnect('123', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {'/foo': ['123']})
+        self.assertFalse(self.bm.is_connected('123', '/foo'))
+        self.bm.pre_disconnect('456', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {'/foo': ['123', '456']})
+        self.assertFalse(self.bm.is_connected('456', '/foo'))
+        self.bm.disconnect('123', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {'/foo': ['456']})
+        self.bm.disconnect('456', '/foo')
+        self.assertEqual(self.bm.pending_disconnect, {})
+
     def test_disconnect(self):
         self.bm.connect('123', '/foo')
         self.bm.connect('456', '/foo')