fix: set change during iteration when dispatching listeners (#1370)

bdraco · web-flow · commit e9f8aa5741ae · 2024-04-12T16:47:46.000-05:00
diff --git a/src/zeroconf/_handlers/record_manager.py b/src/zeroconf/_handlers/record_manager.py
@@ -56,7 +56,7 @@ def async_updates(self, now: _float, records: List[RecordUpdate]) -> None:
 
         This method will be run in the event loop.
         """
-        for listener in self.listeners:
+        for listener in self.listeners.copy():
             listener.async_update_records(self.zc, now, records)
 
     def async_updates_complete(self, notify: bool) -> None:
@@ -67,7 +67,7 @@ def async_updates_complete(self, notify: bool) -> None:
 
         This method will be run in the event loop.
         """
-        for listener in self.listeners:
+        for listener in self.listeners.copy():
             listener.async_update_records_complete()
         if notify:
             self.zc.async_notify_all()
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -1762,3 +1762,77 @@ def async_update_records(self, zc: 'Zeroconf', now: float, records: List[r.Recor
     )
 
     await aiozc.async_close()
+
+
+@pytest.mark.asyncio
+async def test_async_updates_iteration_safe():
+    """Ensure we can safely iterate over the async_updates."""
+
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    zc: Zeroconf = aiozc.zeroconf
+    updated = []
+    good_bye_answer = r.DNSPointer(
+        "myservicelow_tcp._tcp.local.",
+        const._TYPE_PTR,
+        const._CLASS_IN | const._CLASS_UNIQUE,
+        0,
+        'goodbye.local.',
+    )
+
+    class OtherListener(r.RecordUpdateListener):
+        """A RecordUpdateListener that does not implement update_records."""
+
+        def async_update_records(self, zc: 'Zeroconf', now: float, records: List[r.RecordUpdate]) -> None:
+            """Update multiple records in one shot."""
+            updated.extend(records)
+
+    other = OtherListener()
+
+    class ListenerThatAddsListener(r.RecordUpdateListener):
+        """A RecordUpdateListener that does not implement update_records."""
+
+        def async_update_records(self, zc: 'Zeroconf', now: float, records: List[r.RecordUpdate]) -> None:
+            """Update multiple records in one shot."""
+            updated.extend(records)
+            zc.async_add_listener(other, None)
+
+    zc.async_add_listener(ListenerThatAddsListener(), None)
+    await asyncio.sleep(0)  # flush out any call soons
+
+    # This should not raise RuntimeError: set changed size during iteration
+    zc.record_manager.async_updates(
+        now=current_time_millis(), records=[r.RecordUpdate(good_bye_answer, None)]
+    )
+
+    assert len(updated) == 1
+    await aiozc.async_close()
+
+
+@pytest.mark.asyncio
+async def test_async_updates_complete_iteration_safe():
+    """Ensure we can safely iterate over the async_updates_complete."""
+
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    zc: Zeroconf = aiozc.zeroconf
+
+    class OtherListener(r.RecordUpdateListener):
+        """A RecordUpdateListener that does not implement update_records."""
+
+        def async_update_records_complete(self) -> None:
+            """Update multiple records in one shot."""
+
+    other = OtherListener()
+
+    class ListenerThatAddsListener(r.RecordUpdateListener):
+        """A RecordUpdateListener that does not implement update_records."""
+
+        def async_update_records_complete(self) -> None:
+            """Update multiple records in one shot."""
+            zc.async_add_listener(other, None)
+
+    zc.async_add_listener(ListenerThatAddsListener(), None)
+    await asyncio.sleep(0)  # flush out any call soons
+
+    # This should not raise RuntimeError: set changed size during iteration
+    zc.record_manager.async_updates_complete(False)
+    await aiozc.async_close()
