chore: add some async improvements (#229)

* chore: add some async improvements

* chore: improve get_rand_int

roborock/api.py

@@ -23,7 +23,7 @@
     RoborockMessage,
 )
 from .roborock_typing import RoborockCommand
-from .util import RoborockLoggerAdapter, get_running_loop_or_create_one
+from .util import RoborockLoggerAdapter, get_next_int, get_running_loop_or_create_one
 
 _LOGGER = logging.getLogger(__name__)
 KEEPALIVE = 60
@@ -113,6 +113,13 @@ def _async_response(
         self, request_id: int, protocol_id: int = 0
     ) -> Coroutine[Any, Any, tuple[Any, VacuumError | None]]:
         queue = RoborockFuture(protocol_id)
+        if request_id in self._waiting_queue:
+            new_id = get_next_int(10000, 32767)
+            _LOGGER.warning(
+                f"Attempting to create a future with an existing request_id... New id is {new_id}. "
+                f"Code may not function properly."
+            )
+            request_id = new_id
         self._waiting_queue[request_id] = queue
         return self._wait_response(request_id, queue)
 

roborock/roborock_message.py

@@ -4,9 +4,9 @@
 import math
 import time
 from dataclasses import dataclass
-from random import randint
 
 from roborock import RoborockEnum
+from roborock.util import get_next_int
 
 
 class RoborockMessageProtocol(RoborockEnum):
@@ -155,9 +155,9 @@ class MessageRetry:
 class RoborockMessage:
     protocol: RoborockMessageProtocol
     payload: bytes | None = None
-    seq: int = randint(100000, 999999)
+    seq: int = get_next_int(100000, 999999)
     version: bytes = b"1.0"
-    random: int = randint(10000, 99999)
+    random: int = get_next_int(10000, 99999)
     timestamp: int = math.floor(time.time())
     message_retry: MessageRetry | None = None
 

roborock/util.py

@@ -108,3 +108,15 @@ def __init__(self, prefix: str, logger: logging.Logger) -> None:
 
     def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> tuple[str, MutableMapping[str, Any]]:
         return f"[{self.prefix}] {msg}", kwargs
+
+
+counter_map: dict[tuple[int, int], int] = {}
+
+
+def get_next_int(min_val: int, max_val: int):
+    """Gets a random int in the range, precached to help keep it fast."""
+    if (min_val, max_val) not in counter_map:
+        # If we have never seen this range, or if the cache is getting low, make a bunch of preshuffled values.
+        counter_map[(min_val, max_val)] = min_val
+    counter_map[(min_val, max_val)] += 1
+    return counter_map[(min_val, max_val)] % max_val + min_val

roborock/version_1_apis/roborock_client_v1.py

@@ -5,7 +5,6 @@
 import struct
 import time
 from collections.abc import Callable, Coroutine
-from random import randint
 from typing import Any, TypeVar, final
 
 from roborock import (
@@ -54,7 +53,7 @@
     RoborockMessage,
     RoborockMessageProtocol,
 )
-from roborock.util import RepeatableTask, unpack_list
+from roborock.util import RepeatableTask, get_next_int, unpack_list
 
 COMMANDS_SECURED = [
     RoborockCommand.GET_MAP_V1,
@@ -334,7 +333,7 @@ def _get_payload(
         secured=False,
     ):
         timestamp = math.floor(time.time())
-        request_id = randint(10000, 32767)
+        request_id = get_next_int(10000, 32767)
         inner = {
             "id": request_id,
             "method": method,