From 8234360c8d4b54aca3fff531d02fa4bbf7017af2 Mon Sep 17 00:00:00 2001
From: "Dylan K. Taylor" <dktapps@pmmp.io>
Date: Wed, 22 Feb 2023 22:43:10 +0000
Subject: [PATCH] Avoid creating batch buffers just to determine whether a
 batch should be globally compressed Instead, sum together the lengths of
 encoded packet buffers and use that to decide whether to build the buffer or
 not.

---
 src/Server.php                                |  3 ++
 .../mcpe/StandardPacketBroadcaster.php        | 46 +++++++++++++------
 .../mcpe/compression/ThresholdCompressor.php  | 31 +++++++++++++
 .../mcpe/compression/ZlibCompressor.php       | 10 ++--
 4 files changed, 73 insertions(+), 17 deletions(-)
 create mode 100644 src/network/mcpe/compression/ThresholdCompressor.php

diff --git a/src/Server.php b/src/Server.php
index 0fc467bef..116e6f39e 100644
--- a/src/Server.php
+++ b/src/Server.php
@@ -891,6 +891,9 @@ class Server{
 			if($this->configGroup->getPropertyInt("network.batch-threshold", 256) >= 0){
 				$netCompressionThreshold = $this->configGroup->getPropertyInt("network.batch-threshold", 256);
 			}
+			if($netCompressionThreshold < 0){
+				$netCompressionThreshold = null;
+			}
 
 			$netCompressionLevel = $this->configGroup->getPropertyInt("network.compression-level", 6);
 			if($netCompressionLevel < 1 || $netCompressionLevel > 9){
diff --git a/src/network/mcpe/StandardPacketBroadcaster.php b/src/network/mcpe/StandardPacketBroadcaster.php
index c52b8f93d..04b3253f5 100644
--- a/src/network/mcpe/StandardPacketBroadcaster.php
+++ b/src/network/mcpe/StandardPacketBroadcaster.php
@@ -23,21 +23,19 @@ declare(strict_types=1);
 
 namespace pocketmine\network\mcpe;
 
-use pocketmine\network\mcpe\protocol\ClientboundPacket;
+use pocketmine\network\mcpe\compression\ThresholdCompressor;
 use pocketmine\network\mcpe\protocol\serializer\PacketBatch;
 use pocketmine\network\mcpe\protocol\serializer\PacketSerializer;
 use pocketmine\Server;
 use pocketmine\utils\BinaryStream;
-use function array_map;
 use function spl_object_id;
+use function strlen;
 
 final class StandardPacketBroadcaster implements PacketBroadcaster{
 	public function __construct(private Server $server){}
 
 	public function broadcastPackets(array $recipients, array $packets) : void{
-		$batchBuffers = [];
-
-		/** @var string[][] $packetBuffers */
+		$packetBufferTotalLengths = [];
 		$packetBuffers = [];
 		$compressors = [];
 		/** @var NetworkSession[][][] $targetMap */
@@ -45,13 +43,14 @@ final class StandardPacketBroadcaster implements PacketBroadcaster{
 		foreach($recipients as $recipient){
 			$serializerContext = $recipient->getPacketSerializerContext();
 			$bufferId = spl_object_id($serializerContext);
-			if(!isset($batchBuffers[$bufferId])){
-				$packetBuffers[$bufferId] = array_map(function(ClientboundPacket $packet) use ($serializerContext) : string{
-					return NetworkSession::encodePacketTimed(PacketSerializer::encoder($serializerContext), $packet);
-				}, $packets);
-				$stream = new BinaryStream();
-				PacketBatch::encodeRaw($stream, $packetBuffers[$bufferId]);
-				$batchBuffers[$bufferId] = $stream->getBuffer();
+			if(!isset($packetBuffers[$bufferId])){
+				$packetBufferTotalLengths[$bufferId] = 0;
+				$packetBuffers[$bufferId] = [];
+				foreach($packets as $packet){
+					$buffer = NetworkSession::encodePacketTimed(PacketSerializer::encoder($serializerContext), $packet);
+					$packetBufferTotalLengths[$bufferId] += strlen($buffer);
+					$packetBuffers[$bufferId][] = $buffer;
+				}
 			}
 
 			//TODO: different compressors might be compatible, it might not be necessary to split them up by object
@@ -62,10 +61,29 @@ final class StandardPacketBroadcaster implements PacketBroadcaster{
 		}
 
 		foreach($targetMap as $bufferId => $compressorMap){
-			$batchBuffer = $batchBuffers[$bufferId];
 			foreach($compressorMap as $compressorId => $compressorTargets){
 				$compressor = $compressors[$compressorId];
-				if(!$compressor->willCompress($batchBuffer)){
+
+				$batchBuffer = null;
+				if($compressor instanceof ThresholdCompressor){
+					$threshold = $compressor->getCompressionThreshold();
+					if($threshold !== null && $packetBufferTotalLengths[$bufferId] >= $threshold){
+						//do not prepare shared batch unless we're sure it will be compressed
+						$stream = new BinaryStream();
+						PacketBatch::encodeRaw($stream, $packetBuffers[$bufferId]);
+						$batchBuffer = $stream->getBuffer();
+					}
+				}else{
+					//this is a legacy compressor, so we have to encode the batch and check if it will compress
+					$stream = new BinaryStream();
+					PacketBatch::encodeRaw($stream, $packetBuffers[$bufferId]);
+					$tempBatchBuffer = $stream->getBuffer();
+					if($compressor->willCompress($tempBatchBuffer)){
+						$batchBuffer = $tempBatchBuffer;
+					}
+				}
+
+				if($batchBuffer === null){
 					foreach($compressorTargets as $target){
 						foreach($packetBuffers[$bufferId] as $packetBuffer){
 							$target->addToSendBuffer($packetBuffer);
diff --git a/src/network/mcpe/compression/ThresholdCompressor.php b/src/network/mcpe/compression/ThresholdCompressor.php
new file mode 100644
index 000000000..4c3c0ad55
--- /dev/null
+++ b/src/network/mcpe/compression/ThresholdCompressor.php
@@ -0,0 +1,31 @@
+<?php
+
+/*
+ *
+ *  ____            _        _   __  __ _                  __  __ ____
+ * |  _ \ ___   ___| | _____| |_|  \/  (_)_ __   ___      |  \/  |  _ \
+ * | |_) / _ \ / __| |/ / _ \ __| |\/| | | '_ \ / _ \_____| |\/| | |_) |
+ * |  __/ (_) | (__|   <  __/ |_| |  | | | | | |  __/_____| |  | |  __/
+ * |_|   \___/ \___|_|\_\___|\__|_|  |_|_|_| |_|\___|     |_|  |_|_|
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * @author PocketMine Team
+ * @link http://www.pocketmine.net/
+ *
+ *
+ */
+
+declare(strict_types=1);
+
+namespace pocketmine\network\mcpe\compression;
+
+interface ThresholdCompressor extends Compressor{
+	/**
+	 * Returns the minimum size of packet batch that will be compressed. Null means that no packets will be compressed.
+	 */
+	public function getCompressionThreshold() : ?int;
+}
diff --git a/src/network/mcpe/compression/ZlibCompressor.php b/src/network/mcpe/compression/ZlibCompressor.php
index a9ad2718e..1f12c41af 100644
--- a/src/network/mcpe/compression/ZlibCompressor.php
+++ b/src/network/mcpe/compression/ZlibCompressor.php
@@ -32,7 +32,7 @@ use function zlib_decode;
 use function zlib_encode;
 use const ZLIB_ENCODING_RAW;
 
-final class ZlibCompressor implements Compressor{
+final class ZlibCompressor implements ThresholdCompressor{
 	use SingletonTrait;
 
 	public const DEFAULT_LEVEL = 7;
@@ -48,12 +48,16 @@ final class ZlibCompressor implements Compressor{
 
 	public function __construct(
 		private int $level,
-		private int $minCompressionSize,
+		private ?int $minCompressionSize,
 		private int $maxDecompressionSize
 	){}
 
+	public function getCompressionThreshold() : ?int{
+		return $this->minCompressionSize;
+	}
+
 	public function willCompress(string $data) : bool{
-		return $this->minCompressionSize > -1 && strlen($data) >= $this->minCompressionSize;
+		return $this->minCompressionSize !== null && strlen($data) >= $this->minCompressionSize;
 	}
 
 	/**