diff --git a/core/src/main/java/com/seibel/distanthorizons/core/multiplayer/client/ClientNetworkState.java b/core/src/main/java/com/seibel/distanthorizons/core/multiplayer/client/ClientNetworkState.java index a63f8e2db..b84263111 100644 --- a/core/src/main/java/com/seibel/distanthorizons/core/multiplayer/client/ClientNetworkState.java +++ b/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 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 699f3a516..50ba33e05 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 @@ -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, Set>> handlers = new ConcurrentHashMap<>(); - protected boolean hasHandler(Class handlerClass) - { - return this.handlers.containsKey(handlerClass); - } + protected final ConcurrentMap, ConcurrentMap>> handlers = new ConcurrentHashMap<>(); private final ConcurrentMap pendingFutures = new ConcurrentHashMap<>(); private final Set cancelledFutures = Collections.newSetFromMap(CacheBuilder.newBuilder() @@ -60,10 +55,10 @@ public abstract class NetworkEventSource { boolean handled = false; - Set> handlerList = this.handlers.get(message.getClass()); - if (handlerList != null) + ConcurrentMap> handlerMap = this.handlers.get(message.getClass()); + if (handlerMap != null) { - for (Consumer handler : handlerList) + for (Consumer handler : handlerMap.values()) { handled = true; handler.accept(message); @@ -104,25 +99,29 @@ public abstract class NetworkEventSource } } - public void registerHandler(Class handlerClass, Consumer handlerImplementation) + public abstract void registerHandler(Class handlerClass, Consumer handlerImplementation); + + protected void registerHandler(NetworkEventSource instance, Class handlerClass, Consumer 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) handlerImplementation); + .put(instance, (Consumer) handlerImplementation); } - protected void removeHandler(Class handlerClass, Consumer handlerImplementation) + protected void removeAllHandlers(NetworkEventSource childInstance) { - this.handlers.computeIfAbsent(handlerClass, missingHandlerClass -> new HashSet<>()) - .remove(handlerImplementation); + for (ConcurrentMap> handlerMap : this.handlers.values()) + { + handlerMap.remove(childInstance); + } } diff --git a/core/src/main/java/com/seibel/distanthorizons/core/network/event/ScopedNetworkEventSource.java b/core/src/main/java/com/seibel/distanthorizons/core/network/event/ScopedNetworkEventSource.java index 2c99e25dd..cf690aeb0 100644 --- a/core/src/main/java/com/seibel/distanthorizons/core/network/event/ScopedNetworkEventSource.java +++ b/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 handlerClass : this.handlers.keySet()) - { - this.parent.removeHandler(handlerClass, this::handleMessage); - } + this.parent.removeAllHandlers(this); } } \ No newline at end of file diff --git a/core/src/main/java/com/seibel/distanthorizons/core/network/session/Session.java b/core/src/main/java/com/seibel/distanthorizons/core/network/session/Session.java index 064ab5620..3527536be 100644 --- a/core/src/main/java/com/seibel/distanthorizons/core/network/session/Session.java +++ b/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 void registerHandler(Class handlerClass, Consumer handlerImplementation) + { + if (this.closeReason.get() != null) + { + return; + } + + this.registerHandler(this, handlerClass, handlerImplementation); + } + public CompletableFuture sendRequest(TrackableMessage msg, Class responseClass) { msg.setSession(this);