Leaf/leaf-server/minecraft-patches/features/0165-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 f3e9de8716f5e1a72ec465ee897c8f0413f7b1c3..c83ee2137a57e62003b1d20c3ceea9f569350a53 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 - Async switch connection state
+            if (org.dreeam.leaf.config.modules.network.AlternativeJoin.enabled && ca.spottedleaf.moonrise.common.util.TickThread.isTickThread()) {
+                this.channel.config().setAutoRead(false);
+            }
+            // Leaf end - Async switch connection state
             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 - Async switch connection state
+            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 - Async switch connection state
         }
     }
 
@@ -369,7 +381,38 @@ 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 - Async switch connection state
+            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");
+                }
+            }
+            syncAfterConfigurationChange(cf);
+        }
+    }
+
+    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 - Async switch connection state
         }
     }
 
diff --git a/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java b/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java
index 2e9eb04c7c4342393c05339906c267bca9ff29b1..53b9daa909c2b89046d5af515e17afe09ea7015a 100644
--- a/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java
+++ b/net/minecraft/server/network/ServerConfigurationPacketListenerImpl.java
@@ -140,11 +140,34 @@ public class ServerConfigurationPacketListenerImpl extends ServerCommonPacketLis
         }
     }
 
+    private volatile boolean changingState = false; // Leaf - Async switch connection state
     @Override
     public void handleConfigurationFinished(ServerboundFinishConfigurationPacket packet) {
+        // Leaf start - Async switch connection state
+        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;
+        }
         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);
+            this.connection.setupOutboundProtocol(GameProtocols.CLIENTBOUND_TEMPLATE.bind(RegistryFriendlyByteBuf.decorator(this.server.registryAccess())));
+        }
+        // Leaf end - Async switch connection state
 
         try {
             PlayerList playerList = this.server.getPlayerList();
diff --git a/net/minecraft/server/network/ServerLoginPacketListenerImpl.java b/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
index c25f5b92b16f45b960d9b405b2fe1c40ec4e1124..005fee891cdabc7105033be051389557631fb6d2 100644
--- a/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
+++ b/net/minecraft/server/network/ServerLoginPacketListenerImpl.java
@@ -492,11 +492,31 @@ public class ServerLoginPacketListenerImpl implements ServerLoginPacketListener,
         this.disconnect(ServerCommonPacketListenerImpl.DISCONNECT_UNEXPECTED_QUERY);
     }
 
+    private volatile boolean changingState = false; // Leaf - Async switch connection state
     @Override
     public void handleLoginAcknowledgement(ServerboundLoginAcknowledgedPacket packet) {
+        // Leaf start - Async switch connection state
+        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 - Async switch connection state
         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 - Async switch connection state
         CommonListenerCookie commonListenerCookie = CommonListenerCookie.createInitial(Objects.requireNonNull(this.authenticatedProfile), this.transferred);
         ServerConfigurationPacketListenerImpl serverConfigurationPacketListenerImpl = new ServerConfigurationPacketListenerImpl(
             this.server, this.connection, commonListenerCookie, this.player // CraftBukkit