properly detach addon instance on config <-> play transition

fabric-networking-api-v1/src/client/java/net/fabricmc/fabric/impl/networking/client/ClientConfigurationNetworkAddon.java

@@ -53,7 +53,7 @@ public ClientConfigurationNetworkAddon(ClientConfigurationNetworkHandler handler
 		this.registerPendingChannels((ChannelInfoHolder) this.connection, NetworkState.CONFIGURATION);
 
 		// Register global receivers and attach to session
-		this.receiver.startSession(this);
+		this.startSession();
 	}
 
 	@Override
@@ -148,12 +148,12 @@ protected void handleUnregistration(Identifier channelName) {
 	public void handleReady() {
 		ClientConfigurationConnectionEvents.READY.invoker().onConfigurationReady(this.handler, this.client);
 		ClientNetworkingImpl.setClientConfigurationAddon(null);
+		this.endSession();
 	}
 
 	@Override
 	protected void invokeDisconnectEvent() {
 		ClientConfigurationConnectionEvents.DISCONNECT.invoker().onConfigurationDisconnect(this.handler, this.client);
-		this.receiver.endSession(this);
 	}
 
 	@Override

fabric-networking-api-v1/src/client/java/net/fabricmc/fabric/impl/networking/client/ClientLoginNetworkAddon.java

@@ -53,7 +53,7 @@ public ClientLoginNetworkAddon(ClientLoginNetworkHandler handler, MinecraftClien
 		this.client = client;
 
 		ClientLoginConnectionEvents.INIT.invoker().onLoginStart(this.handler, this.client);
-		this.receiver.startSession(this);
+		this.startSession();
 	}
 
 	public boolean handlePacket(LoginQueryRequestS2CPacket packet) {
@@ -114,11 +114,6 @@ protected void handleUnregistration(Identifier channelName) {
 	@Override
 	protected void invokeDisconnectEvent() {
 		ClientLoginConnectionEvents.DISCONNECT.invoker().onLoginDisconnect(this.handler, this.client);
-		this.receiver.endSession(this);
-	}
-
-	public void handleConfigurationTransition() {
-		this.receiver.endSession(this);
 	}
 
 	@Override

fabric-networking-api-v1/src/client/java/net/fabricmc/fabric/impl/networking/client/ClientPlayNetworkAddon.java

@@ -55,7 +55,7 @@ public ClientPlayNetworkAddon(ClientPlayNetworkHandler handler, MinecraftClient
 		this.registerPendingChannels((ChannelInfoHolder) this.connection, NetworkState.PLAY);
 
 		// Register global receivers and attach to session
-		this.receiver.startSession(this);
+		this.startSession();
 	}
 
 	@Override
@@ -148,7 +148,6 @@ protected void handleUnregistration(Identifier channelName) {
 	@Override
 	protected void invokeDisconnectEvent() {
 		ClientPlayConnectionEvents.DISCONNECT.invoker().onPlayDisconnect(this.handler, this.client);
-		this.receiver.endSession(this);
 	}
 
 	@Override

fabric-networking-api-v1/src/client/java/net/fabricmc/fabric/mixin/networking/client/ClientLoginNetworkHandlerMixin.java

@@ -71,7 +71,7 @@ private void invokeLoginDisconnectEvent(Text reason, CallbackInfo ci) {
 
 	@Inject(method = "onSuccess", at = @At("HEAD"))
 	private void handleConfigurationTransition(CallbackInfo ci) {
-		addon.handleConfigurationTransition();
+		this.addon.endSession();
 	}
 
 	@Inject(method = "onSuccess", at = @At("TAIL"))

fabric-networking-api-v1/src/client/java/net/fabricmc/fabric/mixin/networking/client/ClientPlayNetworkHandlerMixin.java

@@ -27,6 +27,7 @@
 import net.minecraft.client.network.ClientConnectionState;
 import net.minecraft.client.network.ClientPlayNetworkHandler;
 import net.minecraft.network.ClientConnection;
+import net.minecraft.network.packet.s2c.play.EnterReconfigurationS2CPacket;
 import net.minecraft.network.packet.s2c.play.GameJoinS2CPacket;
 
 import net.fabricmc.fabric.impl.networking.NetworkHandlerExtensions;
@@ -56,6 +57,11 @@ private void handleServerPlayReady(GameJoinS2CPacket packet, CallbackInfo ci) {
 		this.addon.onServerReady();
 	}
 
+	@Inject(method = "onEnterReconfiguration", at = @At(value = "INVOKE", target = "Lnet/minecraft/network/ClientConnection;setPacketListener(Lnet/minecraft/network/listener/PacketListener;)V"))
+	private void handleConfigurationTransition(EnterReconfigurationS2CPacket packet, CallbackInfo ci) {
+		this.addon.endSession();
+	}
+
 	@Override
 	public ClientPlayNetworkAddon getAddon() {
 		return this.addon;

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

@@ -126,9 +126,18 @@ public Set<Identifier> getReceivableChannels() {
 
 	protected abstract void handleUnregistration(Identifier channelName);
 
+	protected final void startSession() {
+		this.receiver.startSession(this);
+	}
+
+	public final void endSession() {
+		this.receiver.endSession(this);
+	}
+
 	public final void handleDisconnect() {
 		if (disconnected.compareAndSet(false, true)) {
 			invokeDisconnectEvent();
+			endSession();
 		}
 	}
 

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

@@ -26,11 +26,15 @@
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import net.minecraft.network.NetworkState;
 import net.minecraft.util.Identifier;
 
 public final class GlobalReceiverRegistry<H> {
+	private static final Logger LOGGER = LoggerFactory.getLogger(GlobalReceiverRegistry.class);
+
 	private final NetworkState state;
 
 	private final ReadWriteLock lock = new ReentrantReadWriteLock();
@@ -135,6 +139,7 @@ public void startSession(AbstractNetworkAddon<H> addon) {
 
 		try {
 			this.trackedAddons.add(addon);
+			this.logTrackedAddonSize();
 		} finally {
 			lock.unlock();
 		}
@@ -145,17 +150,29 @@ public void endSession(AbstractNetworkAddon<H> addon) {
 		lock.lock();
 
 		try {
+			this.logTrackedAddonSize();
 			this.trackedAddons.remove(addon);
 		} finally {
 			lock.unlock();
 		}
 	}
 
+	/**
+	 * In practice, trackedAddons should never contain more than one instance.
+	 */
+	private void logTrackedAddonSize() {
+		if (LOGGER.isDebugEnabled() && this.trackedAddons.size() > 1) {
+			LOGGER.error("{} receiver registry tracks {} addon instances where it should only tracks one!", state.getId(), trackedAddons.size());
+		}
+	}
+
 	private void handleRegistration(Identifier channelName, H handler) {
 		Lock lock = this.lock.writeLock();
 		lock.lock();
 
 		try {
+			this.logTrackedAddonSize();
+
 			for (AbstractNetworkAddon<H> addon : this.trackedAddons) {
 				addon.registerChannel(channelName, handler);
 			}
@@ -169,6 +186,8 @@ private void handleUnregistration(Identifier channelName) {
 		lock.lock();
 
 		try {
+			this.logTrackedAddonSize();
+
 			for (AbstractNetworkAddon<H> addon : this.trackedAddons) {
 				addon.unregisterChannel(channelName);
 			}

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

@@ -54,7 +54,7 @@ public ServerConfigurationNetworkAddon(ServerConfigurationNetworkHandler handler
 		this.registerPendingChannels((ChannelInfoHolder) this.connection, NetworkState.CONFIGURATION);
 
 		// Register global receivers and attach to session
-		this.receiver.startSession(this);
+		this.startSession();
 	}
 
 	@Override
@@ -177,7 +177,6 @@ protected void handleUnregistration(Identifier channelName) {
 	@Override
 	protected void invokeDisconnectEvent() {
 		ServerConfigurationConnectionEvents.DISCONNECT.invoker().onConfigureDisconnect(handler, server);
-		this.receiver.endSession(this);
 	}
 
 	@Override

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

@@ -69,7 +69,7 @@ public ServerLoginNetworkAddon(ServerLoginNetworkHandler handler) {
 		this.queryIdFactory = QueryIdFactory.create();
 
 		ServerLoginConnectionEvents.INIT.invoker().onLoginInit(handler, this.server);
-		this.receiver.startSession(this);
+		this.startSession();
 	}
 
 	// return true if no longer ticks query
@@ -202,11 +202,6 @@ protected void handleUnregistration(Identifier channelName) {
 	@Override
 	protected void invokeDisconnectEvent() {
 		ServerLoginConnectionEvents.DISCONNECT.invoker().onLoginDisconnect(this.handler, this.server);
-		this.receiver.endSession(this);
-	}
-
-	public void handleConfigurationTransition() {
-		this.receiver.endSession(this);
 	}
 
 	@Override

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

@@ -51,7 +51,7 @@ public ServerPlayNetworkAddon(ServerPlayNetworkHandler handler, MinecraftServer
 		this.registerPendingChannels((ChannelInfoHolder) this.connection, NetworkState.PLAY);
 
 		// Register global receivers and attach to session
-		this.receiver.startSession(this);
+		this.startSession();
 	}
 
 	@Override
@@ -139,7 +139,6 @@ protected void handleUnregistration(Identifier channelName) {
 	@Override
 	protected void invokeDisconnectEvent() {
 		ServerPlayConnectionEvents.DISCONNECT.invoker().onPlayDisconnect(this.handler, this.server);
-		this.receiver.endSession(this);
 	}
 
 	@Override

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

@@ -29,6 +29,7 @@
 
 import net.minecraft.network.ClientConnection;
 import net.minecraft.network.packet.Packet;
+import net.minecraft.network.packet.c2s.config.ReadyC2SPacket;
 import net.minecraft.network.packet.s2c.common.DisconnectS2CPacket;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.network.ConnectedClientData;
@@ -146,6 +147,11 @@ private void handleDisconnection(Text reason, CallbackInfo ci) {
 		this.addon.handleDisconnect();
 	}
 
+	@Inject(method = "onReady", at = @At(value = "INVOKE", target = "Lnet/minecraft/server/PlayerManager;onPlayerConnect(Lnet/minecraft/network/ClientConnection;Lnet/minecraft/server/network/ServerPlayerEntity;Lnet/minecraft/server/network/ConnectedClientData;)V"))
+	private void handlePlayTransition(ReadyC2SPacket packet, CallbackInfo ci) {
+		this.addon.endSession();
+	}
+
 	@Override
 	public ServerConfigurationNetworkAddon getAddon() {
 		return addon;

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

@@ -17,7 +17,6 @@
 package net.fabricmc.fabric.mixin.networking;
 
 import com.mojang.authlib.GameProfile;
-import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
@@ -42,10 +41,6 @@
 
 @Mixin(ServerLoginNetworkHandler.class)
 abstract class ServerLoginNetworkHandlerMixin implements NetworkHandlerExtensions, DisconnectPacketSource, PacketCallbackListener {
-	@Shadow
-	@Final
-	private MinecraftServer server;
-
 	@Shadow
 	protected abstract void tickVerify(GameProfile profile);
 
@@ -88,8 +83,8 @@ private void handleDisconnection(Text reason, CallbackInfo ci) {
 	}
 
 	@Inject(method = "sendSuccessPacket", at = @At("HEAD"))
-	private void handlePlayTransitionNormal(GameProfile profile, CallbackInfo ci) {
-		this.addon.handleConfigurationTransition();
+	private void handleConfigurationTransition(GameProfile profile, CallbackInfo ci) {
+		this.addon.endSession();
 	}
 
 	@Override

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

@@ -24,6 +24,7 @@
 
 import net.minecraft.network.ClientConnection;
 import net.minecraft.network.packet.Packet;
+import net.minecraft.network.packet.c2s.play.AcknowledgeReconfigurationC2SPacket;
 import net.minecraft.network.packet.s2c.common.DisconnectS2CPacket;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.network.ConnectedClientData;
@@ -57,6 +58,11 @@ private void handleDisconnection(Text reason, CallbackInfo ci) {
 		this.addon.handleDisconnect();
 	}
 
+	@Inject(method = "onAcknowledgeReconfiguration", at = @At(value = "INVOKE", target = "Lnet/minecraft/network/ClientConnection;setPacketListener(Lnet/minecraft/network/listener/PacketListener;)V"))
+	private void handleConfigurationTransition(AcknowledgeReconfigurationC2SPacket packet, CallbackInfo ci) {
+		this.addon.endSession();
+	}
+
 	@Override
 	public ServerPlayNetworkAddon getAddon() {
 		return this.addon;