fix: fix potential memory leak in chunk loading when there are except…

…ions thrown
AllayMC · Dec 4, 2024 · d87339b · d87339b
1 parent 7567a23
commit d87339b
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 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
@@ -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<Boolean> 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
@@ -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
@@ -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
@@ -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<Chunk> 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;
     }