From 93e78593045bb08c4ec1291a7b283606414c1209 Mon Sep 17 00:00:00 2001 From: alemidev Date: Fri, 17 Dec 2021 13:35:28 +0100 Subject: [PATCH] implemented packet whitelist to hopefully improve performance --- aiocraft/client.py | 30 ++++++++++++++---------------- aiocraft/dispatcher.py | 29 +++++++++++++++++++++-------- aiocraft/traits/callbacks.py | 5 ++++- 3 files changed, 39 insertions(+), 25 deletions(-) diff --git a/aiocraft/client.py b/aiocraft/client.py index 35d4d6b..834a87b 100644 --- a/aiocraft/client.py +++ b/aiocraft/client.py @@ -27,10 +27,11 @@ LOGGER = logging.getLogger(__name__) @dataclass class ClientOptions: - reconnect : bool - reconnect_delay : float - keep_alive : bool - poll_interval : float + reconnect : bool = True + reconnect_delay : float = 10.0 + keep_alive : bool = True + poll_interval : float = 1.0 + use_packet_whitelist : bool = True class ClientEvent(Enum): CONNECTED = 0 @@ -61,22 +62,13 @@ class MinecraftClient(CallbacksHolder, Runnable): password:Optional[str] = None, token:Optional[Token] = None, online_mode:bool = True, - reconnect:bool = True, - reconnect_delay:float = 10.0, - keep_alive:bool = True, - poll_interval:float = 1.0, - + **kwargs ): super().__init__() self.host = host self.port = port - self.options = ClientOptions( - reconnect=reconnect, - reconnect_delay=reconnect_delay, - keep_alive=keep_alive, - poll_interval=poll_interval - ) + self.options = ClientOptions(**kwargs) self.token = token self.username = username @@ -186,7 +178,13 @@ class MinecraftClient(CallbacksHolder, Runnable): self._logger.error(str(e)) break try: - await self.dispatcher.connect(self.host, self.port) + packet_whitelist = self.callback_keys(filter=Packet) if self.options.use_packet_whitelist else set() + await self.dispatcher.connect( + self.host, + self.port, + queue_timeout=self.options.poll_interval, + packet_whitelist=packet_whitelist + ) await self._handshake() if await self._login(): await self._play() diff --git a/aiocraft/dispatcher.py b/aiocraft/dispatcher.py index 40bea93..9aebbfe 100644 --- a/aiocraft/dispatcher.py +++ b/aiocraft/dispatcher.py @@ -5,7 +5,7 @@ import zlib import logging from asyncio import StreamReader, StreamWriter, Queue, Task from enum import Enum -from typing import List, Dict, Optional, AsyncIterator, Type +from typing import List, Dict, Set, Optional, AsyncIterator, Type from cryptography.hazmat.primitives.ciphers import CipherContext @@ -23,8 +23,6 @@ class InvalidState(Exception): class ConnectionError(Exception): pass -BROKEN_PACKETS = (77, ) # These packets are still not parseable due to missing data type - class Dispatcher: _is_server : bool # True when receiving packets from clients @@ -41,6 +39,9 @@ class Dispatcher: _incoming : Queue _outgoing : Queue + _packet_whitelist : Set[Packet] + _packet_id_whitelist : Set[int] + _host : str _port : int @@ -89,16 +90,26 @@ class Dispatcher: self.encryption = True self._logger.info("Encryption enabled") - def _prepare(self, host:Optional[str] = None, port:Optional[int] = None, queue_timeout:int = 1, queue_size:int = 100): + def _prepare(self, + host:Optional[str] = None, + port:Optional[int] = None, + queue_timeout:int = 1, + queue_size:int = 100, + packet_whitelist : List[Packet] = None + ): self._host = host or self._host or "localhost" self._port = port or self._port or 25565 self._logger = LOGGER.getChild(f"on({self._host}:{self._port})") + self._packet_whitelist = packet_whitelist or set() self.encryption = False self.compression = None self.state = ConnectionState.HANDSHAKING self.proto = 340 # TODO + # This can only happen after we know the connection protocol + self._packet_id_whitelist = set((P(self.proto).id for P in packet_whitelist)) if packet_whitelist else set() + # Make new queues, do set a max size to sorta propagate back pressure self._incoming = Queue(queue_size) self._outgoing = Queue(queue_size) @@ -112,12 +123,13 @@ class Dispatcher: reader : Optional[StreamReader] = None, writer : Optional[StreamWriter] = None, queue_timeout : int = 1, - queue_size : int = 100 + queue_size : int = 100, + packet_whitelist : Set[Packet] = None, ): if self.connected: raise InvalidState("Dispatcher already connected") - self._prepare(host, port, queue_timeout, queue_size) + self._prepare(host, port, queue_timeout, queue_size, packet_whitelist) if reader and writer: self._down, self._up = reader, writer @@ -216,8 +228,9 @@ class Dispatcher: buffer = io.BytesIO(decompressed_data) packet_id = VarInt.read(buffer) - if packet_id in BROKEN_PACKETS: - continue # cheap fix, still need to implement NBT, Slot and EntityMetadata... + if self._packet_id_whitelist and packet_id in self._packet_id_whitelist: + self._logger.debug("[<--] Received | Packet(0x%02x) (ignored)", packet_id) + continue # ignore this packet, we rarely need them all, should improve performance cls = self._packet_type_from_registry(packet_id) packet = cls.deserialize(self.proto, buffer) self._logger.debug("[<--] Received | %s", repr(packet)) diff --git a/aiocraft/traits/callbacks.py b/aiocraft/traits/callbacks.py index 8c203ee..91745ef 100644 --- a/aiocraft/traits/callbacks.py +++ b/aiocraft/traits/callbacks.py @@ -2,7 +2,7 @@ import asyncio import uuid import logging -from typing import Dict, List, Any, Callable +from typing import Dict, List, Set, Any, Callable, Type class CallbacksHolder: @@ -16,6 +16,9 @@ class CallbacksHolder: self._callbacks = {} self._tasks = {} + def callback_keys(self, filter:Type = None) -> Set[Any]: + return set(x for x in self._callbacks.keys() if not filter or isinstance(x, filter)) + def register(self, key:Any, callback:Callable): if key not in self._callbacks: self._callbacks[key] = []