Allow registering large payloads with deferred max packet size

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

@@ -50,7 +50,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 +60,26 @@ 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);
 
+	/**
+	 * Sets the maximum size of an <strong>already registered</strong> payload type.
+	 *
+	 * <p>Larger payloads will be split into multiple packets, allowing to send packets larger than the vanilla limited size.
+	 *
+	 * @param type		    the payload type
+	 * @param maxPacketSize the maximum size of payload packet
+	 * @param <T>           the payload type
+	 */
+	<T extends CustomPacketPayload> void setMaxPacketSize(CustomPacketPayload.Type<T> type, int maxPacketSize);
+
+	/**
+	 * Returns the maximum size of an <strong>already registered</strong> payload type.
+	 *
+	 * @param type		    the payload type
+	 * @param <T>           the payload class
+	 * @return the maximum size of payload packet
+	 */
+	<T extends CustomPacketPayload> int getMaxPacketSize(CustomPacketPayload.Type<T> type);
+
 	/**
 	 * @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

@@ -44,7 +44,7 @@ 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 ConnectionProtocol protocol;
 	private final PacketFlow flow;
 	private final int minimalSplittableSize;
@@ -66,7 +66,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,40 +80,68 @@ 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> void setMaxPacketSize(CustomPacketPayload.Type<T> type, int maxPacketSize) {
+		Identifier id = type.id();
+
+		if (!this.packetTypes.containsKey(id)) {
+			throw new IllegalArgumentException("Packet type " + id + " has not been registered yet!");
+		}
+
+		if (maxPacketSize < 0) {
+			throw new IllegalArgumentException("Provided maxPacketSize needs to be positive!");
+		}
+
+		padAndSetMaxPacketSize(id, maxPacketSize);
+	}
+
+	@Override
+	public <T extends CustomPacketPayload> int getMaxPacketSize(CustomPacketPayload.Type<T> type) {
+		return this.maxPacketSizes.getOrDefault(type.id(), this.minimalSplittableSize);
+	}
+
+	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) {
 			maxPacketSize = Integer.MAX_VALUE;
 		}
 
 		// No need to enable splitting, if packet's max size is smaller than chunk
-		if (maxPacketSize > this.minimalSplittableSize) {
-			this.maxPacketSize.put(type.id(), maxPacketSize);
+		// If requested maxPacketSize <= previous max without padding, then ignore it.
+		if (maxPacketSize > Math.max(this.minimalSplittableSize, this.maxPacketSizes.getOrDefault(id, -1))) {
+			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) {
+		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,21 +34,23 @@
 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_1 = 20 * 1024 * 1024;
+	public static final int DATA_SIZE_2 = 50 * 1024 * 1024;
 
-	// 20 MB of random data source
+	// 20 and 50 MB 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()
+			IntStream.generate(RandomSource.create(24533)::nextInt).limit(DATA_SIZE_1 / 4).toArray(),
+			IntStream.generate(RandomSource.create(24532)::nextInt).limit(DATA_SIZE_2 / 4).toArray()
 	};
 
 	@Override
 	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.clientboundPlay().registerLarge(LargePayload.TYPE, LargePayload.CODEC, DATA_SIZE_1 + 14);
+		PayloadTypeRegistry.serverboundPlay().registerLarge(LargePayload.TYPE, LargePayload.CODEC, DATA_SIZE_1 + 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])));
@@ -57,12 +59,20 @@ public void onInitialize() {
 		// Validate received packet
 		ServerPlayNetworking.registerGlobalReceiver(LargePayload.TYPE, (payload, context) -> {
 			validateLargePacketData(payload.index(), payload.data(), "server");
+
+			// After validating 20 MB packet, try 50 MB.
+			if (payload.index() == 1) {
+				LOGGER.info("Increasing max size of LargePayload to 50MB on server");
+				PayloadTypeRegistry.clientboundPlay().setMaxPacketSize(LargePayload.TYPE, DATA_SIZE_2 + 14);
+				PayloadTypeRegistry.serverboundPlay().setMaxPacketSize(LargePayload.TYPE, DATA_SIZE_2 + 14);
+				context.responseSender().sendPacket(new LargePayload(2, RANDOM_DATA[2]));
+			}
 		});
 	}
 
 	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;
 		}
 

fabric-networking-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/networking/client/splitter/NetworkingSplitterClientTest.java

@@ -18,6 +18,7 @@
 
 import net.fabricmc.api.ClientModInitializer;
 import net.fabricmc.fabric.api.client.networking.v1.ClientPlayNetworking;
+import net.fabricmc.fabric.api.networking.v1.PayloadTypeRegistry;
 import net.fabricmc.fabric.test.networking.splitter.NetworkingSplitterTest;
 
 public class NetworkingSplitterClientTest implements ClientModInitializer {
@@ -26,6 +27,13 @@ public void onInitializeClient() {
 		ClientPlayNetworking.registerGlobalReceiver(NetworkingSplitterTest.LargePayload.TYPE, (payload, context) -> {
 			NetworkingSplitterTest.validateLargePacketData(payload.index(), payload.data(), "client");
 			context.responseSender().sendPacket(payload);
+
+			// After validating 20 MB packet, try 50 MB.
+			if (payload.index() == 1) {
+				NetworkingSplitterTest.LOGGER.info("Increasing max size of LargePayload to 50MB on client");
+				PayloadTypeRegistry.clientboundPlay().setMaxPacketSize(NetworkingSplitterTest.LargePayload.TYPE, NetworkingSplitterTest.DATA_SIZE_2 + 14);
+				PayloadTypeRegistry.serverboundPlay().setMaxPacketSize(NetworkingSplitterTest.LargePayload.TYPE, NetworkingSplitterTest.DATA_SIZE_2 + 14);
+			}
 		});
 	}
 }