From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Aikar <aikar@aikar.co>
Date: Sat, 23 May 2020 01:31:06 -0400
Subject: [PATCH] Fix Non Full Status Chunk NBT Memory Leak

Any full status chunk that was requested for any status less than full
would hold onto their entire nbt tree and every variable in that function.

This was due to use of a lambda that persists on the Chunk object
until that chunk reaches FULL status.

With introduction of no tick, we greatly increased the number of non
full chunks so this was really starting to hurt.

We further improve it by making a copy of the nbt tag with only the memory
it needs, so that we dont have to hold a copy to the entire compound.

diff --git a/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java b/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
index 7c04aef3eac54981ca1e34cb87d97104c3c9685b..b4246524dd11ad1e1dc94c56eee966c5a54d9ecc 100644
--- a/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
+++ b/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
@@ -27,6 +27,7 @@ import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.ListTag;
 import net.minecraft.nbt.LongArrayTag;
 import net.minecraft.nbt.ShortTag;
+import net.minecraft.nbt.Tag;
 import net.minecraft.server.level.ServerChunkCache;
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.server.level.ThreadedLevelLightEngine;
@@ -207,15 +208,9 @@ public class ChunkSerializer {
                 object2 = protochunkticklist1;
             }
 
-            object = new LevelChunk(world.getLevel(), pos, biomestorage, chunkconverter, (TickList) object1, (TickList) object2, k, achunksection, (chunk) -> {
-                ChunkSerializer.postLoadChunk(world, nbttagcompound1, chunk);
-                // CraftBukkit start - load chunk persistent data from nbt
-                net.minecraft.nbt.Tag persistentBase = nbttagcompound1.get("ChunkBukkitValues");
-                if (persistentBase instanceof CompoundTag) {
-                    chunk.persistentDataContainer.putAll((CompoundTag) persistentBase);
-                }
-                // CraftBukkit end
-            });
+            object = new LevelChunk(world.getLevel(), pos, biomestorage, chunkconverter, (TickList) object1, (TickList) object2, k, achunksection, // Paper start - fix massive nbt memory leak due to lambda. move lambda into a container method to not leak scope. Only clone needed NBT keys.
+                createLoadEntitiesConsumer(new SafeNBTCopy(nbttagcompound1, "TileEntities", "Entities", "ChunkBukkitValues")) // Paper - move CB Chunk PDC into here
+            );// Paper end
         } else {
             ProtoChunk protochunk = new ProtoChunk(pos, chunkconverter, achunksection, protochunkticklist, protochunkticklist1, world, world); // Paper - add level
 
@@ -321,6 +316,50 @@ public class ChunkSerializer {
             return new InProgressChunkHolder(protochunk1, tasksToExecuteOnMain); // Paper - Async chunk loading
         }
     }
+    // Paper start
+
+    /**
+     * This wrapper will error out if any key is accessed that wasn't copied so we can catch it easy on an update
+     */
+    private static class SafeNBTCopy extends CompoundTag {
+        private final java.util.Set<String> keys = new java.util.HashSet<String>();
+        public SafeNBTCopy(CompoundTag base, String... keys) {
+            for (String key : keys) {
+                this.keys.add(key);
+                final Tag nbtBase = base.get(key);
+                if (nbtBase != null) {
+                    this.put(key, nbtBase);
+                }
+            }
+        }
+
+        @Override
+        public boolean contains(String key) {
+            if (super.contains(key)) {
+                return true;
+            } else if (keys.contains(key)) {
+                return false;
+            }
+            throw new IllegalStateException("Missing Key " + key + " in SafeNBTCopy");
+        }
+
+        @Override
+        public boolean contains(String key, int type) {
+            return contains(key) && super.contains(key, type);
+        }
+    }
+    private static java.util.function.Consumer<LevelChunk> createLoadEntitiesConsumer(CompoundTag nbt) {
+        return (chunk) -> {
+            postLoadChunk(chunk.level, nbt, chunk);
+            // CraftBukkit start - load chunk persistent data from nbt
+            Tag persistentBase = nbt.get("ChunkBukkitValues");
+            if (persistentBase instanceof CompoundTag) {
+                chunk.persistentDataContainer.putAll((CompoundTag) persistentBase);
+            }
+            // CraftBukkit end
+        };
+    }
+    // Paper end
 
     // Paper start - async chunk save for unload
     public static final class AsyncSaveData {