Improve performance of processing incoming packets (#475)

aioesphomeapi/connection.py

@@ -306,11 +306,12 @@ class APIConnection:
         """Step 3 in connect process: initialize the frame helper and init read loop."""
         fh: Union[APIPlaintextFrameHelper, APINoiseFrameHelper]
         loop = self._loop
+        process_packet = self._process_packet_factory()
 
         if self._params.noise_psk is None:
             _, fh = await loop.create_connection(
                 lambda: APIPlaintextFrameHelper(
-                    on_pkt=self._process_packet,
+                    on_pkt=process_packet,
                     on_error=self._report_fatal_error,
                     client_info=self._params.client_info,
                 ),
@@ -321,7 +322,7 @@ class APIConnection:
                 lambda: APINoiseFrameHelper(
                     noise_psk=self._params.noise_psk,
                     expected_name=self._params.expected_name,
-                    on_pkt=self._process_packet,
+                    on_pkt=process_packet,
                     on_error=self._report_fatal_error,
                     client_info=self._params.client_info,
                 ),
@@ -699,11 +700,21 @@ class APIConnection:
         self._read_exception_futures.clear()
         self._cleanup()
 
-    def _process_packet(self, msg_type_proto: int, data: bytes) -> None:
+    def _process_packet_factory(self) -> Callable[[int, bytes], None]:
+        """Factory to make a packet processor."""
+        message_type_to_proto = MESSAGE_TYPE_TO_PROTO
+        is_enabled_for = _LOGGER.isEnabledFor
+        logging_debug = logging.DEBUG
+        message_handlers = self._message_handlers
+
+        def _process_packet(msg_type_proto: int, data: bytes) -> None:
             """Process a packet from the socket."""
-        debug = _LOGGER.isEnabledFor(logging.DEBUG)
-        if not (class_ := MESSAGE_TYPE_TO_PROTO.get(msg_type_proto)):
-            if debug:
+            try:
+                # python 3.11 has near zero cost exception handling
+                # if we do not raise which is almost never expected
+                # so we can just use a try/except here
+                class_ = message_type_to_proto[msg_type_proto]
+            except KeyError:
                 _LOGGER.debug(
                     "%s: Skipping message type %s", self.log_name, msg_type_proto
                 )
@@ -733,7 +744,7 @@ class APIConnection:
 
             msg_type = type(msg)
 
-        if debug:
+            if is_enabled_for(logging_debug):
                 _LOGGER.debug(
                     "%s: Got message of type %s: %s", self.log_name, msg_type, msg
                 )
@@ -748,7 +759,9 @@ class APIConnection:
                 # since we know the connection is still alive
                 self._send_pending_ping = False
 
-        for handler in self._message_handlers.get(msg_type, [])[:]:
+            handlers = message_handlers.get(msg_type)
+            if handlers is not None:
+                for handler in handlers[:]:
                     handler(msg)
 
             # Pre-check the message type to avoid awaiting
@@ -768,6 +781,8 @@ class APIConnection:
                 resp.epoch_seconds = int(time.time())
                 self.send_message(resp)
 
+        return _process_packet
+
     async def disconnect(self) -> None:
         """Disconnect from the API."""
         if self._connect_task:

tests/test_connection.py

@@ -54,7 +54,7 @@ def socket_socket():
 
 def _get_mock_protocol(conn: APIConnection):
     protocol = APIPlaintextFrameHelper(
-        on_pkt=conn._process_packet,
+        on_pkt=conn._process_packet_factory(),
         on_error=conn._report_fatal_error,
         client_info="mock",
     )