Leaf/leaf-server/minecraft-patches/features/0167-async-switch-connection-state.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: hayanesuru <hayanesuru@outlook.jp>
Date: Fri, 9 May 2025 16:55:34 +0900
Subject: [PATCH] async switch connection state


diff --git a/net/minecraft/network/Connection.java b/net/minecraft/network/Connection.java
index f998cf8d70302a21289de4d84b46d322d0b8a8fe..32f26640664135c9f7f45f8b204b7ff412fe343e 100644
--- a/net/minecraft/network/Connection.java
+++ b/net/minecraft/network/Connection.java
@@ -342,6 +342,11 @@ public class Connection extends SimpleChannelInboundHandler<Packet<?>> {
         if (protocolInfo.flow() != this.getReceiving()) {
             throw new IllegalStateException("Invalid inbound protocol: " + protocolInfo.id());
         } else {
+            // Leaf start
+            if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled && ca.spottedleaf.moonrise.common.util.TickThread.isTickThread()) {
+                this.channel.config().setAutoRead(false);
+            }
+            // Leaf end
             this.packetListener = packetInfo;
             this.disconnectListener = null;
             UnconfiguredPipelineHandler.InboundConfigurationTask inboundConfigurationTask = UnconfiguredPipelineHandler.setupInboundProtocol(protocolInfo);
@@ -351,7 +356,14 @@ public class Connection extends SimpleChannelInboundHandler<Packet<?>> {
                 inboundConfigurationTask = inboundConfigurationTask.andThen(context -> context.pipeline().addAfter("decoder", "bundler", packetBundlePacker));
             }
 
-            syncAfterConfigurationChange(this.channel.writeAndFlush(inboundConfigurationTask));
+            // Leaf start
+            var cf = this.channel.writeAndFlush(inboundConfigurationTask);
+            if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled && ca.spottedleaf.moonrise.common.util.TickThread.isTickThread()) {
+                cf.addListener((ChannelFutureListener) Connection::syncAfterConfigurationChange);
+                return;
+            }
+            syncAfterConfigurationChange(cf);
+            // Leaf end
         }
     }
 
@@ -369,9 +381,41 @@ public class Connection extends SimpleChannelInboundHandler<Packet<?>> {
             }
 
             boolean flag = protocolInfo.id() == ConnectionProtocol.LOGIN;
-            syncAfterConfigurationChange(this.channel.writeAndFlush(outboundConfigurationTask.andThen(context -> this.sendLoginDisconnect = flag)));
+            var cf = this.channel.writeAndFlush(outboundConfigurationTask.andThen(context -> this.sendLoginDisconnect = flag));
+            // Leaf start
+            if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled) {
+                if (ca.spottedleaf.moonrise.common.util.TickThread.isTickThread()) {
+                    throw new IllegalStateException("Thread failed netty thread check: Switching outbound protocol state use setupOutboundProtocolAsync instead");
+                }
+            }
+            // Leaf end
+            syncAfterConfigurationChange(cf);
+        }
+    }
+    // Leaf start
+    public @Nullable ChannelFuture setupOutboundProtocolAsync(ProtocolInfo<?> protocolInfo) {
+        if (protocolInfo.flow() != this.getSending()) {
+            throw new IllegalStateException("Invalid outbound protocol: " + protocolInfo.id());
+        } else {
+            UnconfiguredPipelineHandler.OutboundConfigurationTask outboundConfigurationTask = UnconfiguredPipelineHandler.setupOutboundProtocol(protocolInfo);
+            BundlerInfo bundlerInfo = protocolInfo.bundlerInfo();
+            if (bundlerInfo != null) {
+                PacketBundleUnpacker packetBundleUnpacker = new PacketBundleUnpacker(bundlerInfo);
+                outboundConfigurationTask = outboundConfigurationTask.andThen(
+                    context -> context.pipeline().addAfter("encoder", "unbundler", packetBundleUnpacker)
+                );
+            }
+
+            boolean flag = protocolInfo.id() == ConnectionProtocol.LOGIN;
+            var cf = this.channel.writeAndFlush(outboundConfigurationTask.andThen(context -> this.sendLoginDisconnect = flag));
+            if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled) {
+                cf.addListener((ChannelFutureListener) Connection::syncAfterConfigurationChange);
+                return cf;
+            }
+            return null;
         }
     }
+    // Leaf end
 
     public void setListenerForServerboundHandshake(PacketListener packetListener) {
         if (this.packetListener != null) {
diff --git a/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java b/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java
index 2e9eb04c7c4342393c05339906c267bca9ff29b1..c70d5a0db1dfd01eab323aefd07d6e81dd188927 100644
--- a/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java
+++ b/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java
@@ -140,11 +140,32 @@ public class ServerConfigurationPacketListenerImpl extends ServerCommonPacketLis
         }
     }
 
+    private volatile boolean changingState = false; // Leaf
     @Override
     public void handleConfigurationFinished(ServerboundFinishConfigurationPacket packet) {
+        // Leaf start
+        if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled && !changingState) {
+            changingState = true;
+            this.finishCurrentTask(JoinWorldTask.TYPE);
+            this.connection.setupOutboundProtocolAsync(GameProtocols.CLIENTBOUND_TEMPLATE.bind(RegistryFriendlyByteBuf.decorator(this.server.registryAccess()))).addListener(l -> {
+                try {
+                    PacketUtils.ensureRunningOnSameThread(packet, this, this.server);
+                } catch (net.minecraft.server.RunningOnDifferentThreadException ignored) {
+                } catch (
+                    io.papermc.paper.util.ServerStopRejectedExecutionException ignored) { // Paper - do not prematurely disconnect players on stop
+                } catch (java.util.concurrent.RejectedExecutionException var6) {
+                    this.connection.disconnect(Component.translatable("multiplayer.disconnect.server_shutdown"));
+                } catch (ClassCastException var7) {
+                    LOGGER.error("Received {} that couldn't be processed", packet.getClass(), var7);
+                    this.connection.disconnect(Component.translatable("multiplayer.disconnect.invalid_packet"));
+                }
+            });
+            return;
+        }
+        // Leaf end
         PacketUtils.ensureRunningOnSameThread(packet, this, this.server);
-        this.finishCurrentTask(JoinWorldTask.TYPE);
-        this.connection.setupOutboundProtocol(GameProtocols.CLIENTBOUND_TEMPLATE.bind(RegistryFriendlyByteBuf.decorator(this.server.registryAccess())));
+        if (!org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled) { this.finishCurrentTask(JoinWorldTask.TYPE); } // Leaf
+        if (!org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled) { this.connection.setupOutboundProtocol(GameProtocols.CLIENTBOUND_TEMPLATE.bind(RegistryFriendlyByteBuf.decorator(this.server.registryAccess()))); } // Leaf
 
         try {
             PlayerList playerList = this.server.getPlayerList();
diff --git a/net/minecraft/server/network/ServerLoginPacketListenerImpl.java b/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
index 49f1743db193be1f10bfe6419231eb682e1068f7..b0ffd2077747b2325ab795eef457b9a0fa44754b 100644
--- a/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
+++ b/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
@@ -472,11 +472,32 @@ public class ServerLoginPacketListenerImpl implements ServerLoginPacketListener,
         this.disconnect(ServerCommonPacketListenerImpl.DISCONNECT_UNEXPECTED_QUERY);
     }
 
+    private volatile boolean changingState = false; // Leaf
     @Override
     public void handleLoginAcknowledgement(ServerboundLoginAcknowledgedPacket packet) {
+        // Leaf start
+        if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled && !changingState) {
+            changingState = true;
+            this.connection.setupOutboundProtocolAsync(ConfigurationProtocols.CLIENTBOUND).addListener(l -> {
+                try {
+                    PacketUtils.ensureRunningOnSameThread(packet, this, this.server);
+                } catch (net.minecraft.server.RunningOnDifferentThreadException ignored) {
+                } catch (
+                    io.papermc.paper.util.ServerStopRejectedExecutionException ignored) { // Paper - do not prematurely disconnect players on stop
+                } catch (java.util.concurrent.RejectedExecutionException var6) {
+                    this.connection.disconnect(Component.translatable("multiplayer.disconnect.server_shutdown"));
+                } catch (ClassCastException var7) {
+                    LOGGER.error("Received {} that couldn't be processed", packet.getClass(), var7);
+                    this.connection.disconnect(Component.translatable("multiplayer.disconnect.invalid_packet"));
+                }
+            });
+            return;
+        }
+        // Leaf end
+
         PacketUtils.ensureRunningOnSameThread(packet, this, this.server); // CraftBukkit
         Validate.validState(this.state == ServerLoginPacketListenerImpl.State.PROTOCOL_SWITCHING, "Unexpected login acknowledgement packet");
-        this.connection.setupOutboundProtocol(ConfigurationProtocols.CLIENTBOUND);
+        if (!org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled) { this.connection.setupOutboundProtocol(ConfigurationProtocols.CLIENTBOUND); } // Leaf
         CommonListenerCookie commonListenerCookie = CommonListenerCookie.createInitial(Objects.requireNonNull(this.authenticatedProfile), this.transferred);
         ServerConfigurationPacketListenerImpl serverConfigurationPacketListenerImpl = new ServerConfigurationPacketListenerImpl(
             this.server, this.connection, commonListenerCookie, this.player // CraftBukkit