implemented packet whitelist to hopefully improve performance

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()

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))

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] = []