Allow registering large payloads with deferred max packet size (#5119)

* port from 1.21.11

* change some names and improve impl

* refactor test

* compare against padded size

* tests: update max packet size on both sides separately

* throw if packet type is not registered for getMaxPacketSize()

* change to registerLarge overload

Author: DennisOchulor

fabric-networking-api-v1/src/main/java/net/fabricmc/fabric/api/networking/v1/PayloadTypeRegistry.java

@@ -16,6 +16,8 @@
 
 package net.fabricmc.fabric.api.networking.v1;
 
+import java.util.function.IntSupplier;
+
 import org.jetbrains.annotations.ApiStatus;
 
 import net.minecraft.network.FriendlyByteBuf;
@@ -50,7 +52,7 @@ public interface PayloadTypeRegistry<B extends FriendlyByteBuf> {
 	 * and <strong>before registering a packet handler</strong>.
 	 *
 	 * <p>Payload types registered with this method will be split into multiple packets,
-	 * allowing to send packets larger than vanilla limited size.
+	 * allowing to send packets larger than the vanilla limited size.
 	 *
 	 * @param type          the payload type
 	 * @param codec         the codec for the payload type
@@ -60,6 +62,27 @@ public interface PayloadTypeRegistry<B extends FriendlyByteBuf> {
 	 */
 	<T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super B, T> registerLarge(CustomPacketPayload.Type<T> type, StreamCodec<? super B, T> codec, int maxPacketSize);
 
+	/**
+	 * Registers a large custom payload type.
+	 *
+	 * <p>This must be done on both the sending and receiving side, usually during mod initialization
+	 * and <strong>before registering a packet handler</strong>.
+	 *
+	 * <p>Payload types registered with this method will be split into multiple packets,
+	 * allowing to send packets larger than the vanilla limited size.
+	 *
+	 * <p>The {@code maxPacketSizeSupplier} will be called once, right before the first packet of this payload type
+	 * is sent/received on either side. This allows mods some leeway particularly during mod initialization to
+	 * dynamically determine a suitable max size.
+	 *
+	 * @param type		    the payload type
+	 * @param codec         the codec for the payload type
+	 * @param maxPacketSizeSupplier the function that returns the max size of payload packet
+	 * @param <T>           the payload type
+	 * @return the registered payload type
+	 */
+	<T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super B, T> registerLarge(CustomPacketPayload.Type<T> type, StreamCodec<? super B, T> codec, IntSupplier maxPacketSizeSupplier);
+
 	/**
 	 * @return the {@link PayloadTypeRegistry} instance for the serverbound (client to server) configuration channel.
 	 */

fabric-networking-api-v1/src/main/java/net/fabricmc/fabric/impl/networking/PayloadTypeRegistryImpl.java

@@ -19,10 +19,13 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.IntSupplier;
 
 import io.netty.buffer.ByteBufUtil;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
 import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
 import org.jspecify.annotations.Nullable;
 
 import net.minecraft.network.ConnectionProtocol;
@@ -44,7 +47,8 @@ public class PayloadTypeRegistryImpl<B extends FriendlyByteBuf> implements Paylo
 	public static final PayloadTypeRegistryImpl<RegistryFriendlyByteBuf> SERVERBOUND_PLAY = new PayloadTypeRegistryImpl<>(ConnectionProtocol.PLAY, PacketFlow.SERVERBOUND);
 	public static final PayloadTypeRegistryImpl<RegistryFriendlyByteBuf> CLIENTBOUND_PLAY = new PayloadTypeRegistryImpl<>(ConnectionProtocol.PLAY, PacketFlow.CLIENTBOUND);
 	private final Map<Identifier, CustomPacketPayload.TypeAndCodec<B, ? extends CustomPacketPayload>> packetTypes = new HashMap<>();
-	private final Object2IntMap<Identifier> maxPacketSize = new Object2IntOpenHashMap<>();
+	private final Object2IntMap<Identifier> maxPacketSizes = new Object2IntOpenHashMap<>();
+	private final Object2ObjectMap<Identifier, IntSupplier> pendingMaxPacketSizes = new Object2ObjectOpenHashMap<>();
 	private final ConnectionProtocol protocol;
 	private final PacketFlow flow;
 	private final int minimalSplittableSize;
@@ -66,7 +70,7 @@ public static PayloadTypeRegistryImpl<?> get(ProtocolInfo<?> state) {
 
 	@Override
 	public <T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super B, T> register(CustomPacketPayload.Type<T> type, StreamCodec<? super B, T> codec) {
-		Objects.requireNonNull(type, "id");
+		Objects.requireNonNull(type, "type");
 		Objects.requireNonNull(codec, "codec");
 
 		final CustomPacketPayload.TypeAndCodec<B, T> payloadType = new CustomPacketPayload.TypeAndCodec<>(type, codec.cast());
@@ -80,15 +84,30 @@ public <T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super
 	}
 
 	@Override
-	public <T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super B, T> registerLarge(CustomPacketPayload.Type<T> type, StreamCodec<? super B, T> codec, int maxPayloadSize) {
-		if (maxPayloadSize < 0) {
-			throw new IllegalArgumentException("Provided maxPayloadSize needs to be positive!");
+	public <T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super B, T> registerLarge(CustomPacketPayload.Type<T> type, StreamCodec<? super B, T> codec, int maxPacketSize) {
+		if (maxPacketSize < 0) {
+			throw new IllegalArgumentException("Provided maxPacketSize needs to be positive!");
 		}
 
 		CustomPacketPayload.TypeAndCodec<? super B, T> typeAndCodec = register(type, codec);
+		padAndSetMaxPacketSize(type.id(), maxPacketSize);
+		return typeAndCodec;
+	}
+
+	@Override
+	public <T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super B, T> registerLarge(CustomPacketPayload.Type<T> type, StreamCodec<? super B, T> codec, IntSupplier maxPacketSizeSupplier) {
+		Objects.requireNonNull(maxPacketSizeSupplier, "maxPacketSizeSupplier");
+
+		CustomPacketPayload.TypeAndCodec<? super B, T> typeAndCodec = register(type, codec);
+		pendingMaxPacketSizes.put(type.id(), maxPacketSizeSupplier);
+		return typeAndCodec;
+	}
+
+	private void padAndSetMaxPacketSize(Identifier id, int maxSize) {
 		// Defines max packet size, increased by length of packet's Identifier to cover full size of CustomPayloadX2YPackets.
-		int identifierSize = ByteBufUtil.utf8MaxBytes(type.id().toString());
-		int maxPacketSize = maxPayloadSize + VarInt.getByteSize(identifierSize) + identifierSize + 5 * 2;
+		int identifierSize = ByteBufUtil.utf8MaxBytes(id.toString());
+		int paddingSize = VarInt.getByteSize(identifierSize) + identifierSize + 5 * 2;
+		int maxPacketSize = maxSize + paddingSize;
 
 		// Prevent overflow
 		if (maxPacketSize < 0) {
@@ -97,23 +116,36 @@ public <T extends CustomPacketPayload> CustomPacketPayload.TypeAndCodec<? super
 
 		// No need to enable splitting, if packet's max size is smaller than chunk
 		if (maxPacketSize > this.minimalSplittableSize) {
-			this.maxPacketSize.put(type.id(), maxPacketSize);
+			this.maxPacketSizes.put(id, maxPacketSize);
 		}
-
-		return typeAndCodec;
 	}
 
 	public CustomPacketPayload.@Nullable TypeAndCodec<B, ? extends CustomPacketPayload> get(Identifier id) {
 		return packetTypes.get(id);
 	}
 
-	public <T extends CustomPacketPayload> CustomPacketPayload.@Nullable TypeAndCodec<B, T> get(CustomPacketPayload.Type<T> id) {
+	public <T extends CustomPacketPayload> CustomPacketPayload.@Nullable TypeAndCodec<B, T> get(CustomPacketPayload.Type<T> type) {
 		//noinspection unchecked
-		return (CustomPacketPayload.TypeAndCodec<B, T>) packetTypes.get(id.id());
+		return (CustomPacketPayload.TypeAndCodec<B, T>) packetTypes.get(type.id());
 	}
 
-	public int getMaxPacketSize(Identifier id) {
-		return this.maxPacketSize.getOrDefault(id, -1);
+	/**
+	 * @return the max packet size, or -1 if the payload type does not need splitting.
+	 */
+	public int getMaxPacketSizeForSplitting(Identifier id) {
+		IntSupplier supplier = this.pendingMaxPacketSizes.remove(id);
+
+		if (supplier != null) {
+			int maxPacketSize = supplier.getAsInt();
+
+			if (maxPacketSize < 0) {
+				throw new IllegalArgumentException("maxPacketSize supplier for packet type " + id + ": must be positive!");
+			}
+
+			padAndSetMaxPacketSize(id, maxPacketSize);
+		}
+
+		return this.maxPacketSizes.getOrDefault(id, -1);
 	}
 
 	public ConnectionProtocol getProtocol() {

fabric-networking-api-v1/src/main/java/net/fabricmc/fabric/impl/networking/splitter/FabricPacketMerger.java

@@ -82,7 +82,7 @@ protected void decode(ChannelHandlerContext channelHandlerContext, Packet<?> pac
 			Identifier payloadId = Identifier.STREAM_CODEC.decode(payload.byteBuf());
 
 			buf.readerIndex(readerIndex);
-			int maxSize = payloadTypeRegistry.getMaxPacketSize(payloadId);
+			int maxSize = payloadTypeRegistry.getMaxPacketSizeForSplitting(payloadId);
 
 			if (maxSize == -1) {
 				throw new DecoderException("Received '" + payloadId + "' packet doesn't support splitting, but received split data!");

fabric-networking-api-v1/src/main/java/net/fabricmc/fabric/mixin/networking/ClientboundCustomPayloadPacketMixin.java

@@ -79,7 +79,7 @@ private static StreamCodec<FriendlyByteBuf, CustomPacketPayload> wrapConfigCodec
 
 	@Override
 	public void fabric_split(PayloadTypeRegistryImpl<?> payloadTypeRegistry, ChannelHandlerContext channelHandlerContext, PacketEncoder<?> encoder, Packet<?> packet, Consumer<Packet<?>> consumer) throws Exception {
-		int size = payloadTypeRegistry.getMaxPacketSize(this.payload.type().id());
+		int size = payloadTypeRegistry.getMaxPacketSizeForSplitting(this.payload.type().id());
 
 		if (size == -1) {
 			consumer.accept((Packet<?>) this);

fabric-networking-api-v1/src/main/java/net/fabricmc/fabric/mixin/networking/ServerboundCustomPayloadPacketMixin.java

@@ -70,7 +70,7 @@ private static StreamCodec<FriendlyByteBuf, CustomPacketPayload> wrapCodec(Custo
 
 	@Override
 	public void fabric_split(PayloadTypeRegistryImpl<?> payloadTypeRegistry, ChannelHandlerContext channelHandlerContext, PacketEncoder<?> encoder, Packet<?> packet, Consumer<Packet<?>> consumer) throws Exception {
-		int size = payloadTypeRegistry.getMaxPacketSize(this.payload.type().id());
+		int size = payloadTypeRegistry.getMaxPacketSizeForSplitting(this.payload.type().id());
 
 		if (size == -1) {
 			consumer.accept((Packet<?>) this);

fabric-networking-api-v1/src/testmod/java/net/fabricmc/fabric/test/networking/splitter/NetworkingSplitterTest.java

@@ -34,11 +34,11 @@
 import net.fabricmc.fabric.test.networking.NetworkingTestmods;
 
 public class NetworkingSplitterTest implements ModInitializer {
-	private static final Logger LOGGER = LoggerFactory.getLogger(NetworkingSplitterTest.class);
+	public static final Logger LOGGER = LoggerFactory.getLogger(NetworkingSplitterTest.class);
 
-	private static final int DATA_SIZE = 20 * 1024 * 1024;
+	public static final int DATA_SIZE = 20 * 1024 * 1024;
 
-	// 20 MB of random data source
+	// 20 MiB of random data source
 	private static final int[][] RANDOM_DATA = {
 			IntStream.generate(RandomSource.create(24534)::nextInt).limit(20).toArray(),
 			IntStream.generate(RandomSource.create(24533)::nextInt).limit(DATA_SIZE / 4).toArray()
@@ -48,7 +48,7 @@ public class NetworkingSplitterTest implements ModInitializer {
 	public void onInitialize() {
 		// Register the payload on both sides for play and configuration
 		PayloadTypeRegistry.clientboundPlay().registerLarge(LargePayload.TYPE, LargePayload.CODEC, DATA_SIZE + 14);
-		PayloadTypeRegistry.serverboundPlay().registerLarge(LargePayload.TYPE, LargePayload.CODEC, DATA_SIZE + 14);
+		PayloadTypeRegistry.serverboundPlay().registerLarge(LargePayload.TYPE, LargePayload.CODEC, () -> DATA_SIZE + 14);
 
 		// When the client joins, send a packet expecting it to be validated and echoed back
 		ServerPlayConnectionEvents.JOIN.register((listener, sender, server) -> sender.sendPacket(new LargePayload(0, RANDOM_DATA[0])));
@@ -62,7 +62,7 @@ public void onInitialize() {
 
 	public static void validateLargePacketData(int index, int[] data, String side) {
 		if (Arrays.equals(RANDOM_DATA[index], data)) {
-			NetworkingTestmods.LOGGER.info("Successfully received large packet [" + index + "] on " + side);
+			LOGGER.info("Successfully received large packet [{}] on {}", index, side);
 			return;
 		}