Fix event handler registration

core/src/main/java/com/seibel/distanthorizons/core/multiplayer/client/ClientNetworkState.java

@@ -89,10 +89,10 @@ public class ClientNetworkState implements Closeable
 	public FullDataSourceV2DTO decodeDataSourceAndReleaseBuffer(IFullDataPayloadMessage<?> msg)
 	{
 		CompositeByteBuf composite = this.fullDataBuffers.remove(msg.getDtoBufferId());
+		Objects.requireNonNull(composite);
 		
 		try
 		{
-			Objects.requireNonNull(composite);
 			return INetworkObject.decodeToInstance(new FullDataSourceV2DTO(), composite);
 		}
 		finally

core/src/main/java/com/seibel/distanthorizons/core/network/event/NetworkEventSource.java

@@ -33,7 +33,6 @@ import org.apache.logging.log4j.LogManager;
 
 import java.io.InvalidClassException;
 import java.util.Collections;
-import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.*;
 import java.util.function.Consumer;
@@ -43,11 +42,7 @@ public abstract class NetworkEventSource
 	private static final ConfigBasedLogger LOGGER = new ConfigBasedLogger(LogManager.getLogger(),
 			() -> Config.Client.Advanced.Logging.logNetworkEvent.get());
 	
-	protected final ConcurrentMap<Class<? extends NetworkMessage>, Set<Consumer<NetworkMessage>>> handlers = new ConcurrentHashMap<>();
-	protected boolean hasHandler(Class<? extends NetworkMessage> handlerClass)
-	{
-		return this.handlers.containsKey(handlerClass);
-	}
+	protected final ConcurrentMap<Class<? extends NetworkMessage>, ConcurrentMap<NetworkEventSource, Consumer<NetworkMessage>>> handlers = new ConcurrentHashMap<>();
 	
 	private final ConcurrentMap<Long, FutureResponseData> pendingFutures = new ConcurrentHashMap<>();
 	private final Set<Long> cancelledFutures = Collections.newSetFromMap(CacheBuilder.newBuilder()
@@ -60,10 +55,10 @@ public abstract class NetworkEventSource
 	{
 		boolean handled = false;
 		
-		Set<Consumer<NetworkMessage>> handlerList = this.handlers.get(message.getClass());
-		if (handlerList != null)
+		ConcurrentMap<NetworkEventSource, Consumer<NetworkMessage>> handlerMap = this.handlers.get(message.getClass());
+		if (handlerMap != null)
 		{
-			for (Consumer<NetworkMessage> handler : handlerList)
+			for (Consumer<NetworkMessage> handler : handlerMap.values())
 			{
 				handled = true;
 				handler.accept(message);
@@ -104,25 +99,29 @@ public abstract class NetworkEventSource
 		}
 	}
 	
-	public <T extends NetworkMessage> void registerHandler(Class<T> handlerClass, Consumer<T> handlerImplementation)
+	public abstract <T extends NetworkMessage> void registerHandler(Class<T> handlerClass, Consumer<T> handlerImplementation);
+	
+	protected <T extends NetworkMessage> void registerHandler(NetworkEventSource instance, Class<T> handlerClass, Consumer<T> handlerImplementation)
 	{
 		//noinspection unchecked
 		this.handlers.computeIfAbsent(handlerClass, missingHandlerClass ->
 				{
 					// Will throw if the handler class is not found
-					if (handlerClass != CloseEvent.class)
+					if (missingHandlerClass != CloseEvent.class)
 					{
-						MessageRegistry.INSTANCE.getMessageId(handlerClass);
+						MessageRegistry.INSTANCE.getMessageId(missingHandlerClass);
 					}
-					return new HashSet<>();
+					return new ConcurrentHashMap<>();
 				})
-				.add((Consumer<NetworkMessage>) handlerImplementation);
+				.put(instance, (Consumer<NetworkMessage>) handlerImplementation);
 	}
 	
-	protected <T extends NetworkMessage> void removeHandler(Class<T> handlerClass, Consumer<T> handlerImplementation)
+	protected void removeAllHandlers(NetworkEventSource childInstance)
 	{
-		this.handlers.computeIfAbsent(handlerClass, missingHandlerClass -> new HashSet<>())
-				.remove(handlerImplementation);
+		for (ConcurrentMap<NetworkEventSource, Consumer<NetworkMessage>> handlerMap : this.handlers.values())
+		{
+			handlerMap.remove(childInstance);
+		}
 	}
 	
 	

core/src/main/java/com/seibel/distanthorizons/core/network/event/ScopedNetworkEventSource.java

@@ -42,21 +42,13 @@ public final class ScopedNetworkEventSource extends NetworkEventSource
 			return;
 		}
 		
-		if (!this.hasHandler(handlerClass))
-		{
-			this.parent.registerHandler(handlerClass, this::handleMessage);
-		}
-		
-		super.registerHandler(handlerClass, handlerImplementation);
+		this.parent.registerHandler(this, handlerClass, this::handleMessage);
 	}
 	
 	@Override
 	public void close()
 	{
 		this.isClosed = true;
-		for (Class<? extends NetworkMessage> handlerClass : this.handlers.keySet())
-		{
-			this.parent.removeHandler(handlerClass, this::handleMessage);
-		}
+		this.parent.removeAllHandlers(this);
 	}
 }

core/src/main/java/com/seibel/distanthorizons/core/network/session/Session.java

@@ -15,6 +15,7 @@ import org.jetbrains.annotations.Nullable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 public class Session extends NetworkEventSource
 {
@@ -64,6 +65,17 @@ public class Session extends NetworkEventSource
 		}
 	}
 	
+	@Override
+	public <T extends NetworkMessage> void registerHandler(Class<T> handlerClass, Consumer<T> handlerImplementation)
+	{
+		if (this.closeReason.get() != null)
+		{
+			return;
+		}
+		
+		this.registerHandler(this, handlerClass, handlerImplementation);
+	}
+	
 	public <TResponse extends TrackableMessage> CompletableFuture<TResponse> sendRequest(TrackableMessage msg, Class<TResponse> responseClass)
 	{
 		msg.setSession(this);