Optimize memory usage during broadcasts (#1233)

Author: miguelgrinberg

setup.cfg

@@ -25,7 +25,7 @@ packages = find:
 python_requires = >=3.6
 install_requires =
     bidict >= 0.21.0
-    python-engineio >= 4.3.0
+    python-engineio >= 4.7.0
 
 [options.packages.find]
 where = src

src/socketio/asyncio_manager.py

@@ -1,5 +1,7 @@
 import asyncio
 
+from engineio import packet as eio_packet
+from socketio import packet
 from .base_manager import BaseManager
 
 
@@ -17,18 +19,45 @@ async def emit(self, event, data, namespace, room=None, skip_sid=None,
         """
         if namespace not in self.rooms:
             return
-        tasks = []
+        if isinstance(data, tuple):
+            # tuples are expanded to multiple arguments, everything else is
+            # sent as a single argument
+            data = list(data)
+        elif data is not None:
+            data = [data]
+        else:
+            data = []
         if not isinstance(skip_sid, list):
             skip_sid = [skip_sid]
-        for sid, eio_sid in self.get_participants(namespace, room):
-            if sid not in skip_sid:
-                if callback is not None:
+        tasks = []
+        if not callback:
+            # when callbacks aren't used the packets sent to each recipient are
+            # identical, so they can be generated once and reused
+            pkt = self.server.packet_class(
+                packet.EVENT, namespace=namespace, data=[event] + data)
+            encoded_packet = pkt.encode()
+            if not isinstance(encoded_packet, list):
+                encoded_packet = [encoded_packet]
+            eio_pkt = [eio_packet.Packet(eio_packet.MESSAGE, p)
+                       for p in encoded_packet]
+            for sid, eio_sid in self.get_participants(namespace, room):
+                if sid not in skip_sid:
+                    for p in eio_pkt:
+                        tasks.append(asyncio.create_task(
+                            self.server._send_eio_packet(eio_sid, p)))
+        else:
+            # callbacks are used, so each recipient must be sent a packet that
+            # contains a unique callback id
+            # note that callbacks when addressing a group of people are
+            # implemented but not tested or supported
+            for sid, eio_sid in self.get_participants(namespace, room):
+                if sid not in skip_sid:  # pragma: no branch
                     id = self._generate_ack_id(sid, callback)
-                else:
-                    id = None
-                tasks.append(asyncio.create_task(
-                    self.server._emit_internal(eio_sid, event, data,
-                                               namespace, id)))
+                    pkt = self.server.packet_class(
+                        packet.EVENT, namespace=namespace, data=[event] + data,
+                        id=id)
+                    tasks.append(asyncio.create_task(
+                        self.server._send_packet(eio_sid, pkt)))
         if tasks == []:  # pragma: no cover
             return
         await asyncio.wait(tasks)

src/socketio/asyncio_server.py

@@ -424,19 +424,6 @@ async def sleep(self, seconds=0):
         """
         return await self.eio.sleep(seconds)
 
-    async def _emit_internal(self, sid, event, data, namespace=None, id=None):
-        """Send a message to a client."""
-        # tuples are expanded to multiple arguments, everything else is sent
-        # as a single argument
-        if isinstance(data, tuple):
-            data = list(data)
-        elif data is not None:
-            data = [data]
-        else:
-            data = []
-        await self._send_packet(sid, self.packet_class(
-            packet.EVENT, namespace=namespace, data=[event] + data, id=id))
-
     async def _send_packet(self, eio_sid, pkt):
         """Send a Socket.IO packet to a client."""
         encoded_packet = pkt.encode()
@@ -446,6 +433,10 @@ async def _send_packet(self, eio_sid, pkt):
         else:
             await self.eio.send(eio_sid, encoded_packet)
 
+    async def _send_eio_packet(self, eio_sid, eio_pkt):
+        """Send a raw Engine.IO packet to a client."""
+        await self.eio.send_packet(eio_sid, eio_pkt)
+
     async def _handle_connect(self, eio_sid, namespace, data):
         """Handle a client connection request."""
         namespace = namespace or '/'

src/socketio/base_manager.py

@@ -2,6 +2,8 @@
 import logging
 
 from bidict import bidict, ValueDuplicationError
+from engineio import packet as eio_packet
+from socketio import packet
 
 default_logger = logging.getLogger('socketio')
 
@@ -161,15 +163,42 @@ def emit(self, event, data, namespace, room=None, skip_sid=None,
         connected to the namespace."""
         if namespace not in self.rooms:
             return
+        if isinstance(data, tuple):
+            # tuples are expanded to multiple arguments, everything else is
+            # sent as a single argument
+            data = list(data)
+        elif data is not None:
+            data = [data]
+        else:
+            data = []
         if not isinstance(skip_sid, list):
             skip_sid = [skip_sid]
-        for sid, eio_sid in self.get_participants(namespace, room):
-            if sid not in skip_sid:
-                if callback is not None:
+        if not callback:
+            # when callbacks aren't used the packets sent to each recipient are
+            # identical, so they can be generated once and reused
+            pkt = self.server.packet_class(
+                packet.EVENT, namespace=namespace, data=[event] + data)
+            encoded_packet = pkt.encode()
+            if not isinstance(encoded_packet, list):
+                encoded_packet = [encoded_packet]
+            eio_pkt = [eio_packet.Packet(eio_packet.MESSAGE, p)
+                       for p in encoded_packet]
+            for sid, eio_sid in self.get_participants(namespace, room):
+                if sid not in skip_sid:
+                    for p in eio_pkt:
+                        self.server._send_eio_packet(eio_sid, p)
+        else:
+            # callbacks are used, so each recipient must be sent a packet that
+            # contains a unique callback id
+            # note that callbacks when addressing a group of people are
+            # implemented but not tested or supported
+            for sid, eio_sid in self.get_participants(namespace, room):
+                if sid not in skip_sid:  # pragma: no branch
                     id = self._generate_ack_id(sid, callback)
-                else:
-                    id = None
-                self.server._emit_internal(eio_sid, event, data, namespace, id)
+                    pkt = self.server.packet_class(
+                        packet.EVENT, namespace=namespace, data=[event] + data,
+                        id=id)
+                    self.server._send_packet(eio_sid, pkt)
 
     def trigger_callback(self, sid, id, data):
         """Invoke an application callback."""

src/socketio/server.py

@@ -185,7 +185,7 @@ def on(self, event, handler=None, namespace=None):
         Example usage::
 
             # as a decorator:
-            @socket_io.on('connect', namespace='/chat')
+            @sio.on('connect', namespace='/chat')
             def connect_handler(sid, environ):
                 print('Connection request')
                 if environ['REMOTE_ADDR'] in blacklisted:
@@ -194,7 +194,7 @@ def connect_handler(sid, environ):
             # as a method:
             def message_handler(sid, msg):
                 print('Received message: ', msg)
-                eio.send(sid, 'response')
+                sio.send(sid, 'response')
             socket_io.on('message', namespace='/chat', handler=message_handler)
 
         The handler function receives the ``sid`` (session ID) for the
@@ -633,19 +633,6 @@ def sleep(self, seconds=0):
         """
         return self.eio.sleep(seconds)
 
-    def _emit_internal(self, eio_sid, event, data, namespace=None, id=None):
-        """Send a message to a client."""
-        # tuples are expanded to multiple arguments, everything else is sent
-        # as a single argument
-        if isinstance(data, tuple):
-            data = list(data)
-        elif data is not None:
-            data = [data]
-        else:
-            data = []
-        self._send_packet(eio_sid, self.packet_class(
-            packet.EVENT, namespace=namespace, data=[event] + data, id=id))
-
     def _send_packet(self, eio_sid, pkt):
         """Send a Socket.IO packet to a client."""
         encoded_packet = pkt.encode()
@@ -655,6 +642,10 @@ def _send_packet(self, eio_sid, pkt):
         else:
             self.eio.send(eio_sid, encoded_packet)
 
+    def _send_eio_packet(self, eio_sid, eio_pkt):
+        """Send a raw Engine.IO packet to a client."""
+        self.eio.send_packet(eio_sid, eio_pkt)
+
     def _handle_connect(self, eio_sid, namespace, data):
         """Handle a client connection request."""
         namespace = namespace or '/'