code cleanup and improved partial fragment removal

tkulin · tkulin · commit 04ae89baf761 · 2025-03-03T11:13:55.000-05:00
diff --git a/bellows/ezsp/fragmentation.py b/bellows/ezsp/fragmentation.py
@@ -1,12 +1,10 @@
-"""
-Implements APS fragmentation reassembly on the EZSP Host side,
+"""Implements APS fragmentation reassembly on the EZSP Host side,
 mirroring the logic from fragmentation.c in the EmberZNet stack.
 """
 
 import asyncio
 import logging
-from collections import defaultdict
-from typing import Optional, Dict, Tuple
+from typing import Dict, Optional, Tuple
 
 LOGGER = logging.getLogger(__name__)
 
@@ -17,13 +15,14 @@
 # store partial data keyed by (sender, aps_sequence, profile_id, cluster_id)
 FragmentKey = Tuple[int, int, int, int]
 
+
 class _FragmentEntry:
     def __init__(self, fragment_count: int):
         self.fragment_count = fragment_count
         self.fragments_received = 0
         self.fragment_data = {}
         self.start_time = asyncio.get_event_loop().time()
-    
+
     def add_fragment(self, index: int, data: bytes) -> None:
         if index not in self.fragment_data:
             self.fragment_data[index] = data
@@ -33,32 +32,41 @@ def is_complete(self) -> bool:
         return self.fragments_received == self.fragment_count
 
     def assemble(self) -> bytes:
-        return b''.join(self.fragment_data[i] for i in sorted(self.fragment_data.keys()))
+        return b"".join(
+            self.fragment_data[i] for i in sorted(self.fragment_data.keys())
+        )
+
 
 class FragmentManager:
     def __init__(self):
         self._partial: Dict[FragmentKey, _FragmentEntry] = {}
-    
-    def handle_incoming_fragment(self, sender_nwk: int, aps_sequence: int, profile_id: int, cluster_id: int,
-                                 group_id: int, payload: bytes) -> Tuple[bool, Optional[bytes], int, int]:
-        """
-        Handle a newly received fragment. The group_id field
-        encodes high byte = total fragment count, low byte = current fragment index.
+        self._cleanup_timers: Dict[FragmentKey, asyncio.TimerHandle] = {}
+
+    def handle_incoming_fragment(
+        self,
+        sender_nwk: int,
+        aps_sequence: int,
+        profile_id: int,
+        cluster_id: int,
+        fragment_count: int,
+        fragment_index: int,
+        payload: bytes,
+    ) -> Tuple[bool, Optional[bytes], int, int]:
+        """Handle a newly received fragment.
 
         :param sender_nwk: NWK address or the short ID of the sender.
         :param aps_sequence: The APS sequence from the incoming APS frame.
         :param profile_id: The APS frame's profileId.
         :param cluster_id: The APS frame's clusterId.
-        :param group_id: The APS frame's groupId (used to store fragment # / total).
+        :param fragment_count: The total number of expected message fragments.
+        :param fragment_index: The index of the current fragment being processed.
         :param payload: The fragment of data for this message.
         :return: (complete, reassembled_data, fragment_count, fragment_index)
                  complete = True if we have all fragments now, else False
                  reassembled_data = the final complete payload (bytes) if complete is True
                  fragment_coutn = the total number of fragments holding the complete packet
                  fragment_index = the index of the current received fragment
         """
-        fragment_count = (group_id >> 8) & 0xFF
-        fragment_index = group_id & 0xFF
 
         key: FragmentKey = (sender_nwk, aps_sequence, profile_id, cluster_id)
 
@@ -69,30 +77,44 @@ def handle_incoming_fragment(self, sender_nwk: int, aps_sequence: int, profile_i
         else:
             entry = self._partial[key]
 
-        LOGGER.debug("Received fragment %d/%d from %s (APS seq=%d, cluster=0x%04X)",
-                     fragment_index, fragment_count, sender_nwk, aps_sequence, cluster_id)
+        LOGGER.debug(
+            "Received fragment %d/%d from %s (APS seq=%d, cluster=0x%04X)",
+            fragment_index + 1,
+            fragment_count,
+            sender_nwk,
+            aps_sequence,
+            cluster_id,
+        )
 
         entry.add_fragment(fragment_index, payload)
 
+        loop = asyncio.get_running_loop()
+        self._cleanup_timers[key] = loop.call_later(
+            FRAGMENT_TIMEOUT, self.cleanup_partial, key
+        )
+
         if entry.is_complete():
             reassembled = entry.assemble()
             del self._partial[key]
-            LOGGER.debug("Message reassembly complete. Total length=%d", len(reassembled))
+            timer = self._cleanup_timers.pop(key, None)
+            if timer:
+                timer.cancel()
+            LOGGER.debug(
+                "Message reassembly complete. Total length=%d", len(reassembled)
+            )
             return (True, reassembled, fragment_count, fragment_index)
         else:
             return (False, None, fragment_count, fragment_index)
 
-    def cleanup_expired(self) -> None:
+    def cleanup_partial(self, key: FragmentKey):
+        # Called when FRAGMENT_TIMEOUT passes with no new fragments for that key.
+        LOGGER.debug(
+            "Timeout for partial reassembly of fragmented message, discarding key=%s",
+            key,
+        )
+        self._partial.pop(key, None)
+        self._cleanup_timers.pop(key, None)
 
-        now = asyncio.get_event_loop().time()
-        to_remove = []
-        for k, entry in self._partial.items():
-            if now - entry.start_time > FRAGMENT_TIMEOUT:
-                to_remove.append(k)
-        for k in to_remove:
-            del self._partial[k]
-            LOGGER.debug("Removed stale fragment reassembly for key=%s", k)
 
 # Create a single global manager instance
 fragment_manager = FragmentManager()
-
diff --git a/bellows/ezsp/protocol.py b/bellows/ezsp/protocol.py
@@ -53,7 +53,6 @@ def __init__(self, cb_handler: Callable, gateway: Gateway) -> None:
 
         # Cached by `set_extended_timeout` so subsequent calls are a little faster
         self._address_table_size: int | None = None
-        self._cleanup_fragments_periodically()
 
     def _ezsp_frame(self, name: str, *args: Any, **kwargs: Any) -> bytes:
         """Serialize the named frame and data."""
@@ -182,27 +181,41 @@ def __call__(self, data: bytes) -> None:
         if data:
             LOGGER.debug("Frame contains trailing data: %s", data)
 
-        if frame_name == "incomingMessageHandler" and result[1].options & 0x8000:  # incoming message with APS_OPTION_FRAGMENT raised
+        if (
+            frame_name == "incomingMessageHandler" and result[1].options & 0x8000
+        ):  # incoming message with APS_OPTION_FRAGMENT raised
             from bellows.ezsp.fragmentation import fragment_manager
 
             # Extract received APS frame and sender
             aps_frame = result[1]
-            sender = result[4] 
+            sender = result[4]
 
             group_id = aps_frame.groupId
             profile_id = aps_frame.profileId
             cluster_id = aps_frame.clusterId
             aps_seq = aps_frame.sequence
-            
-            complete, reassembled, frag_count, frag_index = fragment_manager.handle_incoming_fragment(
+
+            fragment_count = (group_id >> 8) & 0xFF
+            fragment_index = group_id & 0xFF
+
+            (
+                complete,
+                reassembled,
+                frag_count,
+                frag_index,
+            ) = fragment_manager.handle_incoming_fragment(
                 sender_nwk=sender,
                 aps_sequence=aps_seq,
                 profile_id=profile_id,
                 cluster_id=cluster_id,
-                group_id=group_id,
-                payload=result[7]
+                fragment_count=fragment_count,
+                fragment_index=fragment_index,
+                payload=result[7],
             )
-            asyncio.create_task(self._send_fragment_ack(sender, aps_frame, frag_count, frag_index))  # APS Ack
+            ack_task = asyncio.create_task(
+                self._send_fragment_ack(sender, aps_frame, frag_count, frag_index)
+            )  # APS Ack
+            ack_task.add_done_callback(self._ack_tasks.remove)
 
             if not complete:
                 # Do not pass partial data up the stack
@@ -211,8 +224,10 @@ def __call__(self, data: bytes) -> None:
             else:
                 # Replace partial data with fully reassembled data
                 result[7] = reassembled
-                
-                LOGGER.debug("Reassembled fragmented message. Proceeding with normal handling.")
+
+                LOGGER.debug(
+                    "Reassembled fragmented message. Proceeding with normal handling."
+                )
 
         if sequence in self._awaiting:
             expected_id, schema, future = self._awaiting.pop(sequence)
@@ -238,27 +253,32 @@ def __call__(self, data: bytes) -> None:
         else:
             self._handle_callback(frame_name, result)
 
-    async def _send_fragment_ack(self, sender: int, incoming_aps: t.EmberApsFrame, fragment_count: int, fragment_index: int):
-       
+    async def _send_fragment_ack(
+        self,
+        sender: int,
+        incoming_aps: t.EmberApsFrame,
+        fragment_count: int,
+        fragment_index: int,
+    ) -> t.EmberStatus:
         ackFrame = t.EmberApsFrame(
             profileId=incoming_aps.profileId,
             clusterId=incoming_aps.clusterId,
             sourceEndpoint=incoming_aps.destinationEndpoint,
             destinationEndpoint=incoming_aps.sourceEndpoint,
             options=incoming_aps.options,
             groupId=((0xFF00) | (fragment_index & 0xFF)),
-            sequence=incoming_aps.sequence
+            sequence=incoming_aps.sequence,
+        )
+
+        LOGGER.debug(
+            "Sending fragment ack to 0x%04X for fragment index=%d/%d",
+            sender,
+            fragment_index + 1,
+            fragment_count,
         )
+        status = await self.sendReply(sender, ackFrame, b"")
+        return status
 
-        LOGGER.debug("Sending fragment ack to 0x%04X for fragment index=%d/%d", sender, fragment_index, fragment_count)
-        await self.sendReply(sender, ackFrame, b'')
-       
-    async def _cleanup_fragments_periodically(self):
-        from bellows.ezsp.fragmentation import fragment_manager
-        while True:
-            await asyncio.sleep(5)
-            fragment_manager.cleanup_expired()
-        
     def __getattr__(self, name: str) -> Callable:
         if name not in self.COMMANDS:
             raise AttributeError(f"{name} not found in COMMANDS")
