From 83c1a2fd6315174dccedaaba8650305e07cca09e Mon Sep 17 00:00:00 2001
From: s809 <11816467-s809@users.noreply.gitlab.com>
Date: Sat, 20 Jul 2024 23:56:24 +0500
Subject: [PATCH] Fix handler registration

---
 .../core/level/DhServerLevel.java             |  5 ++--
 .../network/event/NetworkEventSource.java     | 26 +++++++++++++------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/core/src/main/java/com/seibel/distanthorizons/core/level/DhServerLevel.java b/core/src/main/java/com/seibel/distanthorizons/core/level/DhServerLevel.java
index 5c805aa96..a4ee82e46 100644
--- a/core/src/main/java/com/seibel/distanthorizons/core/level/DhServerLevel.java
+++ b/core/src/main/java/com/seibel/distanthorizons/core/level/DhServerLevel.java
@@ -265,12 +265,11 @@ public class DhServerLevel extends AbstractDhLevel implements IDhServerLevel
 				continue;
 			}
 			
-			// Prevent adding or removing requests
+			// Make this group unavailable for adding into
+			this.requestGroupsByPos.remove(entry.getKey());
 			requestGroup.requestRemoveSemaphore.acquireUninterruptibly(Short.MAX_VALUE);
 			requestGroup.requestAddSemaphore.acquireUninterruptibly(Short.MAX_VALUE);
 			
-			this.requestGroupsByPos.remove(entry.getKey());
-			
 			ThreadPoolExecutor executor = ThreadPoolUtil.getNetworkCompressionExecutor();
 			if (executor == null)
 			{
diff --git a/core/src/main/java/com/seibel/distanthorizons/core/network/event/NetworkEventSource.java b/core/src/main/java/com/seibel/distanthorizons/core/network/event/NetworkEventSource.java
index 50ba33e05..f1c0f3dd5 100644
--- a/core/src/main/java/com/seibel/distanthorizons/core/network/event/NetworkEventSource.java
+++ b/core/src/main/java/com/seibel/distanthorizons/core/network/event/NetworkEventSource.java
@@ -42,7 +42,13 @@ 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>, ConcurrentMap<NetworkEventSource, Consumer<NetworkMessage>>> handlers = new ConcurrentHashMap<>();
+	protected final ConcurrentMap<
+			Class<? extends NetworkMessage>,
+			ConcurrentMap<
+					NetworkEventSource,
+					Set<Consumer<NetworkMessage>>
+			>
+	> handlers = new ConcurrentHashMap<>();
 	
 	private final ConcurrentMap<Long, FutureResponseData> pendingFutures = new ConcurrentHashMap<>();
 	private final Set<Long> cancelledFutures = Collections.newSetFromMap(CacheBuilder.newBuilder()
@@ -55,13 +61,16 @@ public abstract class NetworkEventSource
 	{
 		boolean handled = false;
 		
-		ConcurrentMap<NetworkEventSource, Consumer<NetworkMessage>> handlerMap = this.handlers.get(message.getClass());
-		if (handlerMap != null)
+		ConcurrentMap<NetworkEventSource, Set<Consumer<NetworkMessage>>> handlersByEventSource = this.handlers.get(message.getClass());
+		if (handlersByEventSource != null)
 		{
-			for (Consumer<NetworkMessage> handler : handlerMap.values())
+			for (Set<Consumer<NetworkMessage>> handlerSet : handlersByEventSource.values())
 			{
-				handled = true;
-				handler.accept(message);
+				for (Consumer<NetworkMessage> handler : handlerSet)
+				{
+					handled = true;
+					handler.accept(message);
+				}
 			}
 		}
 		
@@ -113,12 +122,13 @@ public abstract class NetworkEventSource
 					}
 					return new ConcurrentHashMap<>();
 				})
-				.put(instance, (Consumer<NetworkMessage>) handlerImplementation);
+				.computeIfAbsent(instance, _instance -> ConcurrentHashMap.newKeySet())
+				.add((Consumer<NetworkMessage>) handlerImplementation);
 	}
 	
 	protected void removeAllHandlers(NetworkEventSource childInstance)
 	{
-		for (ConcurrentMap<NetworkEventSource, Consumer<NetworkMessage>> handlerMap : this.handlers.values())
+		for (ConcurrentMap<NetworkEventSource, Set<Consumer<NetworkMessage>>> handlerMap : this.handlers.values())
 		{
 			handlerMap.remove(childInstance);
 		}