From 09b27827e70ce0ed6cfb54d4b21bd229af6e1dc7 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Wed, 20 Dec 2023 10:17:51 -0500
Subject: [PATCH] Implement command priority (#602)

---
 bellows/ezsp/__init__.py | 25 +++++++++++++++++++++++--
 pyproject.toml           |  2 +-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py
index b444d5c6..715226ae 100644
--- a/bellows/ezsp/__init__.py
+++ b/bellows/ezsp/__init__.py
@@ -12,6 +12,8 @@
 from typing import Any, Callable, Generator
 import urllib.parse
 
+from zigpy.datastructures import PriorityDynamicBoundedSemaphore
+
 if sys.version_info[:2] < (3, 11):
     from async_timeout import timeout as asyncio_timeout  # pragma: no cover
 else:
@@ -39,6 +41,8 @@
 NETWORK_OPS_TIMEOUT = 10
 NETWORK_COORDINATOR_STARTUP_RESET_WAIT = 1
 
+MAX_COMMAND_CONCURRENCY = 4
+
 
 class EZSP:
     _BY_VERSION = {
@@ -60,6 +64,7 @@ def __init__(self, device_config: dict):
         self._ezsp_version = v4.EZSPv4.VERSION
         self._gw = None
         self._protocol = None
+        self._send_sem = PriorityDynamicBoundedSemaphore(value=MAX_COMMAND_CONCURRENCY)
 
         self._stack_status_listeners: collections.defaultdict[
             t.EmberStatus, list[asyncio.Future]
@@ -184,14 +189,30 @@ def close(self):
             self._gw.close()
             self._gw = None
 
-    def _command(self, name: str, *args: tuple[Any, ...]) -> asyncio.Future:
+    def _get_command_priority(self, name: str) -> int:
+        return {
+            # Deprioritize any commands that send packets
+            "setSourceRoute": -1,
+            "setExtendedTimeout": -1,
+            "sendUnicast": -1,
+            "sendMulticast": -1,
+            "sendBroadcast": -1,
+            # Prioritize watchdog commands
+            "nop": 999,
+            "readCounters": 999,
+            "readAndClearCounters": 999,
+            "getValue": 999,
+        }.get(name, 0)
+
+    async def _command(self, name: str, *args: tuple[Any, ...]) -> Any:
         if not self.is_ezsp_running:
             LOGGER.debug(
                 "Couldn't send command %s(%s). EZSP is not running", name, args
             )
             raise EzspError("EZSP is not running")
 
-        return self._protocol.command(name, *args)
+        async with self._send_sem(priority=self._get_command_priority(name)):
+            return await self._protocol.command(name, *args)
 
     async def _list_command(self, name, item_frames, completion_frame, spos, *args):
         """Run a command, returning result callbacks as a list"""
diff --git a/pyproject.toml b/pyproject.toml
index a2bf9d2d..674197b3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,7 +18,7 @@ dependencies = [
     "click-log>=0.2.1",
     "pure_pcapy3==1.0.1",
     "voluptuous",
-    "zigpy>=0.60.0",
+    "zigpy>=0.60.2",
     'async-timeout; python_version<"3.11"',
 ]