Do not block for broadcast and multicast requests

Author: puddly

tests/api/test_request.py

@@ -92,7 +92,8 @@ async def test_callback_rsp_cleanup_timeout_external(connected_znp):
     assert not znp._listeners
 
 
-async def test_callback_rsp_cleanup_timeout_internal(connected_znp):
+@pytest.mark.parametrize("background", [False, True])
+async def test_callback_rsp_cleanup_timeout_internal(background, connected_znp):
     znp, znp_server = connected_znp
     znp._config[conf.CONF_ZNP_CONFIG][conf.CONF_SREQ_TIMEOUT] = 0.1
     znp._config[conf.CONF_ZNP_CONFIG][conf.CONF_ARSP_TIMEOUT] = 0.1
@@ -104,6 +105,45 @@ async def test_callback_rsp_cleanup_timeout_internal(connected_znp):
         await znp.request_callback_rsp(
             request=c.Util.TimeAlive.Req(),
             callback=c.SYS.ResetInd.Callback(partial=True),
+            background=background,
+        )
+
+    # We should be cleaned up
+    assert not znp._listeners
+
+
+async def test_callback_rsp_cleanup_background_error(connected_znp):
+    znp, znp_server = connected_znp
+    znp._config[conf.CONF_ZNP_CONFIG][conf.CONF_SREQ_TIMEOUT] = 0.1
+    znp._config[conf.CONF_ZNP_CONFIG][conf.CONF_ARSP_TIMEOUT] = 0.1
+
+    assert not znp._listeners
+
+    # This request will timeout because we didn't send anything back
+    with pytest.raises(asyncio.TimeoutError):
+        await znp.request_callback_rsp(
+            request=c.Util.TimeAlive.Req(),
+            callback=c.SYS.ResetInd.Callback(partial=True),
+            background=True,
+        )
+
+    # We should be cleaned up
+    assert not znp._listeners
+
+
+async def test_callback_rsp_cleanup_background_timeout(connected_znp):
+    znp, znp_server = connected_znp
+    znp._config[conf.CONF_ZNP_CONFIG][conf.CONF_SREQ_TIMEOUT] = 0.1
+    znp._config[conf.CONF_ZNP_CONFIG][conf.CONF_ARSP_TIMEOUT] = 0.1
+
+    assert not znp._listeners
+
+    # This request will timeout because we didn't send anything back
+    with pytest.raises(asyncio.TimeoutError):
+        await znp.request_callback_rsp(
+            request=c.Util.TimeAlive.Req(),
+            callback=c.SYS.ResetInd.Callback(partial=True),
+            background=True,
         )
 
     # We should be cleaned up

tests/application/test_requests.py

@@ -267,6 +267,39 @@ async def test_mrequest(device, make_application, mocker):
     await app.shutdown()
 
 
+@pytest.mark.parametrize("device", [FormedLaunchpadCC26X2R1])
+async def test_mrequest_doesnt_block(device, make_application, event_loop):
+    app, znp_server = make_application(server_cls=device)
+
+    znp_server.reply_once_to(
+        request=c.AF.DataRequestExt.Req(
+            DstAddrModeAddress=t.AddrModeAddress(mode=t.AddrMode.Group, address=0x1234),
+            ClusterId=0x0006,
+            partial=True,
+        ),
+        responses=[
+            # Confirm the request immediately but do not send a callback response until
+            # *after* the group request is "done".
+            c.AF.DataRequestExt.Rsp(Status=t.Status.SUCCESS),
+        ],
+    )
+
+    data_confirm_rsp = c.AF.DataConfirm.Callback(
+        Status=t.Status.SUCCESS, Endpoint=1, TSN=2
+    )
+
+    request_sent = event_loop.create_future()
+    request_sent.add_done_callback(lambda _: znp_server.send(data_confirm_rsp))
+
+    await app.startup(auto_form=False)
+
+    group = app.groups.add_group(0x1234, "test group")
+    await group.endpoint.on_off.on()
+    request_sent.set_result(True)
+
+    await app.shutdown()
+
+
 @pytest.mark.parametrize("device", [FormedLaunchpadCC26X2R1])
 async def test_unimplemented_zdo_converter(device, make_application, mocker):
     app, znp_server = make_application(server_cls=device)

zigpy_znp/api.py

@@ -1,3 +1,4 @@
+import time
 import typing
 import asyncio
 import logging
@@ -398,26 +399,13 @@ async def capture_responses(self, responses):
         """
 
         queue = asyncio.Queue()
-        listener = self.callback_for_responses(responses, queue.put_nowait)
+        listener = self.callback_for_responses(responses, callback=queue.put_nowait)
 
         try:
             yield queue
         finally:
             self.remove_listener(listener)
 
-    @contextlib.asynccontextmanager
-    async def capture_responses_once(self, responses):
-        """
-        Captures all matched responses in a queue within the context manager.
-        """
-
-        future, listener = self.wait_for_responses(responses, context=True)
-
-        try:
-            yield future
-        finally:
-            self.remove_listener(listener)
-
     def callback_for_responses(self, responses, callback) -> CallbackResponseListener:
         """
         Creates a callback listener that matches any of the provided responses.
@@ -546,7 +534,7 @@ async def request(self, request: t.CommandBase, **response_params) -> t.CommandB
         return response
 
     async def request_callback_rsp(
-        self, *, request, callback, timeout=None, **response_params
+        self, *, request, callback, timeout=None, background=False, **response_params
     ):
         """
         Sends an SREQ, gets its SRSP confirmation, and waits for its real AREQ response.
@@ -559,12 +547,41 @@ async def request_callback_rsp(
         from the UART and be handled in the same event loop step by ZNP.
         """
 
+        # Every request should have a timeout to prevent deadlocks
         if timeout is None:
             timeout = self._config[conf.CONF_ZNP_CONFIG][conf.CONF_ARSP_TIMEOUT]
 
-        # The async context manager allows us to clean up resources upon cancellation
-        async with self.capture_responses_once([callback]) as callback_rsp:
-            await self.request(request, **response_params)
+        callback_rsp, listener = self.wait_for_responses([callback], context=True)
+
+        if not background:
+            try:
+                async with async_timeout.timeout(timeout):
+                    await self.request(request, **response_params)
+
+                    return await callback_rsp
+            finally:
+                self.remove_listener(listener)
+
+        start_time = time.time()
 
+        # If the SREQ/SRSP pair fails, we must cancel the AREQ listener
+        try:
             async with async_timeout.timeout(timeout):
-                return await callback_rsp
+                request_rsp = await self.request(request, **response_params)
+        except Exception:
+            self.remove_listener(listener)
+            raise
+
+        async def callback_handler(timeout):
+            try:
+                async with async_timeout.timeout(timeout):
+                    await callback_rsp
+            finally:
+                self.remove_listener(listener)
+
+        # If it succeeds, create a background task to receive the AREQ but take into
+        # account the time it took to start the SREQ to ensure we do not grossly exceed
+        # the timeout
+        asyncio.create_task(callback_handler(time.time() - start_time))
+
+        return request_rsp

zigpy_znp/zigbee/application.py

@@ -31,16 +31,19 @@
 
 ZDO_ENDPOINT = 0
 
-PROBE_TIMEOUT = 5  # seconds
-STARTUP_TIMEOUT = 5  # seconds
-ZDO_REQUEST_TIMEOUT = 15  # seconds
-DATA_CONFIRM_TIMEOUT = 8  # seconds
-DEVICE_JOIN_MAX_DELAY = 5  # seconds
-NETWORK_COMMISSIONING_TIMEOUT = 30  # seconds
-WATCHDOG_PERIOD = 30  # seconds
+# All of these are in seconds
+PROBE_TIMEOUT = 5
+STARTUP_TIMEOUT = 5
+ZDO_REQUEST_TIMEOUT = 15
+DATA_CONFIRM_TIMEOUT = 8
+DEVICE_JOIN_MAX_DELAY = 5
+NETWORK_COMMISSIONING_TIMEOUT = 30
+WATCHDOG_PERIOD = 30
+BROADCAST_SEND_WAIT_DURATION = 3
+MULTICAST_SEND_WAIT_DURATION = 3
 
 REQUEST_MAX_RETRIES = 5
-REQUEST_ERROR_RETRY_DELAY = 0.5  # second
+REQUEST_ERROR_RETRY_DELAY = 0.5  # seconds
 
 # Errors that go away on their own after waiting for a bit
 REQUEST_TRANSIENT_ERRORS = {
@@ -1353,17 +1356,14 @@ async def _send_request_raw(
             )
 
         if dst_addr.mode == t.AddrMode.Broadcast:
-            # Broadcasts will not receive a confirmation but they still take time
-            # and use up concurrency slots
+            # Broadcasts will not receive a confirmation
             response = await self._znp.request(
                 request=request, RspStatus=t.Status.SUCCESS
             )
-
-            await asyncio.sleep(0.1 * 30)
         else:
             async with async_timeout.timeout(DATA_CONFIRM_TIMEOUT):
-                # Shield from cancellation to prevent requests that time out
-                # in higher layers from missing expected responses
+                # Shield from cancellation to prevent requests that time out in higher
+                # layers from missing expected responses
                 response = await asyncio.shield(
                     self._znp.request_callback_rsp(
                         request=request,
@@ -1374,6 +1374,9 @@ async def _send_request_raw(
                             # XXX: can this ever not match?
                             # Endpoint=src_ep,
                         ),
+                        # Multicasts eventually receive a confirmation but waiting for
+                        # it is unnecessary
+                        background=(dst_addr.mode == t.AddrMode.Group),
                     )
                 )