From d87339b4b06519c3d776f28174303c25cb79ff16 Mon Sep 17 00:00:00 2001 From: daoge_cmd <3523206925@qq.com> Date: Wed, 4 Dec 2024 17:52:35 +0800 Subject: [PATCH] fix: fix potential memory leak in chunk loading when there are exceptions thrown --- .../server/world/chunk/AllayChunk.java | 27 +++++++++----- .../server/world/chunk/AllayUnsafeChunk.java | 6 +++- .../world/generator/AllayWorldGenerator.java | 3 +- .../world/service/AllayChunkService.java | 35 +++++++++++-------- 4 files changed, 45 insertions(+), 26 deletions(-) diff --git a/server/src/main/java/org/allaymc/server/world/chunk/AllayChunk.java b/server/src/main/java/org/allaymc/server/world/chunk/AllayChunk.java index 529ce0881..72ba810c5 100644 --- a/server/src/main/java/org/allaymc/server/world/chunk/AllayChunk.java +++ b/server/src/main/java/org/allaymc/server/world/chunk/AllayChunk.java @@ -29,6 +29,7 @@ import javax.annotation.concurrent.ThreadSafe; import java.util.*; import java.util.concurrent.locks.StampedLock; +import java.util.function.Consumer; import java.util.function.Predicate; /** @@ -51,8 +52,9 @@ public class AllayChunk implements Chunk { @Getter protected boolean loaded = false; // The callback to be called when the chunk is loaded into the world + // The provided boolean value indicated whether the chunk is set successfully @Setter - protected Runnable chunkSetCallback; + protected Consumer chunkSetCallback; protected int autoSaveTimer = 0; private static void checkXZ(int x, int z) { @@ -396,19 +398,26 @@ public ChunkSection[] getSections() { public void beforeSetChunk(Dimension dimension) { unsafeChunk.beforeSetChunk(dimension); - unsafeChunk.setBlockChangeCallback((x, y, z, blockState, layer) -> { - if (layer != 0) return; - ((AllayLightService) dimension.getLightService()).onBlockChange(x + (unsafeChunk.x << 4), y, z + (unsafeChunk.z << 4), blockState.getBlockStateData().lightEmission(), blockState.getBlockStateData().lightDampening()); - }); } - public void afterSetChunk(Dimension dimension) { + public void afterSetChunk(Dimension dimension, boolean success) { if (chunkSetCallback != null) { - chunkSetCallback.run(); + chunkSetCallback.accept(success); } - loaded = true; - unsafeChunk.afterSetChunk(dimension); + unsafeChunk.afterSetChunk(dimension, success); + + if (!success) { + return; + } + + unsafeChunk.setBlockChangeCallback((x, y, z, blockState, layer) -> { + if (layer != 0) { + return; + } + ((AllayLightService) dimension.getLightService()).onBlockChange(x + (unsafeChunk.x << 4), y, z + (unsafeChunk.z << 4), blockState.getBlockStateData().lightEmission(), blockState.getBlockStateData().lightDampening()); + }); ((AllayLightService) dimension.getLightService()).onChunkLoad(this); + loaded = true; } @Override diff --git a/server/src/main/java/org/allaymc/server/world/chunk/AllayUnsafeChunk.java b/server/src/main/java/org/allaymc/server/world/chunk/AllayUnsafeChunk.java index e22fe8079..e99ae4532 100644 --- a/server/src/main/java/org/allaymc/server/world/chunk/AllayUnsafeChunk.java +++ b/server/src/main/java/org/allaymc/server/world/chunk/AllayUnsafeChunk.java @@ -132,7 +132,11 @@ public void beforeSetChunk(Dimension dimension) { } } - public void afterSetChunk(Dimension dimension) { + public void afterSetChunk(Dimension dimension, boolean success) { + if (!success) { + return; + } + if (entityNbtList != null && !entityNbtList.isEmpty()) { for (var nbt : entityNbtList) { var entity = EntityHelper.fromNBT(dimension, nbt); diff --git a/server/src/main/java/org/allaymc/server/world/generator/AllayWorldGenerator.java b/server/src/main/java/org/allaymc/server/world/generator/AllayWorldGenerator.java index 3d481d32c..ff541bcec 100644 --- a/server/src/main/java/org/allaymc/server/world/generator/AllayWorldGenerator.java +++ b/server/src/main/java/org/allaymc/server/world/generator/AllayWorldGenerator.java @@ -130,7 +130,8 @@ private void processPopulationQueue() { statusPopulatedToFinished(chunk); var chunkHash = HashUtils.hashXZ(chunk.getX(), chunk.getZ()); // Remove recorded futures - ((AllayChunk) chunk).setChunkSetCallback(() -> { + ((AllayChunk) chunk).setChunkSetCallback(success -> { + // The stored futures should always being removed chunkNoiseFutures.remove(chunkHash); chunkFutures.remove(chunkHash); }); diff --git a/server/src/main/java/org/allaymc/server/world/service/AllayChunkService.java b/server/src/main/java/org/allaymc/server/world/service/AllayChunkService.java index e1162f10b..ff5b357ef 100644 --- a/server/src/main/java/org/allaymc/server/world/service/AllayChunkService.java +++ b/server/src/main/java/org/allaymc/server/world/service/AllayChunkService.java @@ -123,8 +123,9 @@ private void removeUnusedChunks() { private void setChunk(int x, int z, Chunk chunk) { var chunkHash = HashUtils.hashXZ(x, z); - if (loadedChunks.putIfAbsent(chunkHash, chunk) != null) + if (loadedChunks.putIfAbsent(chunkHash, chunk) != null) { throw new IllegalStateException("Trying to set a chunk (" + x + "," + z + ") which is already loaded"); + } } @Override @@ -201,20 +202,24 @@ public CompletableFuture loadChunk(int x, int z) { }).exceptionally(t -> { log.error("Error while generating chunk ({},{}) !", x, z, t); return AllayUnsafeChunk.builder().newChunk(x, z, dimension.getDimensionInfo()).toSafeChunk(); - }).thenApply(preparedChunk -> { - ((AllayChunk) preparedChunk).beforeSetChunk(dimension); - setChunk(x, z, preparedChunk); - ((AllayChunk) preparedChunk).afterSetChunk(dimension); - future.complete(preparedChunk); - loadingChunks.remove(hashXZ); - - var chunkLoadEvent = new ChunkLoadEvent(dimension, preparedChunk); - chunkLoadEvent.call(); - - return preparedChunk; - }).exceptionally(t -> { - log.error("Error while setting chunk ({},{}) !", x, z, t); - return AllayUnsafeChunk.builder().newChunk(x, z, dimension.getDimensionInfo()).toSafeChunk(); + }).thenAccept(preparedChunk -> { + boolean success = true; + try { + ((AllayChunk) preparedChunk).beforeSetChunk(dimension); + setChunk(x, z, preparedChunk); + } catch (Throwable t) { + log.error("Error while setting chunk ({},{}) !", x, z, t); + success = false; + } finally { + loadingChunks.remove(hashXZ); + ((AllayChunk) preparedChunk).afterSetChunk(dimension, success); + if (success) { + future.complete(preparedChunk); + new ChunkLoadEvent(dimension, preparedChunk).call(); + } else { + future.complete(null); + } + } }); return future; }