Skip to content

Commit

Permalink
chore: add some async improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
Lash-L committed Oct 28, 2024
1 parent 2731bce commit afcaa7f
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 7 deletions.
9 changes: 8 additions & 1 deletion roborock/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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_rand_int, get_running_loop_or_create_one

_LOGGER = logging.getLogger(__name__)
KEEPALIVE = 60
Expand Down Expand Up @@ -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_rand_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)

Expand Down
6 changes: 3 additions & 3 deletions roborock/roborock_message.py
Original file line number Diff line number Diff line change
Expand Up @@ -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_rand_int


class RoborockMessageProtocol(RoborockEnum):
Expand Down Expand Up @@ -155,9 +155,9 @@ class MessageRetry:
class RoborockMessage:
protocol: RoborockMessageProtocol
payload: bytes | None = None
seq: int = randint(100000, 999999)
seq: int = get_rand_int(100000, 999999)
version: bytes = b"1.0"
random: int = randint(10000, 99999)
random: int = get_rand_int(10000, 99999)
timestamp: int = math.floor(time.time())
message_retry: MessageRetry | None = None

Expand Down
13 changes: 13 additions & 0 deletions roborock/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import datetime
import functools
import logging
import random
from asyncio import AbstractEventLoop, TimerHandle
from collections.abc import Callable, Coroutine, MutableMapping
from typing import Any, TypeVar
Expand Down Expand Up @@ -108,3 +109,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


cached_numbers: dict[tuple[int, int], list[int]] = {}


def get_rand_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 cached_numbers or len(cached_numbers[(min_val, max_val)]) < 25:
# If we have never seen this range, or if the cache is getting low, make a bunch of preshuffled values.
cached_numbers[(min_val, max_val)] = list(range(min_val, max_val + 1))
random.shuffle(cached_numbers[(min_val, max_val)])
return cached_numbers[(min_val, max_val)].pop()
5 changes: 2 additions & 3 deletions roborock/version_1_apis/roborock_client_v1.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -54,7 +53,7 @@
RoborockMessage,
RoborockMessageProtocol,
)
from roborock.util import RepeatableTask, unpack_list
from roborock.util import RepeatableTask, get_rand_int, unpack_list

COMMANDS_SECURED = [
RoborockCommand.GET_MAP_V1,
Expand Down Expand Up @@ -334,7 +333,7 @@ def _get_payload(
secured=False,
):
timestamp = math.floor(time.time())
request_id = randint(10000, 32767)
request_id = get_rand_int(10000, 32767)
inner = {
"id": request_id,
"method": method,
Expand Down

0 comments on commit afcaa7f

Please sign in to comment.