Fix some issue caused by refactor

Freesia-Velocity/src/main/java/meow/kikir/freesia/velocity/network/ysm/MapperSessionProcessor.java

@@ -1,6 +1,7 @@
 package meow.kikir.freesia.velocity.network.ysm;
 
 import ca.spottedleaf.concurrentutil.collection.MultiThreadedQueue;
+import ca.spottedleaf.concurrentutil.util.ConcurrentUtil;
 import com.velocitypowered.api.proxy.Player;
 import com.velocitypowered.api.proxy.messages.MinecraftChannelIdentifier;
 import io.netty.buffer.ByteBuf;
@@ -17,17 +18,27 @@
 import org.geysermc.mcprotocollib.protocol.packet.common.serverbound.ServerboundPongPacket;
 import org.geysermc.mcprotocollib.protocol.packet.ingame.clientbound.ClientboundLoginPacket;
 
+import java.lang.invoke.VarHandle;
 import java.util.Optional;
 import java.util.UUID;
 
 public class MapperSessionProcessor implements SessionListener {
     private final Player bindPlayer;
     private final YsmPacketProxy packetProxy;
     private final YsmMapperPayloadManager mapperPayloadManager;
+
+    // Callbacks for packet processing and tracker updates
     private final MultiThreadedQueue<PendingPacket> pendingYsmPacketsInbound = new MultiThreadedQueue<>();
     private final MultiThreadedQueue<UUID> pendingTrackerUpdatesTo = new MultiThreadedQueue<>();
+
+    // Controlled by the VarHandles following
     private volatile Session session;
-    private volatile boolean kickMasterWhenDisconnect = true;
+    private boolean kickMasterWhenDisconnect = true;
+    private boolean destroyed = false;
+
+    private static final VarHandle KICK_MASTER_HANDLE = ConcurrentUtil.getVarHandle(MapperSessionProcessor.class, "kickMasterWhenDisconnect", boolean.class);
+    private static final VarHandle SESSION_HANDLE = ConcurrentUtil.getVarHandle(MapperSessionProcessor.class, "session", Session.class);
+    private static final VarHandle DESTROYED_HANDLE = ConcurrentUtil.getVarHandle(MapperSessionProcessor.class, "destroyed", boolean.class);
 
     public MapperSessionProcessor(Player bindPlayer, YsmPacketProxy packetProxy, YsmMapperPayloadManager mapperPayloadManager) {
         this.bindPlayer = bindPlayer;
@@ -58,16 +69,21 @@ protected YsmPacketProxy getPacketProxy() {
         return this.packetProxy;
     }
 
-    protected Session getSession() {
-        return this.session;
-    }
-
     protected void setKickMasterWhenDisconnect(boolean kickMasterWhenDisconnect) {
-        this.kickMasterWhenDisconnect = kickMasterWhenDisconnect;
+        KICK_MASTER_HANDLE.setVolatile(this, kickMasterWhenDisconnect);
     }
 
     protected void processPlayerPluginMessage(byte[] packetData) {
         final ProxyComputeResult result = this.packetProxy.processC2S(YsmMapperPayloadManager.YSM_CHANNEL_KEY_ADVENTURE, Unpooled.copiedBuffer(packetData));
+        final Session sessionObject = (Session) SESSION_HANDLE.getVolatile(this);
+
+        // This case should never happen because player's ysm packet won't come in
+        // until we forward the handshake packet from the worker side
+        // And when the handshake packet is reached, the session was already set before
+        // see YsmMapperPayloadManager#createMapperSession
+        if (sessionObject == null) {
+            throw new IllegalStateException("Processing plugin message on non-connected mapper");
+        }
 
         switch (result.result()) {
             case MODIFY -> {
@@ -77,11 +93,11 @@ protected void processPlayerPluginMessage(byte[] packetData) {
                 byte[] data = new byte[finalData.readableBytes()];
                 finalData.readBytes(data);
 
-                this.session.send(new ServerboundCustomPayloadPacket(YsmMapperPayloadManager.YSM_CHANNEL_KEY_ADVENTURE, data));
+                sessionObject.send(new ServerboundCustomPayloadPacket(YsmMapperPayloadManager.YSM_CHANNEL_KEY_ADVENTURE, data));
             }
 
             case PASS ->
-                    this.session.send(new ServerboundCustomPayloadPacket(YsmMapperPayloadManager.YSM_CHANNEL_KEY_ADVENTURE, packetData));
+                    sessionObject.send(new ServerboundCustomPayloadPacket(YsmMapperPayloadManager.YSM_CHANNEL_KEY_ADVENTURE, packetData));
         }
     }
 
@@ -160,7 +176,6 @@ public void packetError(PacketErrorEvent event) {
 
     @Override
     public void connected(ConnectedEvent event) {
-        this.session = event.getSession();
     }
 
     @Override
@@ -178,13 +193,36 @@ public void disconnected(DisconnectedEvent event) {
         }
 
         // Remove callback
-        this.mapperPayloadManager.onWorkerSessionDisconnect(this, this.kickMasterWhenDisconnect, event.getReason()); // Fire events
-        this.session = null; //Set session to null to finalize the mapper connection
+        this.mapperPayloadManager.onWorkerSessionDisconnect(this, (boolean) KICK_MASTER_HANDLE.getVolatile(this), event.getReason()); // Fire events
+        SESSION_HANDLE.setVolatile(this, null); //Set session to null to finalize the mapper connection
+    }
+
+    protected void setSession(Session session) {
+        SESSION_HANDLE.setVolatile(this, session);
+    }
+
+    public void destroyAndAwaitDisconnected() {
+        // Prevent multiple disconnect calls
+        if (!DESTROYED_HANDLE.compareAndSet(this, false, true)) {
+            // Wait for fully disconnected
+            this.waitForDisconnected();
+            return;
+        }
+
+        final Session sessionObject = (Session) SESSION_HANDLE.getVolatile(this);
+
+        // Destroy the session
+        if (sessionObject != null) {
+            sessionObject.disconnect("DESTROYED");
+        }
+
+        // Wait for fully disconnected
+        this.waitForDisconnected();
     }
 
     protected void waitForDisconnected() {
         // We will set the session to null after finishing all disconnect logics
-        while (this.session != null) {
+        while (SESSION_HANDLE.getVolatile(this) != null) {
             Thread.onSpinWait(); // Spin wait instead of block waiting
         }
     }

Freesia-Velocity/src/main/java/meow/kikir/freesia/velocity/network/ysm/YsmMapperPayloadManager.java

@@ -217,10 +217,9 @@ public CompletableFuture<Boolean> removeVirtualPlayer(UUID playerUUID) {
         return callback;
     }
 
-    private void disconnectMapper(@NotNull MapperSessionProcessor connection) {
+    private void disconnectMapperWithoutKickingMaster(@NotNull MapperSessionProcessor connection) {
         connection.setKickMasterWhenDisconnect(false);
-        connection.getSession().disconnect("RECONNECT");
-        connection.waitForDisconnected();
+        connection.destroyAndAwaitDisconnected();
     }
 
     public void autoCreateMapper(Player player) {
@@ -241,9 +240,7 @@ public void onPlayerDisconnect(@NotNull Player player) {
         final MapperSessionProcessor mapperSession = this.mapperSessions.remove(player);
 
         if (mapperSession != null) {
-            mapperSession.setKickMasterWhenDisconnect(false); // Player already offline, so we don't disconnect again
-            mapperSession.getSession().disconnect("PLAYER DISCONNECTED");
-            mapperSession.waitForDisconnected();
+            this.disconnectMapperWithoutKickingMaster(mapperSession);
         }
     }
 
@@ -291,7 +288,7 @@ public boolean disconnectAlreadyConnected(Player player) {
         }
 
         // Will do remove in the callback
-        this.disconnectMapper(current);
+        this.disconnectMapperWithoutKickingMaster(current);
         return true;
     }
 
@@ -331,6 +328,7 @@ public void createMapperSession(@NotNull Player player, @NotNull InetSocketAddre
             throw new IllegalStateException("Mapper session not found or ready for player " + player.getUsername());
         }
 
+        packetProcessor.setSession(mapperSession);
         mapperSession.addListener(packetProcessor);
 
         // Default as Minecraft client

gradle.properties

@@ -1,5 +1,5 @@
 kotlin.code.style=official
-version=2.5.0+2.4
+version=2.5.0+2.4.1
 minecraft_version=1.21.1
 parchment_version=2024.11.17
 fabric_loader_version=0.16.13