Skip to content

Commit e987c17

Browse files
authored
chore: add some async improvements (#229)
* chore: add some async improvements * chore: improve get_rand_int
1 parent 5331fb2 commit e987c17

File tree

4 files changed

+25
-7
lines changed

4 files changed

+25
-7
lines changed

roborock/api.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
RoborockMessage,
2424
)
2525
from .roborock_typing import RoborockCommand
26-
from .util import RoborockLoggerAdapter, get_running_loop_or_create_one
26+
from .util import RoborockLoggerAdapter, get_next_int, get_running_loop_or_create_one
2727

2828
_LOGGER = logging.getLogger(__name__)
2929
KEEPALIVE = 60
@@ -113,6 +113,13 @@ def _async_response(
113113
self, request_id: int, protocol_id: int = 0
114114
) -> Coroutine[Any, Any, tuple[Any, VacuumError | None]]:
115115
queue = RoborockFuture(protocol_id)
116+
if request_id in self._waiting_queue:
117+
new_id = get_next_int(10000, 32767)
118+
_LOGGER.warning(
119+
f"Attempting to create a future with an existing request_id... New id is {new_id}. "
120+
f"Code may not function properly."
121+
)
122+
request_id = new_id
116123
self._waiting_queue[request_id] = queue
117124
return self._wait_response(request_id, queue)
118125

roborock/roborock_message.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44
import math
55
import time
66
from dataclasses import dataclass
7-
from random import randint
87

98
from roborock import RoborockEnum
9+
from roborock.util import get_next_int
1010

1111

1212
class RoborockMessageProtocol(RoborockEnum):
@@ -155,9 +155,9 @@ class MessageRetry:
155155
class RoborockMessage:
156156
protocol: RoborockMessageProtocol
157157
payload: bytes | None = None
158-
seq: int = randint(100000, 999999)
158+
seq: int = get_next_int(100000, 999999)
159159
version: bytes = b"1.0"
160-
random: int = randint(10000, 99999)
160+
random: int = get_next_int(10000, 99999)
161161
timestamp: int = math.floor(time.time())
162162
message_retry: MessageRetry | None = None
163163

roborock/util.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,3 +108,15 @@ def __init__(self, prefix: str, logger: logging.Logger) -> None:
108108

109109
def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> tuple[str, MutableMapping[str, Any]]:
110110
return f"[{self.prefix}] {msg}", kwargs
111+
112+
113+
counter_map: dict[tuple[int, int], int] = {}
114+
115+
116+
def get_next_int(min_val: int, max_val: int):
117+
"""Gets a random int in the range, precached to help keep it fast."""
118+
if (min_val, max_val) not in counter_map:
119+
# If we have never seen this range, or if the cache is getting low, make a bunch of preshuffled values.
120+
counter_map[(min_val, max_val)] = min_val
121+
counter_map[(min_val, max_val)] += 1
122+
return counter_map[(min_val, max_val)] % max_val + min_val

roborock/version_1_apis/roborock_client_v1.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
import struct
66
import time
77
from collections.abc import Callable, Coroutine
8-
from random import randint
98
from typing import Any, TypeVar, final
109

1110
from roborock import (
@@ -54,7 +53,7 @@
5453
RoborockMessage,
5554
RoborockMessageProtocol,
5655
)
57-
from roborock.util import RepeatableTask, unpack_list
56+
from roborock.util import RepeatableTask, get_next_int, unpack_list
5857

5958
COMMANDS_SECURED = [
6059
RoborockCommand.GET_MAP_V1,
@@ -334,7 +333,7 @@ def _get_payload(
334333
secured=False,
335334
):
336335
timestamp = math.floor(time.time())
337-
request_id = randint(10000, 32767)
336+
request_id = get_next_int(10000, 32767)
338337
inner = {
339338
"id": request_id,
340339
"method": method,

0 commit comments

Comments
 (0)